Tail call optimisation

1 files changed, 139 insertions, 1 deletions

diff --git a/Optimiser.hs b/Optimiser.hs
index c4c60cb..a349803 100644
--- a/Optimiser.hs
+++ b/Optimiser.hs
@@ -1,12 +1,23 @@
+{-# LANGUAGE TupleSections #-}
 module Optimiser(optimise) where
 
 import Data.List
+import qualified Data.Map.Strict as Map
+import Data.Maybe
+import qualified Data.Set as Set
 
+import AST (Name)
 import Intermediate
 
 
 optimise :: IRProgram -> IRProgram
-optimise (IRProgram bbs gfds datas) = IRProgram (mergeBlocks bbs) gfds datas
+optimise (IRProgram bbs gfds datas) =
+    let optf = foldl (.) id
+                     [ tailCallIntro
+                     , deadStoreElim, deadBBElim gfds
+                     , map propAssigns
+                     , mergeRets, mergeBlocks]
+    in IRProgram (optf bbs) gfds datas
 
 mergeBlocks :: [BB] -> [BB]
 mergeBlocks [] = []
@@ -22,3 +33,130 @@ mergeBlocks allbbs@(BB startb _ _ : _) =
 
     hasJumpTo bid (IJmp a) = a == bid
     hasJumpTo _ _ = False
+
+mergeRets :: [BB] -> [BB]
+mergeRets bbs =
+    let rets = Map.fromList [(bid, ret) | BB bid [] ret@(IRet _) <- bbs]
+    in [case bb of
+            BB bid inss (IJmp target) | Just ret <- Map.lookup target rets ->
+                BB bid inss ret
+            _ ->
+                bb
+       | bb <- bbs]
+
+propAssigns :: BB -> BB
+propAssigns (BB bid inss term) =
+    let (state, inss') = mapFoldl propagateI Map.empty inss
+        term' = propagateT state term
+    in BB bid inss' term'
+  where
+    propagateI mp (d@(RTemp i), IAssign r) = let r' = propR mp r
+                                             in (Map.insert i r' mp, (d, IAssign r'))
+    propagateI mp (d, IAssign r)      = (mp, (d, IAssign (propR mp r)))
+    propagateI mp ins@(_, IParam _)   = (mp, ins)
+    propagateI mp ins@(_, IClosure _) = (mp, ins)
+    propagateI mp ins@(_, IData _)    = (mp, ins)
+    propagateI mp (d, ICallC r rs)    = (Map.empty, (d, ICallC (propR mp r) (map (propR mp) rs)))
+    propagateI mp (d, IAllocClo n rs) = (mp, (d, IAllocClo n (map (propR mp) rs)))
+    propagateI mp (d, IDiscard r)     = (mp, (d, IDiscard (propR mp r)))
+
+    propagateT mp (IBr r a b)   = IBr (propR mp r) a b
+    propagateT _  t@(IJmp _)    = t
+    propagateT mp (IRet r)      = IRet (propR mp r)
+    propagateT mp (ITailC r rs) = ITailC (propR mp r) (map (propR mp) rs)
+    propagateT _  t@IExit       = t
+    propagateT _  t@IUnknown    = t
+
+    propR mp ref@(RTemp i) = fromMaybe ref (Map.lookup i mp)
+    propR _ ref = ref
+
+deadBBElim :: Map.Map Name GlobFuncDef -> [BB] -> [BB]
+deadBBElim gfds bbs =
+    let callable = 0 : [bid | GlobFuncDef bid _ _ <- Map.elems gfds]
+        jumpable = concatMap outEdges bbs
+        reachable = Set.fromList (jumpable ++ callable)
+    in filter (\bb -> bidOf bb `Set.member` reachable) 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)
+    elim = alltemps Set.\\ readtemps
+
+    shouldRemove :: Instruction -> Bool
+    shouldRemove (RNone, IDiscard RNone) = True
+    shouldRemove (RNone, IDiscard (RTemp i)) = i `Set.member` elim
+    shouldRemove (RTemp i, ins) = pureIC ins && i `Set.member` elim
+    shouldRemove _ = False
+
+    pureIC :: InsCode -> Bool
+    pureIC (IAssign _)     = True
+    pureIC (IParam _)      = True
+    pureIC (IClosure _)    = True
+    pureIC (IData _)       = True
+    pureIC (IAllocClo _ _) = True
+    pureIC (ICallC _ _)    = False
+    pureIC (IDiscard _)    = False
+
+tailCallIntro :: [BB] -> [BB]
+tailCallIntro bbs = map introduce bbs
+  where
+    readInBB = map (Set.fromList . readTempsBB) bbs
+    readBefore = init $ scanl (<>) Set.empty readInBB
+    readAfter = tail $ scanr (<>) Set.empty readInBB
+    readInOthers = Map.fromList [(bid, before <> after)
+                                | (BB bid _ _, before, after) <- zip3 bbs readBefore readAfter]
+
+    introduce orig@(BB _ [] _) = orig
+    introduce orig@(BB bid inss@(_:_) term) =
+        case (last inss, term) of
+            ((RTemp i1, ICallC cl as), IRet (RTemp i2))
+                | i1 == i2
+                , i1 `Set.notMember` (readInOthers Map.! bid)
+                , i1 `notElem` concatMap (readTempsIC . snd) (init inss) ->
+                    BB bid (init inss) (ITailC cl as)
+            _ -> orig
+
+outEdges :: BB -> [Int]
+outEdges (BB _ _ term) = outEdgesT term
+
+outEdgesT :: Terminator -> [Int]
+outEdgesT (IBr _ a b) = [a, b]
+outEdgesT (IJmp a) = [a]
+outEdgesT (IRet _) = []
+outEdgesT (ITailC _ _) = []
+outEdgesT IExit = []
+outEdgesT IUnknown = []
+
+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, [])