Tail call optimisation

3 files changed, 159 insertions, 17 deletions

diff --git a/Intermediate.hs b/Intermediate.hs
index 0e181a4..95f6cc7 100644
--- a/Intermediate.hs
+++ b/Intermediate.hs
@@ -41,6 +41,7 @@ data Terminator
     = IBr Ref Int Int
     | IJmp Int
     | IRet Ref
+    | ITailC Ref [Ref]
     | IExit
     | IUnknown
   deriving Eq
@@ -91,5 +92,6 @@ instance Show Terminator where
     show (IBr r b1 b2) = "br " ++ show r ++ " " ++ show b1 ++ " " ++ show b2
     show (IJmp b) = "jmp " ++ show b
     show (IRet r) = "ret " ++ show r
+    show (ITailC r as) = "tailc " ++ show r ++ " " ++ show as
     show IExit = "exit"
     show IUnknown = "<<UNKNOWN>>"
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, [])
diff --git a/VM.hs b/VM.hs
index fb79ebd..b66cfc5 100644
--- a/VM.hs
+++ b/VM.hs
@@ -60,7 +60,7 @@ vmRunBB info state (BB _ inss term) = do
 
 vmRunInstr :: Info -> State -> Instruction -> IO State
 -- vmRunInstr _ _ ins | traceShow ins False = undefined
-vmRunInstr info@(Info bbmap gfds datas) state@(State { sTempMap = tmap, sArgs = args, sCloVals = closure }) (dest, instr) = case instr of
+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))
     IParam i ->
         if i < length args then return (assignRef state dest (args !! i))
@@ -71,21 +71,9 @@ vmRunInstr info@(Info bbmap gfds datas) state@(State { sTempMap = tmap, sArgs =
     IData i ->
         if i < length datas then return (assignRef state dest (toRunValue (datas !! i)))
             else error "data-out-of-range"
-    ICallC cl as ->
-        -- trace ("callc " ++ show (findRef tmap cl) ++ "  " ++ show (map (findRef tmap) as)) $
-        case findRef tmap cl of
-            RVClosure clname clvals -> case Map.lookup clname gfds of
-                Just (GlobFuncDef b _ _) ->
-                    let Just bb = Map.lookup b bbmap
-                    in do
-                        -- traceM ("call " ++ show cl ++ " with arguments " ++ show (map (findRef tmap) as))
-                        (rv, _) <- vmRunBB info (state { sArgs = map (findRef tmap) as, sCloVals = clvals }) bb
-                        return (assignRef state dest rv)
-                Nothing -> do
-                    -- Take 'tail as' to skip the first self-link argument
-                    (rv, state') <- vmRunBuiltin state clname (map (findRef tmap) (tail as))
-                    return (assignRef state' dest rv)
-            obj -> error $ "VM: Cannot call non-closure object: " ++ show obj
+    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
 
@@ -94,6 +82,7 @@ vmRunTerm info@(Info bbmap _ _) state@(State { sTempMap = tmap }) term = case te
     IBr ref b1 b2 -> vmRunBB info state . (bbmap !) $ if truthy (findRef tmap ref) then b1 else b2
     IJmp b -> vmRunBB info state (bbmap ! b)
     IRet ref -> return (findRef tmap ref, state)
+    ITailC cl as -> callClosure info state cl as
     IExit -> IO.ioError (IO.userError kErrorExit)
     IUnknown -> undefined
 
@@ -107,6 +96,19 @@ assignRef :: State -> Ref -> RunValue -> State
 assignRef state (RTemp i) rv = state { sTempMap = Map.insert i rv (sTempMap state) }
 assignRef _ ref _ = error $ "VM: Unexpected ref " ++ show ref ++ " assigned"
 
+callClosure :: Info -> State -> Ref -> [Ref] -> IO (RunValue, State)
+callClosure info@(Info bbmap gfds _) state@(State { sTempMap = tmap }) cl as =
+    case findRef tmap cl of
+        RVClosure clname clvals -> case Map.lookup clname gfds of
+            Just (GlobFuncDef b _ _) -> do
+                (rv, state') <- vmRunBB info (state { sArgs = map (findRef tmap) as, sCloVals = clvals })
+                                        (bbmap Map.! b)
+                return (rv, state' { sArgs = sArgs state, sCloVals = sCloVals state })
+            Nothing ->
+                -- Take 'tail as' to skip the first self-link argument
+                vmRunBuiltin state clname (map (findRef tmap) (tail as))
+        obj -> error $ "VM: Cannot call non-closure object: " ++ show obj
+
 vmRunBuiltin :: State -> Name -> [RunValue] -> IO (RunValue, State)
 -- vmRunBuiltin _ name args | trace (name ++ " " ++ show args) False = undefined
 vmRunBuiltin state "print" as = putStrLn (intercalate ", " (map printshow as)) >> return (RVList [], state)