{-# LANGUAGE OverloadedStrings #-} module Grin.HPT.Environment ( mkEnvironment , Equations , Rhs(..) , RhsValue(..) , Lhs(..) , singleton ) where import CompactString import Grin.Types import qualified Data.Map as Map import Control.Monad.RWS import Control.Monad.State import Control.Monad.Reader import Control.Monad.Writer type HeapPointer = Int data Lhs = HeapEntry HeapPointer | VarEntry Renamed deriving (Eq,Ord) data RhsValue = Extract Renamed Renamed Int | ExtractVector Renamed Int | Eval Renamed | Update Renamed Renamed | Apply Renamed Renamed | PartialApply Renamed Renamed | Ident Renamed | Fetch Renamed | Base | Heap HeapPointer | Tag Renamed NodeType Int [Rhs] | VectorTag [Rhs] deriving (Eq,Ord,Show) -- A set of possible rhs values newtype Rhs = Rhs [RhsValue] deriving (Eq,Ord,Show) singleton :: RhsValue -> Rhs singleton value = Rhs [value] type Equations = Map.Map Lhs Rhs instance Monoid Rhs where mempty = Rhs [] mappend (Rhs a) (Rhs b) = Rhs (worker a b) where worker [] lst = lst worker lst [] = lst worker (Tag tag1 nt1 missing1 args1:xs) (Tag tag2 nt2 missing2 args2:ys) = case tag1 `compare` tag2 of LT -> Tag tag1 nt1 missing1 args1 : worker xs (Tag tag2 nt2 missing2 args2:ys) GT -> Tag tag2 nt2 missing2 args2 : worker (Tag tag1 nt1 missing1 args1:xs) ys EQ -> Tag tag1 nt1 (min missing1 missing2) (zipJoin args1 args2):worker xs ys worker (y@Tag{}:ys) (x:xs) = x:worker (y:ys) xs worker (y:ys) (x@Tag{}:xs) = y:worker ys (x:xs) worker (y:ys) (x:xs) = case y `compare` x of LT -> y:worker ys (x:xs) GT -> x:worker (y:ys) xs EQ -> x:worker ys xs zipJoin :: Monoid a => [a] -> [a] -> [a] zipJoin [] [] = [] zipJoin [] lst = zipWith mappend (repeat mempty) lst zipJoin lst [] = zipWith mappend lst (repeat mempty) zipJoin (x:xs) (y:ys) = mappend x y : zipJoin xs ys type GenReader = Map.Map Renamed [Renamed] type GenM a = RWS GenReader (Endo Equations) Int a applications :: Renamed applications = Builtin "applications" updates :: Renamed updates = Builtin "updates" mkEnvironment :: Grin -> Equations mkEnvironment grin = case execRWS (setupEnvGrin grin) reader 0 of (st, eqsEndo) -> appEndo eqsEndo Map.empty where reader = Map.fromList [ (funcDefName func, funcDefArgs func) | func <- grinFunctions grin ] setupEnvGrin :: Grin -> GenM () setupEnvGrin grin = do forM_ (grinCAFs grin) $ \caf -> do hp <- store =<< processVal (cafValue caf) addEquation (VarEntry (cafName caf)) (singleton $ Heap hp) forM_ (grinFunctions grin) $ \function -> do rhs <- setupEnv (funcDefBody function) addEquation (VarEntry (funcDefName function)) rhs forM_ (zip (funcDefArgs function) [0..]) $ \(arg, n) -> addEquation (VarEntry arg) (singleton $ Extract applications (funcDefName function) n) -- FIXME: Put these in order. baseBuiltins, vectorBuiltins, unsupportedBuiltins :: [CompactString] baseBuiltins = ["<#",">#","<=#",">=#","-#","+#","*#","narrow32Int#" ,"uncheckedIShiftRA#","and#","==#", "remInt#", "noDuplicate#" ,"narrow8Word#", "writeInt8OffAddr#" ,"narrow8Int#", "byteArrayContents#","touch#" ,"uncheckedIShiftL#", "negateInt#" ,"indexCharOffAddr#","minusWord#","geWord#","eqWord#","narrow16Word#" ,"ord#","chr#","or#","narrow32Word#","uncheckedShiftL#","plusWord#" ,"uncheckedShiftRL#","neChar#","gcdInt#","narrow16Int#","timesWord#" ,"writeAddrOffAddr#","writeInt32OffAddr#","quotInt#" ,"leWord#","/=#","writeCharArray#","xor#" ] vectorBuiltins = ["unsafeFreezeByteArray#", "newAlignedPinnedByteArray#" , "word2Integer#","integer2Int#", "newPinnedByteArray#" ,"readInt8OffAddr#","readInt32OffAddr#","readAddrOffAddr#","readInt32OffAddr#"] unsupportedBuiltins = ["raise#","atomicModifyMutVar#","waitWrite#","mkWeak#","writeTVar#" ,"raiseIO#","fork#","atomically#"] setupEnv :: Expression -> GenM Rhs setupEnv (Store val) = do hp <- store =<< processVal val return $ singleton $ Heap hp setupEnv (exp :>>= bind :-> rest) = do expRhs <- setupEnv exp addEquation (VarEntry bind) expRhs setupEnv rest setupEnv (exp :>> rest) = do setupEnv exp setupEnv rest setupEnv (Unit val) = processVal val setupEnv (Case val alts) = do let valRhs = singleton $ Ident val rets <- forM alts $ \(l :> alt) -> case l of Node tag _ _ args -> do forM_ (zip [0..] args) $ \(n,arg) -> addEquation (VarEntry arg) (singleton $ Extract val tag n) setupEnv alt Vector args -> do forM_ (zip [0..] args) $ \(n,arg) -> addEquation (VarEntry arg) (singleton $ ExtractVector val n) setupEnv alt Lit{} -> setupEnv alt Variable v -> do addEquation (VarEntry v) valRhs setupEnv alt _ -> error $ "setupEnv: Invalid case: " ++ show l return $ mconcat rets setupEnv (Application External{} args) = return $ singleton (VectorTag [mempty, singleton Base]) setupEnv (Application (Builtin "eval") [arg]) = do return $ singleton (Eval arg) setupEnv (Application (Builtin "apply") [arg1, arg2]) = do addEquation (VarEntry applications) (singleton $ PartialApply arg1 arg2) return $ singleton (Apply arg1 arg2) setupEnv (Application (Builtin "update") [ptr,val]) = do addEquation (VarEntry updates) (singleton $ Update ptr val) return mempty setupEnv (Application (Builtin fn) args) | fn `elem` baseBuiltins = return $ singleton Base setupEnv (Application (Builtin fn) args) | fn `elem` vectorBuiltins = return $ singleton $ VectorTag [mempty, singleton Base] setupEnv (Application (Builtin fn) args) | fn `elem` unsupportedBuiltins = return mempty setupEnv (Application (Builtin "makeStablePtr#") [val,realworld]) = do hp <- store (singleton $ Ident val) return $ singleton $ VectorTag [mempty, singleton $ Heap hp] setupEnv (Application (Builtin "deRefStablePtr#") [ptr,realworld]) = do return $ singleton $ VectorTag [mempty, singleton $ Fetch ptr] setupEnv (Application (Builtin "unblockAsyncExceptions#") [fn, realworld]) = do return $ singleton $ Apply fn realworld setupEnv (Application (Builtin "blockAsyncExceptions#") [fn, realworld]) = do return $ singleton $ Apply fn realworld setupEnv (Application (Builtin "fetch") [a]) = return $ singleton $ Fetch a setupEnv (Application (Builtin "newArray#") [size, elt, realworld]) = do hp <- store (singleton $ Ident elt) return $ singleton $ VectorTag [mempty, singleton $ Heap hp] setupEnv (Application (Builtin "readArray#") [arr, nth, realworld]) = return $ singleton $ VectorTag [mempty, singleton $ Fetch arr] setupEnv (Application (Builtin "writeArray#") [arr, nth, elt, realworld]) = do addEquation (VarEntry updates) (singleton $ Update arr elt) return mempty setupEnv (Application (Builtin builtin) args) = error $ "unknown builtin: " ++ show builtin setupEnv (Application fn args) = do funcArgs <- lookupFuncArgs fn forM_ (zip funcArgs args) $ \(var, arg) -> addEquation (VarEntry var) (singleton $ Ident arg) return $ singleton (Ident fn) processVal :: Value -> GenM Rhs processVal (Node name nt missing args) = do case nt of FunctionNode -> do funcArgs <- lookupFuncArgs name forM_ (zip funcArgs args) $ \(funcArg,arg) -> addEquation (VarEntry funcArg) (singleton $ Ident arg) ConstructorNode -> do return () return $ singleton $ Tag name nt missing (map (singleton . Ident) args) processVal (Variable var) = return $ singleton $ Ident var processVal Lit{} = return $ singleton Base processVal Hole{} = return mempty processVal Empty = return mempty processVal (Vector vs) = return $ singleton $ VectorTag (map (singleton . Ident) vs) store :: Rhs -> GenM Int store rhs = do u <- get put $ u+1 addEquation (HeapEntry u) rhs return u addEquation :: Lhs -> Rhs -> GenM () addEquation lhs rhs = tell $ Endo $ Map.insertWith mappend lhs rhs lookupFuncArgs :: Renamed -> GenM [Renamed] lookupFuncArgs func = asks $ \funcs -> case Map.lookup func funcs of Nothing -> error $ "Couldn't find function: " ++ show func Just args -> args