Language/Atom/UeMap.hs

-- | 
-- Module: UeMap
-- Description: Sharing for UEs, based on IntMaps.
-- Copyright: (c) 2013 Tom Hawkins & Lee Pike
--
-- Sharing for 'UE's, based on IntMaps. The idea is to share subexpressions
-- of 'UE's.

module Language.Atom.UeMap
 ( UeElem (..)
 , MUV (..)
 , UeMap
 , emptyMap
 , Hash
 , typeOf
 , UeState
 , recoverUE
 , getUE
 , newUE
 , newUV
 , maybeUpdate
 , ueUpstream
 , nearestUVs
 , arrayIndices
 , isMathHCall
 ) where

import Control.Monad.State.Strict
import qualified Data.Bimap as M
import Data.List (nub)

import Language.Atom.Expressions hiding (typeOf)
import qualified Language.Atom.Expressions as E

type Hash = Int

-- | Untyped variables map.
data MUV
 = MUV Int String Const
 | MUVArray UA Hash
 | MUVExtern String Type
 deriving (Show, Eq, Ord)

-- | Transforms a 'UV' into a 'MUV', returning the possibly updated map.
newUV :: UV -> UeMap -> (MUV, UeMap)
newUV u mp =
 case u of
 UV i j k -> (MUV i j k, mp)
 UVExtern i j -> (MUVExtern i j, mp)
 UVArray arr ue_ -> let (h,mp') = newUE ue_ mp in
 (MUVArray arr h, mp')

-- | Corresponds to 'UE's --- the elements in the sharing structure.
data UeElem
 = MUVRef !MUV
 | MUConst !Const
 | MUCast !Type !Hash
 | MUAdd !Hash !Hash
 | MUSub !Hash !Hash
 | MUMul !Hash !Hash
 | MUDiv !Hash !Hash
 | MUMod !Hash !Hash
 | MUNot !Hash
 | MUAnd [Hash]
 | MUBWNot !Hash
 | MUBWAnd !Hash !Hash
 | MUBWOr !Hash !Hash
 | MUBWXor !Hash !Hash
 | MUBWShiftL !Hash !Hash
 | MUBWShiftR !Hash !Hash
 | MUEq !Hash !Hash
 | MULt !Hash !Hash
 | MUMux !Hash !Hash !Hash
 | MUF2B !Hash
 | MUD2B !Hash
 | MUB2F !Hash
 | MUB2D !Hash
-- math.h:
 | MUPi
 | MUExp !Hash
 | MULog !Hash
 | MUSqrt !Hash
 | MUPow !Hash !Hash
 | MUSin !Hash
 | MUAsin !Hash
 | MUCos !Hash
 | MUAcos !Hash
 | MUSinh !Hash
 | MUCosh !Hash
 | MUAsinh !Hash
 | MUAcosh !Hash
 | MUAtan !Hash
 | MUAtanh !Hash
 deriving (Show, Eq, Ord)

typeOf :: Hash -> UeMap -> Type
typeOf h mp = case getUE h mp of
 MUVRef (MUV _ _ a) -> E.typeOf a
 MUVRef (MUVArray a _) -> E.typeOf a
 MUVRef (MUVExtern _ t) -> t
 MUCast t _ -> t
 MUConst c -> E.typeOf c
 MUAdd a _ -> typeOf' a
 MUSub a _ -> typeOf' a
 MUMul a _ -> typeOf' a
 MUDiv a _ -> typeOf' a
 MUMod a _ -> typeOf' a
 MUNot _ -> Bool
 MUAnd _ -> Bool
 MUBWNot a -> typeOf' a
 MUBWAnd a _ -> typeOf' a
 MUBWOr a _ -> typeOf' a
 MUBWXor a _ -> typeOf' a
 MUBWShiftL a _ -> typeOf' a
 MUBWShiftR a _ -> typeOf' a
 MUEq _ _ -> Bool
 MULt _ _ -> Bool
 MUMux _ a _ -> typeOf' a
 MUF2B _ -> Word32
 MUD2B _ -> Word64
 MUB2F _ -> Float
 MUB2D _ -> Double
-- math.h:
 MUPi -> Double
 MUExp a -> typeOf' a
 MULog a -> typeOf' a
 MUSqrt a -> typeOf' a
 MUPow a _ -> typeOf' a
 MUSin a -> typeOf' a
 MUAsin a -> typeOf' a
 MUCos a -> typeOf' a
 MUAcos a -> typeOf' a
 MUSinh a -> typeOf' a
 MUCosh a -> typeOf' a
 MUAsinh a -> typeOf' a
 MUAcosh a -> typeOf' a
 MUAtan a -> typeOf' a
 MUAtanh a -> typeOf' a
 where
 typeOf' h' = typeOf h' mp

-- | An entry in the Map.
type UeMap = (Hash, M.Bimap Int UeElem)

-- | Wrapped in the State Monad.
type UeState a = State UeMap a

-- | Get the element associated with a 'Hash' value. It's an error if the
-- element is not in the map.
getUE :: Hash -> UeMap -> UeElem
getUE h (_,mp) =
 case M.lookup h mp of
 Nothing -> error $ "Error looking up hash " ++ show h ++ " in the UE map\n" ++ show mp
 Just e -> e

-- | Put a new 'UE' in the map, unless it's already in there, and return the
-- hash pointing to the 'UE' and a new map.
newUE :: UE -> UeMap -> (Hash, UeMap)
newUE ue_ mp = runState (share ue_) mp

emptyMap :: UeMap
emptyMap = (0, M.empty)

-- | Create the sharing map.
share :: UE -> UeState Hash
share e = case e of
 UVRef (UV i j k) -> maybeUpdate (MUVRef $ MUV i j k)
 UVRef (UVExtern i j) -> maybeUpdate (MUVRef $ MUVExtern i j)
 UVRef (UVArray arr a) -> unOp a (\x -> MUVRef (MUVArray arr x))
 UConst a -> maybeUpdate (MUConst a)
 UCast t a -> unOp a (MUCast t)
 UAdd a b -> binOp (a,b) MUAdd
 USub a b -> binOp (a,b) MUSub
 UMul a b -> binOp (a,b) MUMul
 UDiv a b -> binOp (a,b) MUDiv
 UMod a b -> binOp (a,b) MUMod
 UNot a -> unOp a MUNot
 UAnd ls -> listOp ls MUAnd
 UBWNot a -> unOp a MUBWNot
 UBWAnd a b -> binOp (a,b) MUBWAnd
 UBWOr a b -> binOp (a,b) MUBWOr
 UBWXor a b -> binOp (a,b) MUBWXor
 UBWShiftL a b -> binOp (a,b) MUBWShiftL
 UBWShiftR a b -> binOp (a,b) MUBWShiftR
 UEq a b -> binOp (a,b) MUEq
 ULt a b -> binOp (a,b) MULt
 UMux a b c -> triOp (a,b,c) MUMux
 UF2B a -> unOp a MUF2B
 UD2B a -> unOp a MUD2B
 UB2F a -> unOp a MUB2F
 UB2D a -> unOp a MUB2D
-- math.h:
 UPi -> maybeUpdate (MUPi)
 UExp a -> unOp a MUExp
 ULog a -> unOp a MULog
 USqrt a -> unOp a MUSqrt
 UPow a b -> binOp (a,b) MUPow
 USin a -> unOp a MUSin
 UAsin a -> unOp a MUAsin
 UCos a -> unOp a MUCos
 UAcos a -> unOp a MUAcos
 USinh a -> unOp a MUSinh
 UCosh a -> unOp a MUCosh
 UAsinh a -> unOp a MUAsinh
 UAcosh a -> unOp a MUAcosh
 UAtan a -> unOp a MUAtan
 UAtanh a -> unOp a MUAtanh

-- XXX I could combine some of the following functions (unOp, binOp, etc.) to
-- slightly reduce code...
unOp :: UE -> (Hash -> UeElem) -> UeState Hash
unOp e code = do
 h <- share e
 maybeUpdate (code h)

binOp :: (UE, UE) -> (Hash -> Hash -> UeElem) -> UeState Hash
binOp (e0,e1) code = do
 h0 <- share e0
 h1 <- share e1
 maybeUpdate (code h0 h1)

triOp :: (UE, UE, UE) -> (Hash -> Hash -> Hash -> UeElem) -> UeState Hash
triOp (e0,e1,e2) code = do
 h0 <- share e0
 h1 <- share e1
 h2 <- share e2
 maybeUpdate (code h0 h1 h2)

listOp :: [UE] -> ([Hash] -> UeElem) -> UeState Hash
listOp es code = do
 hashes <- foldM (\hashes e -> do h <- share e
 return (h:hashes)
 ) [] es
 maybeUpdate (code hashes)

-- | Lookup an element in the map, and if it's in there, do nothing, but return
-- its hash value. Otherwise, update the map and return the new hash value
-- for the inserted element.
maybeUpdate :: UeElem -> UeState Hash
maybeUpdate e = do
 st <- get
 let mp = snd st
 case M.lookupR e mp of
 Nothing -> do let hash = fst st + 1
 put (hash, M.insert hash e mp)
 return hash
 Just h -> return h

-- -- Lookup an elem, returning 'Nothing' if no hash exists in the map and 'Just'
-- -- the hash value otherwise.
-- getHash :: UeElem -> UeMap -> Maybe Hash
-- getHash e mp = M.lookupR e


-- ((k,e'):_) | e == e' = Just k
-- getHash e (_:es) | otherwise = getHash e es
-- getHash _ [] = Nothing

-- | Get a 'UE' back out of the 'UeMap'.
recoverUE :: UeMap -> Hash -> UE
recoverUE st h = case getUE h st of
 MUVRef (MUV i j k) -> UVRef (UV i j k)
 MUVRef (MUVArray i a) -> UVRef (UVArray i (recover' a))
 MUVRef (MUVExtern i j) -> UVRef (UVExtern i j)
 MUCast t a -> UCast t (recover' a)
 MUConst a -> UConst a
 MUAdd a b -> UAdd (recover' a) (recover' b)
 MUSub a b -> USub (recover' a) (recover' b)
 MUMul a b -> UMul (recover' a) (recover' b)
 MUDiv a b -> UDiv (recover' a) (recover' b)
 MUMod a b -> UMod (recover' a) (recover' b)
 MUNot a -> UNot (recover' a)
 MUAnd a -> UAnd $ map recover' a
 MUBWNot a -> UBWNot (recover' a)
 MUBWAnd a b -> UBWAnd (recover' a) (recover' b)
 MUBWOr a b -> UBWOr (recover' a) (recover' b)
 MUBWXor a b -> UBWXor (recover' a) (recover' b)
 MUBWShiftL a b -> UBWShiftL (recover' a) (recover' b)
 MUBWShiftR a b -> UBWShiftR (recover' a) (recover' b)
 MUEq a b -> UEq (recover' a) (recover' b)
 MULt a b -> ULt (recover' a) (recover' b)
 MUMux a b c -> UMux (recover' a) (recover' b) (recover' c)
 MUF2B a -> UF2B (recover' a)
 MUD2B a -> UD2B (recover' a)
 MUB2F a -> UB2F (recover' a)
 MUB2D a -> UB2D (recover' a)
-- math.h:
 MUPi -> UPi
 MUExp a -> UExp (recover' a)
 MULog a -> ULog (recover' a)
 MUSqrt a -> USqrt (recover' a)
 MUPow a b -> UPow (recover' a) (recover' b)
 MUSin a -> USin (recover' a)
 MUAsin a -> UAsin (recover' a)
 MUCos a -> UCos (recover' a)
 MUAcos a -> UAcos (recover' a)
 MUSinh a -> USinh (recover' a)
 MUCosh a -> UCosh (recover' a)
 MUAsinh a -> UAsinh (recover' a)
 MUAcosh a -> UAcosh (recover' a)
 MUAtan a -> UAtan (recover' a)
 MUAtanh a -> UAtanh (recover' a)
 where recover' h' = recoverUE st h'

-- | The list of Hashes to adjacent upstream of a UE.
ueUpstream :: Hash -> UeMap -> [Hash]
ueUpstream h t = case getUE h t of
 MUVRef (MUV _ _ _) -> []
 MUVRef (MUVArray _ a) -> [a]
 MUVRef (MUVExtern _ _) -> []
 MUCast _ a -> [a]
 MUConst _ -> []
 MUAdd a b -> [a, b]
 MUSub a b -> [a, b]
 MUMul a b -> [a, b]
 MUDiv a b -> [a, b]
 MUMod a b -> [a, b]
 MUNot a -> [a]
 MUAnd a -> a
 MUBWNot a -> [a]
 MUBWAnd a b -> [a, b]
 MUBWOr a b -> [a, b]
 MUBWXor a b -> [a, b]
 MUBWShiftL a b -> [a, b]
 MUBWShiftR a b -> [a, b]
 MUEq a b -> [a, b]
 MULt a b -> [a, b]
 MUMux a b c -> [a, b, c]
 MUF2B a -> [a]
 MUD2B a -> [a]
 MUB2F a -> [a]
 MUB2D a -> [a]
-- math.h:
 MUPi -> []
 MUExp a -> [a]
 MULog a -> [a]
 MUSqrt a -> [a]
 MUPow a b -> [a, b]
 MUSin a -> [a]
 MUAsin a -> [a]
 MUCos a -> [a]
 MUAcos a -> [a]
 MUSinh a -> [a]
 MUCosh a -> [a]
 MUAsinh a -> [a]
 MUAcosh a -> [a]
 MUAtan a -> [a]
 MUAtanh a -> [a]

-- | The list of all UVs that directly control the value of an expression.
nearestUVs :: Hash -> UeMap -> [MUV]
nearestUVs h mp = nub $ f h
 where
 f :: Hash -> [MUV]
 f hash = case getUE hash mp of
 (MUVRef u@(MUVArray _ h')) -> [u] ++ f h'
 (MUVRef u) -> [u]
 _ -> concatMap f $ ueUpstream hash mp

-- | All array indexing subexpressions.
arrayIndices :: Hash -> UeMap -> [(UA, Hash)]
arrayIndices h mp = nub $ f h
 where
 f :: Hash -> [(UA, Hash)]
 f hash = case getUE hash mp of
 (MUVRef (MUVArray ua h')) -> (ua, h') : f h'
 _ -> concatMap f $ ueUpstream hash mp

-- XXX can put this back after making UE map---won't be expensive.
isMathHCall :: UeElem -> Bool
isMathHCall fc =
 case fc of
 MUPi -> True
 MUExp _ -> True
 MULog _ -> True
 MUSqrt _ -> True
 MUPow _ _ -> True
 MUSin _ -> True
 MUAsin _ -> True
 MUCos _ -> True
 MUAcos _ -> True
 MUSinh _ -> True
 MUCosh _ -> True
 MUAsinh _ -> True
 MUAcosh _ -> True
 MUAtan _ -> True
 MUAtanh _ -> True
 _ -> False