Optimise some more

2 files changed, 39 insertions, 12 deletions

diff --git a/Intermediate.hs b/Intermediate.hs
index 6f4c360..b0f12b9 100644
--- a/Intermediate.hs
+++ b/Intermediate.hs
@@ -26,7 +26,7 @@ data Ref
     | RTemp Int
     | RSClo Name  -- static closure object of a function
     | RNone
-  deriving Eq
+  deriving (Eq, Ord)
 
 data InsCode
     = IAssign Ref
diff --git a/Optimiser.hs b/Optimiser.hs
index b874fc6..01267e2 100644
--- a/Optimiser.hs
+++ b/Optimiser.hs
@@ -1,4 +1,4 @@
-{-# LANGUAGE TupleSections #-}
+{-# LANGUAGE LambdaCase, TupleSections #-}
 module Optimiser(optimise) where
 
 import Data.Function (on)
@@ -12,18 +12,18 @@ import Intermediate
 
 
 optimise :: IRProgram -> IRProgram
-optimise prog =
-    let progoptf = foldl (.) id . reverse $
-            [ dedupDatas ]
-        optf = foldl (.) id . reverse $
+optimise (IRProgram bbs gfds datas) =
+    let optf = foldl (.) id . reverse $
             [ mergeBlocks, mergeRets
+            , map propAssigns, globalPropAssigns
+            , deadBBElim gfds, deadStoreElim, mergeRets
             , map propAssigns
-            , deadBBElim gfds, deadStoreElim
-            , mergeRets, deadBBElim gfds
-            , tailCallIntro ]
-
-        IRProgram bbs gfds datas = progoptf prog
-    in IRProgram (optf bbs) gfds datas
+            , deadBBElim gfds, deadStoreElim, mergeRets
+            , tailCallIntro
+            ]
+        progoptf = foldl (.) id . reverse $
+            [ dedupDatas ]
+    in progoptf $ IRProgram (optf bbs) gfds datas
 
 mergeBlocks :: [BB] -> [BB]
 mergeBlocks [] = []
@@ -91,6 +91,31 @@ propAssigns (BB bid inss term) =
     propR mp ref@(RTemp i) = fromMaybe ref (Map.lookup i mp)
     propR _ ref = ref
 
+globalPropAssigns :: [BB] -> [BB]
+globalPropAssigns bbs =
+    let asgmap = map ((,) <$> fst . head <*> map snd)
+                 . groupBy ((==) `on` fst)
+                 $ [pair | BB _ inss _ <- bbs, pair <- inss]
+        replacements = concatMap (\(dest, inss) -> case inss of
+                                                       [IAssign ref@(RConst _)] -> [(dest, ref)]
+                                                       [IAssign ref@(RSClo _)] -> [(dest, ref)]
+                                                       _ -> [])
+                                 asgmap
+        replMap = Map.fromList replacements
+        replace r = case Map.lookup r replMap of { Just r2 -> r2 ; Nothing -> r }
+       -- Explicitly do not replace the assignment itself; that will be
+       -- handled by deadStoreElim
+    in flip map bbs $ \(BB bid inss term) ->
+           let inss' = flip map inss $ \case
+                           (d, IAssign r) -> (d, IAssign (replace r))
+                           (d, ICallC r rs) -> (d, ICallC (replace r) (map replace rs))
+                           (d, IAllocClo n rs) -> (d, IAllocClo n (map replace rs))
+                           (d, IDiscard r) -> (d, IDiscard (replace r))
+                           ins@(_, IParam _) -> ins
+                           ins@(_, IClosure _) -> ins
+                           ins@(_, IData _) -> ins
+           in BB bid inss' term
+
 deadBBElim :: Map.Map Name GlobFuncDef -> [BB] -> [BB]
 deadBBElim gfds bbs =
     let callable = 0 : [bid | GlobFuncDef bid _ _ <- Map.elems gfds]
@@ -107,6 +132,8 @@ deadStoreElim bbs = [BB bid (filter (not . shouldRemove) inss) term | BB bid ins
 
     shouldRemove :: Instruction -> Bool
     shouldRemove (RNone, IDiscard RNone) = True
+    shouldRemove (RNone, IDiscard (RConst _)) = True
+    shouldRemove (RNone, IDiscard (RSClo _)) = True
     shouldRemove (RNone, IDiscard (RTemp i)) = i `Set.member` elim
     shouldRemove (RTemp i, ins) = pureIC ins && i `Set.member` elim
     shouldRemove _ = False