Follow caller-save convention using stack, not full state restore

8 files changed, 159 insertions, 66 deletions

diff --git a/Intermediate.hs b/Intermediate.hs
index 7984176..a020509 100644
--- a/Intermediate.hs
+++ b/Intermediate.hs
@@ -38,6 +38,8 @@ data InsCode
     | ICallC Ref [Ref]
     | IAllocClo Name [Ref]
     | IDiscard Ref
+    | IPush [Ref]  -- pushes references on the stack; should be matched with an IPop with the same number of references
+    | IPop [Ref]
   deriving Eq
 
 data Terminator
@@ -77,6 +79,8 @@ instance AllRefs InsCode where
     allRefs (ICallC r rs) = sortUniq (r : rs)
     allRefs (IAllocClo _ rs) = sortUniq rs
     allRefs (IDiscard r) = [r]
+    allRefs (IPush rs) = sortUniq rs
+    allRefs (IPop rs) = sortUniq rs
 
 instance AllRefs Terminator where
     allRefs (IBr r _ _) = [r]
@@ -106,6 +110,8 @@ icReadTemps = \case
     ICallC r rs -> onlyTemporaries (r : rs)
     IAllocClo _ rs -> onlyTemporaries rs
     IDiscard r -> onlyTemporaries [r]
+    IPush rs -> onlyTemporaries rs
+    IPop _ -> []
 
 termReadTemps :: Terminator -> [Int]
 termReadTemps = \case
@@ -123,6 +129,7 @@ bbWrittenTemps :: BB -> [Int]
 bbWrittenTemps (BB _ inss _) = concatMap insWrittenTemps inss
 
 insWrittenTemps :: Instruction -> [Int]
+insWrittenTemps (_, IPop rs) = onlyTemporaries rs
 insWrittenTemps (d, _) = onlyTemporaries [d]
 
 onlyTemporaries :: [Ref] -> [Int]
@@ -171,6 +178,8 @@ instance Show InsCode where
     show (ICallC r as) = "callc " ++ show r ++ " " ++ show as
     show (IAllocClo name vs) = "alloc-closure \"" ++ name ++ "\" " ++ show vs
     show (IDiscard r) = "discard " ++ show r
+    show (IPush rs) = "push " ++ show rs
+    show (IPop rs) = "pop " ++ show rs
 
 instance Show Terminator where
     show (IBr r b1 b2) = "br " ++ show r ++ " " ++ show b1 ++ " " ++ show b2
diff --git a/Main.hs b/Main.hs
index 2173653..e3ed454 100644
--- a/Main.hs
+++ b/Main.hs
@@ -1,17 +1,14 @@
 module Main where
 
 import Control.Monad
-import qualified Data.Map.Strict as Map
-import qualified Data.Set as Set
 import System.Environment
 import System.Exit
 
 import Compiler
 import CompilerMacros
-import Intermediate
-import Liveness
 import Optimiser
 import Parser
+import Stackify
 import VM
 
 
@@ -20,57 +17,26 @@ usage = do
     progname <- getProgName
     putStrLn $ "Usage: " ++ progname ++ " [-h] [-ast] [-ir] [filename.lisp]"
     putStrLn $ "When no filename is given, will read from stdin."
-    putStrLn $ "  -h     Show this help"
-    putStrLn $ "  -ast   Print AST after compiler macro's"
-    putStrLn $ "  -ir    Print IR after optimisation"
+    putStrLn $ "  -h      Show this help"
+    putStrLn $ "  -ast    Print AST after compiler macro's"
+    putStrLn $ "  -irpre  Print IR after optimisation, before stackification"
+    putStrLn $ "  -ir     Print IR after optimisation and stackification"
 
-data Options = Options { optAST :: Bool, optIR :: Bool }
+data Options = Options { optAST :: Bool, optIRPre :: Bool, optIR :: Bool }
 
+-- TODO: clean this function up
 parseOptions' :: (Options, Maybe FilePath) -> [String] -> IO (Options, Maybe FilePath)
 parseOptions' pair [] = return pair
 parseOptions' _ ("-h":_) = usage >> exitSuccess
 parseOptions' (opt, fp) ("-ast":as) = parseOptions' (opt { optAST = True }, fp) as
+parseOptions' (opt, fp) ("-irpre":as) = parseOptions' (opt { optIRPre = True }, fp) as
 parseOptions' (opt, fp) ("-ir":as) = parseOptions' (opt { optIR = True }, fp) as
 parseOptions' _ (('-':a):_) = putStrLn ("Unknown option '" ++ a ++ "'") >> usage >> exitFailure
 parseOptions' (opt, Nothing) (f:as) = parseOptions' (opt, Just f) as
 parseOptions' (_, Just _) (_:_) = putStrLn "At most one filename argument expected" >> usage >> exitFailure
 
 parseOptions :: [String] -> IO (Options, Maybe FilePath)
