General cleanup

3 files changed, 55 insertions, 43 deletions

diff --git a/Intermediate.hs b/Intermediate.hs
index 6a11bf0..7984176 100644
--- a/Intermediate.hs
+++ b/Intermediate.hs
@@ -1,3 +1,4 @@
+{-# LANGUAGE LambdaCase, FlexibleInstances #-}
 module Intermediate where
 
 import Data.List
@@ -65,6 +66,9 @@ instance AllRefs BB where
     allRefs (BB _ inss term) =
         sortUniq $ concatMap (allRefs . snd) inss ++ allRefs term
 
+instance AllRefs [BB] where
+    allRefs = sortUniq . concatMap allRefs
+
 instance AllRefs InsCode where
     allRefs (IAssign r) = [r]
     allRefs (IParam _) = []
@@ -93,6 +97,37 @@ outEdgesT (ITailC _ _) = []
 outEdgesT IExit = []
 outEdgesT IUnknown = []
 
+icReadTemps :: InsCode -> [Int]
+icReadTemps = \case
+    IAssign r -> onlyTemporaries [r]
+    IParam _ -> []
+    IClosure _ -> []
+    IData _ -> []
+    ICallC r rs -> onlyTemporaries (r : rs)
+    IAllocClo _ rs -> onlyTemporaries rs
+    IDiscard r -> onlyTemporaries [r]
+
+termReadTemps :: Terminator -> [Int]
+termReadTemps = \case
+    IBr r _ _ -> onlyTemporaries [r]
+    IJmp _ -> []
+    IRet r -> onlyTemporaries [r]
+    ITailC r rs -> onlyTemporaries (r : rs)
+    IExit -> []
+    IUnknown -> []
+
+bbReadTemps :: BB -> [Int]
+bbReadTemps (BB _ inss term) = onlyTemporaries (concatMap (allRefs . snd) inss ++ allRefs term)
+
+bbWrittenTemps :: BB -> [Int]
+bbWrittenTemps (BB _ inss _) = concatMap insWrittenTemps inss
+
+insWrittenTemps :: Instruction -> [Int]
+insWrittenTemps (d, _) = onlyTemporaries [d]
+
+onlyTemporaries :: [Ref] -> [Int]
+onlyTemporaries rs = [i | RTemp i <- rs]
+
 
 instance Show IRProgram where
     show (IRProgram bbs gfds datas) = intercalate "\n" $
@@ -115,9 +150,9 @@ instance Show IRProgram where
         bbannot bid = maybe "" ("entry point of " ++) (Map.lookup bid bidToName)
 
 instance Show GlobFuncDef where
-    show (GlobFuncDef bbid na []) = "BB " ++ show bbid ++ " (" ++ show na ++ ")"
+    show (GlobFuncDef bbid na []) = "GFD at BB " ++ show bbid ++ " (" ++ show na ++ ")"
     show (GlobFuncDef bbid na cs) =
-        "BB " ++ show bbid ++ " (" ++ show na ++ ") (closure slots: " ++ intercalate ", " cs ++ ")"
+        "GFD at BB " ++ show bbid ++ " (" ++ show na ++ ") (closure slots: " ++ intercalate ", " cs ++ ")"
 
 instance Show BB where
     show = genericShowBB (const "") show show
diff --git a/Optimiser.hs b/Optimiser.hs
index f09128f..42ebdb1 100644
--- a/Optimiser.hs
+++ b/Optimiser.hs
@@ -127,8 +127,8 @@ deadBBElim gfds bbs =
 deadStoreElim :: [BB] -> [BB]
 deadStoreElim bbs = [BB bid (filter (not . shouldRemove) inss) term | BB bid inss term <- bbs]
   where
-    readtemps = Set.fromList (concatMap readTempsBB bbs)
-    alltemps = readtemps <> Set.fromList (concatMap writtenTempsBB bbs)
+    readtemps = Set.fromList (concatMap bbReadTemps bbs)
+    alltemps = readtemps <> Set.fromList (concatMap bbWrittenTemps bbs)
     elim = alltemps Set.\\ readtemps
 
     shouldRemove :: Instruction -> Bool
@@ -151,7 +151,7 @@ deadStoreElim bbs = [BB bid (filter (not . shouldRemove) inss) term | BB bid ins
 tailCallIntro :: [BB] -> [BB]
 tailCallIntro bbs = map introduce bbs
   where
-    readInBB = map (Set.fromList . readTempsBB) bbs
+    readInBB = map (Set.fromList . bbReadTemps) bbs
     readBefore = init $ scanl (<>) Set.empty readInBB
     readAfter = tail $ scanr (<>) Set.empty readInBB
     readInOthers = Map.fromList [(bid, before <> after)
@@ -163,7 +163,7 @@ tailCallIntro bbs = map introduce bbs
             ((RTemp i1, ICallC cl as), IRet (RTemp i2))
                 | i1 == i2
                 , i1 `Set.notMember` (readInOthers Map.! bid)
-                , i1 `notElem` concatMap (readTempsIC . snd) (init inss) ->
+                , i1 `notElem` onlyTemporaries (concatMap (allRefs . snd) (init inss)) ->
                     BB bid (init inss) (ITailC cl as)
             _ -> orig
 
@@ -178,34 +178,5 @@ dedupDatas (IRProgram origbbs gfds datatbl) = IRProgram (map goBB origbbs) gfds
     goI (ref, IData i) = (ref, IData (valueIdx Map.! (datatbl !! i)))
     goI ins = ins
 
-readTempsBB :: BB -> [Int]
-readTempsBB (BB _ inss term) = concatMap (readTempsIC . snd) inss ++ readTempsT term
-
-writtenTempsBB :: BB -> [Int]
-writtenTempsBB (BB _ inss _) = concatMap (readTempsR . fst) inss
-
-readTempsIC :: InsCode -> [Int]
-readTempsIC (IAssign r) = readTempsR r
-readTempsIC (IParam _) = []
-readTempsIC (IClosure _) = []
-readTempsIC (IData _) = []
-readTempsIC (ICallC r rs) = readTempsR r ++ concatMap readTempsR rs
-readTempsIC (IAllocClo _ rs) = concatMap readTempsR rs
-readTempsIC (IDiscard _) = []
-
-readTempsT :: Terminator -> [Int]
-readTempsT (IBr r _ _) = readTempsR r
-readTempsT (IJmp _) = []
-readTempsT (IRet r) = readTempsR r
-readTempsT (ITailC r rs) = readTempsR r ++ concatMap readTempsR rs
-readTempsT IExit = []
-readTempsT IUnknown = []
-
-readTempsR :: Ref -> [Int]
-readTempsR (RConst _) = []
-readTempsR (RTemp i) = [i]
-readTempsR (RSClo _) = []
-readTempsR RNone = []
-
 mapFoldl :: (s -> a -> (s, b)) -> s -> [a] -> (s, [b])
 mapFoldl f s = fmap reverse . foldl' (\(s', yet) x -> fmap (: yet) (f s' x)) (s, [])
diff --git a/VM.hs b/VM.hs
index 3d78519..08e7c63 100644
--- a/VM.hs
+++ b/VM.hs
@@ -60,22 +60,28 @@ vmRunBB info state (BB _ inss term) = do
 
 vmRunInstr :: Info -> State -> Instruction -> IO State
 -- vmRunInstr _ _ ins | traceShow ins False = undefined
-vmRunInstr info@(Info _ _ datas) state@(State { sTempMap = tmap, sArgs = args, sCloVals = closure }) (dest, instr) = case instr of
-    IAssign ref -> return (assignRef state dest (findRef tmap ref))
+vmRunInstr info@(Info _ _ datas) state (dest, instr) = case instr of
+    IAssign ref ->
+        return (assignRef state dest (findRef (sTempMap state) ref))
     IParam i ->
-        if i < length args then return (assignRef state dest (args !! i))
-            else error $ show args ++ ", " ++ show i ++ ", param-out-of-range"
+        let args = sArgs state
+        in if i < length args then return (assignRef state dest (args !! i))
+               else error $ show args ++ ", " ++ show i ++ ", param-out-of-range"
     IClosure i ->
-        if i < length closure then return (assignRef state dest (closure !! i))
-            else error $ show closure ++ ", " ++ show i ++ ", closure-out-of-range"
+        let closure = sCloVals state
+        in if i < length closure then return (assignRef state dest (closure !! i))
+               else error $ show closure ++ ", " ++ show i ++ ", closure-out-of-range"
     IData i ->
         if i < length datas then return (assignRef state dest (toRunValue (datas !! i)))
             else error "data-out-of-range"
     ICallC cl as -> do
         (rv, state') <- callClosure info state cl as
         return (assignRef state' dest rv)
-    IAllocClo name clrefs -> return (assignRef state dest (RVClosure name (map (findRef tmap) clrefs)))
-    IDiscard _ -> return state
+    IAllocClo name clrefs ->
+        let cloVals = (map (findRef (sTempMap state)) clrefs)
+        in return (assignRef state dest (RVClosure name cloVals))
+    IDiscard _ ->
+        return state  -- TODO: erase temporary from state for IDiscard (RTemp i)
 
 vmRunTerm :: Info -> State -> Terminator -> IO (RunValue, State)
 vmRunTerm info@(Info bbmap _ _) state@(State { sTempMap = tmap }) term = case term of