1 files changed, 80 insertions, 9 deletions

diff --git a/Stackify.hs b/Stackify.hs
index a09c6f0..df43358 100644
--- a/Stackify.hs
+++ b/Stackify.hs
@@ -5,32 +5,103 @@ module Stackify (
 import qualified Data.Map.Strict as Map
 import qualified Data.Set as Set
 
+import AST (Name)
 import Intermediate
+import Liveness
 
 
-data Funcinfo =
-    FuncInfo { fiGFD :: GlobFuncDef
+data FuncInfo =
+    FuncInfo { fiInit :: Int
              , fiBBs :: [BB]
              , fiTemps :: [Int] }
 
+-- Note about stackification. Temporaries need to be pushed before a call
+-- if they're live after it and they could, conceivably, be wrongly mutated
+-- by the called function otherwise. This has a couple of interesting
+-- consequences:
+-- 1. No temporaries ever need to be pushed from the global context, or
+--    "main function", since they can never be wrongly mutated: that
+--    requires re-entering the function in which they were defined, but one
+--    can never re-enter the "main function".
+-- 2. No temporaries ever need to be pushed before a tail call; since no
+--    local variables are live after it (obviously). (Global variables are
+--    covered by point (1.).)
 stackify :: IRProgram -> IRProgram
 stackify (IRProgram origbbs gfds datas) =
-    let infos = [FuncInfo gfd bbs [i | RTemp i <- allRefs bbs]
-                | (gfd, bbs) <- partitionFunctions gfds origbbs]
-    -- LIVENESS ANALYSIS within the function!
+    let (mainCode, funcCodes) = partitionFunctions gfds origbbs
+        infos = FuncInfo 0 mainCode (onlyTemporaries (allRefs mainCode))
+                : [FuncInfo bid0 bbs (onlyTemporaries (allRefs bbs))
+                  | (GlobFuncDef bid0 _ _, bbs) <- funcCodes]
+        infos' = map stackifyF infos
+        resbbs = concatMap fiBBs infos'
+    in optimise (IRProgram resbbs gfds datas)
+  where
+    -- TODO: initBid and temps are actually unused, I think
+    stackifyF :: FuncInfo -> FuncInfo
+    stackifyF (FuncInfo initBid bbs temps) =
+        let livemap = liveness bbs
+            bbs' = [stackifyBB (livemap Map.! bid) bb | bb@(BB bid _ _) <- bbs]
+        in FuncInfo initBid bbs' temps
+
+    stackifyBB :: [(Set.Set Int, Set.Set Int)] -> BB -> BB
+    stackifyBB live (BB bid inss term) =
+        -- Note that no temporaries need to be pushed for a tail call, so
+        -- never at all for a terminator.
+        BB bid (concatMap (uncurry stackifyIns) (zip live inss)) term
 
-partitionFunctions :: Map.Map Name GlobFuncDef -> [BB] -> [(GlobFuncDef, [BB])]
+    stackifyIns :: (Set.Set Int, Set.Set Int) -> Instruction -> [Instruction]
+    stackifyIns (liveBefore, liveAfter) orig@(d, ins) = case ins of
+        ICallC r rs   -> let refs = map RTemp (Set.toList (liveBefore `Set.intersection` liveAfter))
+                         in if null refs
+                                then [(d, ICallC r rs)]
+                                else [(RNone, IPush refs)
+                                     ,(d,     ICallC r rs)
+                                     ,(RNone, IPop refs)]
+        IAssign _     -> [orig]
+        IParam _      -> [orig]
+        IClosure _    -> [orig]
+        IData _       -> [orig]
+        IAllocClo _ _ -> [orig]
+        IDiscard _    -> [orig]
+        IPush _       -> error "Unexpected stack operation before Stackify"
+        IPop _        -> error "Unexpected stack operation before Stackify"
+
+optimise :: IRProgram -> IRProgram
+optimise (IRProgram bbs gfds datas) =
+    IRProgram [BB bid (go inss) term | BB bid inss term <- bbs] gfds datas
+  where
+    go [] = []
+    go ((_, IPop l1) : (_, IPush l2) : inss) | l1 == l2 = go inss
+    go (ins : inss) = ins : go inss
+
+-- Returns the global main code, and for each function its local code.
+partitionFunctions :: Map.Map Name GlobFuncDef -> [BB] -> ([BB], [(GlobFuncDef, [BB])])
 partitionFunctions gfds bbs =
     let bbmap = Map.fromList [(bidOf bb, bb) | bb <- bbs]
+        mainReach = floodFill bbmap 0
         pairs = [(gfd, floodFill bbmap bid) | gfd@(GlobFuncDef bid _ _) <- Map.elems gfds]
-        accums = scanl (<>) Set.empty (map snd pairs)
+        accums = scanl (<>) mainReach (map snd pairs)
         triples = zip3 accums pairs (tail accums)
     in if all (\(from, (_, bbset), to) -> Set.size from + Set.size bbset == Set.size to) triples
             && Set.size (last accums) == length bbs
-           then [(gfd, map (bbmap Map.!) (Set.toList bbset)) | (gfd, bbset) <- pairs]
+           then (map (bbmap Map.!) (Set.toList mainReach)
+                ,[(gfd, map (bbmap Map.!) (Set.toList bbset)) | (gfd, bbset) <- pairs])
            else error "Non-partitionable BBs in partitionFunctions"
   where
-    floodFill bbmap origin = go Set.empty origin
+    floodFill bbmap origin = go (Set.singleton origin) origin
       where
         go seen at = let newseen = seen <> Set.fromList (outEdges (bbmap Map.! at))
                      in foldl go newseen (Set.toList (newseen Set.\\ seen))
+
+-- Returns, for each BB, for each instruction the set of temporaries live
+-- before and after that instruction.
+liveness :: [BB] -> Map.Map Int [(Set.Set Int, Set.Set Int)]
+liveness bbs =
+    let sets = livenessAnalysis bbs bidOf itemsOf outEdges fread fwrite
+    in Map.fromList (zip (map bidOf bbs) sets)
+  where
+    itemsOf (BB _ inss term) = map Right inss ++ [Left term]
+    fread (Right (_, ins)) = icReadTemps ins
+    fread (Left term) = termReadTemps term
+    fwrite (Right (r, _)) = onlyTemporaries [r]
+    fwrite (Left _) = []