-parseOptions = parseOptions' (Options False False, Nothing)
-
--- TODO: 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.).)
-liveness :: IRProgram -> Map.Map Int [Set.Set Int]
-liveness (IRProgram 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 (_, IAssign r)) = collect [r]
-    fread (Right (_, IParam _)) = []
-    fread (Right (_, IClosure _)) = []
-    fread (Right (_, IData _)) = []
-    fread (Right (_, ICallC r rs)) = collect (r : rs)
-    fread (Right (_, IAllocClo _ rs)) = collect rs
-    fread (Right (_, IDiscard r)) = collect [r]
-    fread (Left (IBr r _ _)) = collect [r]
-    fread (Left (IJmp _)) = []
-    fread (Left (IRet r)) = collect [r]
-    fread (Left (ITailC r rs)) = collect (r : rs)
-    fread (Left IExit) = []
-    fread (Left IUnknown) = []
-    fwrite (Right (r, _)) = collect [r]
-    fwrite (Left _) = []
-    collect rs = [i | RTemp i <- rs]
+parseOptions = parseOptions' (Options False False False, Nothing)
 
 main :: IO ()
 main = do
@@ -78,10 +44,29 @@ main = do
     source <- maybe getContents readFile mfname
 
     prog <- parseProgram mfname source >>= either (die . show) return
+
     let prog' = compilerMacros prog
     when (optAST opts) $ print prog'
+
     irprog <- either die return (compileProgram prog')
+
     let opt = optimise irprog
-    when (optIR opts) $ print opt
-    print (liveness opt)
-    vmRun opt
+    when (optIRPre opts) $ print opt
+
+    let optS = stackify opt
+    when (optIR opts) $ print optS
+
+    -- TODO: do we want to run the optimiser again now? In a situation as
+    -- follows, stuff might be inlined still:
+    --   t2 <- assign t1
+    --   push [t2]        ; could've been push [t1]
+    --   callc ...
+    --   pop [t2]
+    --   ... use t2 again ...
+
+    -- TODO: this raises the question of using liveness for optimisation.
+    -- In the example above, the pop instruction writes to t2, breaking up
+    -- its lifetime, so that the assignment becomes dead. The current
+    -- optimiser would not be able to catch this.
+
+    vmRun optS
diff --git a/Optimiser.hs b/Optimiser.hs
index 42ebdb1..c59c40e 100644
--- a/Optimiser.hs
+++ b/Optimiser.hs
@@ -20,7 +20,7 @@ optimise (IRProgram bbs gfds datas) =
             , deadBBElim gfds, deadStoreElim, mergeRets
             , map propAssigns
             , deadBBElim gfds, deadStoreElim, mergeRets
-            , tailCallIntro
+            , tailCallIntro, deadBBElim gfds
             ]
         progoptf = foldl (.) id . reverse $
             [ dedupDatas ]
@@ -81,6 +81,8 @@ propAssigns (BB bid inss term) =
     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)))
+    propagateI mp (d, IPush rs)       = (mp, (d, IPush (map (propR mp) rs)))
+    propagateI mp (d, IPop rs)        = (foldr Map.delete mp (onlyTemporaries rs), (d, IPop rs))
 
     propagateT mp (IBr r a b)   = IBr (propR mp r) a b
     propagateT _  t@(IJmp _)    = t
@@ -112,9 +114,12 @@ globalPropAssigns bbs =
                            (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))
+                           (d, IPush rs) -> (d, IPush (map replace rs))
                            ins@(_, IParam _) -> ins
                            ins@(_, IClosure _) -> ins
                            ins@(_, IData _) -> ins
+                           -- Cannot replace in an IPop, because its arguments are output parameters
+                           ins@(_, IPop _) -> ins
            in BB bid inss' term
 
 deadBBElim :: Map.Map Name GlobFuncDef -> [BB] -> [BB]
