WIP push temporaries before calls

Requires liveness analysis

4 files changed, 66 insertions, 17 deletions

diff --git a/Intermediate.hs b/Intermediate.hs
index b0f12b9..2705431 100644
--- a/Intermediate.hs
+++ b/Intermediate.hs
@@ -5,6 +5,7 @@ import qualified Data.Map.Strict as Map
 import Data.Maybe
 
 import AST
+import Util
 
 
 data IRProgram =
@@ -60,20 +61,24 @@ termOf (BB _ _ t) = t
 class AllRefs a where
     allRefs :: a -> [Ref]
 
+instance AllRefs BB where
+    allRefs (BB _ inss term) =
+        sortUniq $ concatMap (allRefs . snd) inss ++ allRefs term
+
 instance AllRefs InsCode where
     allRefs (IAssign r) = [r]
     allRefs (IParam _) = []
     allRefs (IClosure _) = []
     allRefs (IData _) = []
-    allRefs (ICallC r rs) = r : rs
-    allRefs (IAllocClo _ rs) = rs
+    allRefs (ICallC r rs) = sortUniq (r : rs)
+    allRefs (IAllocClo _ rs) = sortUniq rs
     allRefs (IDiscard r) = [r]
 
 instance AllRefs Terminator where
     allRefs (IBr r _ _) = [r]
     allRefs (IJmp _) = []
     allRefs (IRet r) = [r]
-    allRefs (ITailC r rs) = r : rs
+    allRefs (ITailC r rs) = sortUniq (r : rs)
     allRefs IExit = []
     allRefs IUnknown = []
 
diff --git a/Optimiser.hs b/Optimiser.hs
index 01267e2..626cb50 100644
--- a/Optimiser.hs
+++ b/Optimiser.hs
@@ -9,6 +9,7 @@ import qualified Data.Set as Set
 
 import AST (Name)
 import Intermediate
+import Util
 
 
 optimise :: IRProgram -> IRProgram
@@ -219,17 +220,3 @@ 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, [])
-
-uniq :: Eq a => [a] -> [a]
-uniq (x:y:zs) | x == y = uniq (y:zs)
-              | otherwise = x : uniq (y:zs)
-uniq l = l
-
-oppositeGraph :: (Show a, Ord a) => Map.Map a [a] -> Map.Map a [a]
-oppositeGraph graph =
-    let nodes = concat [k : vs | (k, vs) <- Map.assocs graph]
-        edges = map ((,) <$> fst . head <*> map snd)
-                . groupBy ((==) `on` fst)
-                . sortOn fst
-                $ [(to, from) | (from, tos) <- Map.assocs graph, to <- tos]
-    in Map.fromList (map (,[]) nodes ++ edges)
diff --git a/Stackify.hs b/Stackify.hs
new file mode 100644
index 0000000..a09c6f0
--- /dev/null
+++ b/Stackify.hs
@@ -0,0 +1,36 @@
+module Stackify (
+    stackify
+) where
+
+import qualified Data.Map.Strict as Map
+import qualified Data.Set as Set
+
+import Intermediate
+
+
+data Funcinfo =
+    FuncInfo { fiGFD :: GlobFuncDef
+             , fiBBs :: [BB]
+             , fiTemps :: [Int] }
+
+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!
+
+partitionFunctions :: Map.Map Name GlobFuncDef -> [BB] -> [(GlobFuncDef, [BB])]
+partitionFunctions gfds bbs =
+    let bbmap = Map.fromList [(bidOf bb, bb) | bb <- bbs]
+        pairs = [(gfd, floodFill bbmap bid) | gfd@(GlobFuncDef bid _ _) <- Map.elems gfds]
+        accums = scanl (<>) Set.empty (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]
+           else error "Non-partitionable BBs in partitionFunctions"
+  where
+    floodFill bbmap origin = go Set.empty origin
+      where
+        go seen at = let newseen = seen <> Set.fromList (outEdges (bbmap Map.! at))
+                     in foldl go newseen (Set.toList (newseen Set.\\ seen))
diff --git a/Util.hs b/Util.hs
new file mode 100644
index 0000000..11b5658
--- /dev/null
+++ b/Util.hs
@@ -0,0 +1,21 @@
+module Util where
+
+import qualified Data.Map.Strict as Map
+
+
+uniq :: Eq a => [a] -> [a]
+uniq (x:y:zs) | x == y = uniq (y:zs)
+              | otherwise = x : uniq (y:zs)
+uniq l = l
+
+sortUniq :: Ord a => [a] -> [a]
+sortUniq = uniq . sort
+
+oppositeGraph :: (Show a, Ord a) => Map.Map a [a] -> Map.Map a [a]
+oppositeGraph graph =
+    let nodes = concat [k : vs | (k, vs) <- Map.assocs graph]
+        edges = map ((,) <$> fst . head <*> map snd)
+                . groupBy ((==) `on` fst)
+                . sortOn fst
+                $ [(to, from) | (from, tos) <- Map.assocs graph, to <- tos]
+    in Map.fromList (map (,[]) nodes ++ edges)