@@ -147,6 +152,8 @@ deadStoreElim bbs = [BB bid (filter (not . shouldRemove) inss) term | BB bid ins
     pureIC (IAllocClo _ _) = True
     pureIC (ICallC _ _)    = False
     pureIC (IDiscard _)    = False
+    pureIC (IPush _)       = False
+    pureIC (IPop _)        = False
 
 tailCallIntro :: [BB] -> [BB]
 tailCallIntro bbs = map introduce bbs
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 _) = []
diff --git a/VM.hs b/VM.hs
index 08e7c63..bd7cbef 100644
--- a/VM.hs
+++ b/VM.hs
@@ -25,6 +25,7 @@ data State = State
     { sTempMap :: TempMap
     , sArgs :: [RunValue] {- current arguments -}
     , sCloVals :: [RunValue] {- current closure -}
+    , sStack :: [RunValue] {- IPush/IPop stack -}
     , sHandles :: IMap.IntMap Handle
     , sUniq :: Int
     }
@@ -46,7 +47,7 @@ vmRun :: IRProgram -> IO ()
 vmRun (IRProgram bbs gfds datas) =
     let bbmap = Map.fromList [(bidOf bb, bb) | bb <- bbs]
         info = Info bbmap gfds datas
-        state = State IMap.empty [] [] (IMap.fromList [(-1, stdin), (-2, stdout), (-3, stderr)]) 0
+        state = State IMap.empty [] [] [] (IMap.fromList [(-1, stdin), (-2, stdout), (-3, stderr)]) 0
     in IO.catchIOError (void $ vmRunBB info state (head bbs)) vmErrorHandler
 
 vmErrorHandler :: IOError -> IO ()
@@ -82,6 +83,14 @@ vmRunInstr info@(Info _ _ datas) state (dest, instr) = case instr of
         in return (assignRef state dest (RVClosure name cloVals))
     IDiscard _ ->
         return state  -- TODO: erase temporary from state for IDiscard (RTemp i)
+    IPush refs ->
+        return (state { sStack = map (findRef (sTempMap state)) refs ++ sStack state })
+    IPop refs ->
+        if length (sStack state) >= length refs
+            then let (values, newStack) = splitAt (length refs) (sStack state)
+                     state' = state { sStack = newStack }
+                 in return (foldl (uncurry . assignRef) state' (zip refs values))
+            else error "VM: IPop on too-small stack"
 
 vmRunTerm :: Info -> State -> Terminator -> IO (RunValue, State)
 vmRunTerm info@(Info bbmap _ _) state@(State { sTempMap = tmap }) term = case term of
@@ -111,14 +120,7 @@ callClosure info@(Info bbmap gfds _) state@(State { sTempMap = tmap }) cl as =
             Just (GlobFuncDef b _ _) -> do
                 (rv, state') <- vmRunBB info (state { sArgs = map (findRef tmap) as, sCloVals = clvals })
                                         (bbmap Map.! b)
-                -- TODO: we restore tempmap here after the call, and that's
-                -- _really_ dodgy. I think it works because we never mutate
-                -- a reference after first assigning it, and without it
-                -- recursion fails, but it feels ugly. A proper solution
-                -- would be to further compile the IR down to something
-                -- that stores its values on a stack, so that recursion is
-                -- handled automatically.
-                return (rv, state' { sArgs = sArgs state, sCloVals = sCloVals state, sTempMap = sTempMap state })
+                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))
diff --git a/lisphs.cabal b/lisphs.cabal
index 32aebbe..131deb4 100644
--- a/lisphs.cabal
+++ b/lisphs.cabal
@@ -13,4 +13,4 @@ executable            lisp
     ghc-options:      -Wall -O2
     build-depends:    base >= 4 && < 5,
                       containers, filepath, mtl, parsec, text
-    other-modules:    AST, Compiler, CompilerMacros, DString, Intermediate, Liveness, Optimiser, Parser, Util, VM
+    other-modules:    AST, Compiler, CompilerMacros, DString, Intermediate, Liveness, Optimiser, Parser, Stackify, Util, VM
diff --git a/tests/stack-test.lisp b/tests/stack-test.lisp
new file mode 100644
index 0000000..cabbaa3
--- /dev/null
+++ b/tests/stack-test.lisp
@@ -0,0 +1,18 @@
+(define g (x) (+ x 1))
+
+(define f (x)
+    (if (<= x 0)
+        0
+        (let ((y (g x)))
+            (+ y (f (/ x 2))))))
+
+; f 10
+; = 11 + f 5
+; = 11 + 6 + f 2
+; = 11 + 6 + 3 + f 1
+; = 11 + 6 + 3 + 2
+; = 22
+(print (f 10))
+
+; Without Stackify, the deepest y value, i.e. 2, overwrites all y values above,
+; resulting in 2 + 2 + 2 + 2 = 8.
diff --git a/tests/stack-test.out b/tests/stack-test.out
new file mode 100644
index 0000000..2bd5a0a
--- /dev/null
+++ b/tests/stack-test.out
@@ -0,0 +1 @@
+22