cleanupTrampolines

1 files changed, 47 insertions, 0 deletions

diff --git a/codegen.hs b/codegen.hs
index 1bdf7fd..e7adf75 100644
--- a/codegen.hs
+++ b/codegen.hs
@@ -152,6 +152,7 @@ genFunctions (Program decs) = mapM gen $ filter isDecFunction decs
     gen :: Declaration -> CGMonad A.Definition
     gen (DecFunction rettype name args body) = do
         firstbb <- genBlock' body
+        cleanupTrampolines
         blockmap <- liftM allBlocks get
         let bbs' = map snd $ filter (\x -> fst x /= firstbb) $ Map.toList blockmap
             bbs = fromJust (Map.lookup firstbb blockmap) : bbs'
@@ -276,6 +277,52 @@ castOperand orig@(A.ConstantOperand (A.C.Int s1 val)) t2@(TypeInt s2)
 castOperand _ _ = undefined
 
 
+cleanupTrampolines :: CGMonad ()
+cleanupTrampolines = do
+    st <- get
+    let newblocks = go (allBlocks st)
+    put $ st {allBlocks = newblocks}
+  where
+    go :: Map.Map LLName A.BasicBlock -> Map.Map LLName A.BasicBlock
+    go bbs = folder bbs (Map.toList bbs)
+      where
+        folder :: Map.Map LLName A.BasicBlock -> [(LLName, A.BasicBlock)] -> Map.Map LLName A.BasicBlock
+        folder whole [] = whole
+        folder whole ((name, (A.BasicBlock (A.Name name2) [] (A.Do (A.Br (A.Name dst) [])))) : _)
+            | name /= name2 = error "INTERNAL ERROR: name /= name2"
+            | otherwise = let res = eliminate name dst $ Map.delete name whole
+                          in folder res (Map.toList res)
+        folder whole (_:rest) = folder whole rest
+
+    eliminate :: LLName -> LLName -> Map.Map LLName A.BasicBlock -> Map.Map LLName A.BasicBlock
+    eliminate name dst bbs = Map.fromList $ map (\(n,bb) -> (n,goBB bb)) $ Map.toList bbs
+      where
+        goBB :: A.BasicBlock -> A.BasicBlock
+        goBB (A.BasicBlock nm instrs (A.Name n A.:= term))
+            = A.BasicBlock nm instrs (A.Name n A.:= (goT term))
+        goBB (A.BasicBlock _ _ (A.UnName _ A.:= _))
+            = undefined
+        goBB (A.BasicBlock nm instrs (A.Do term))
+            = A.BasicBlock nm instrs (A.Do (goT term))
+
+        goT :: A.Terminator -> A.Terminator
+        goT (A.CondBr cond d1 d2 []) = A.CondBr cond (changeName name dst d1)
+                                                     (changeName name dst d2) []
+        goT (A.Br d []) = A.Br (changeName name dst d) []
+        goT (A.Switch op d1 ds []) = A.Switch op (changeName name dst d1)
+                                              (map (\(c,n) -> (c, changeName name dst n)) ds) 
+                                              []
+        goT (A.IndirectBr {}) = undefined
+        goT (A.Invoke {}) = undefined
+        goT bb = bb
+
+    changeName :: LLName -> LLName -> A.Name -> A.Name
+    changeName from to (A.Name x)
+        | x == from = A.Name to
+        | otherwise = A.Name x
+    changeName _ _ (A.UnName _) = undefined
+
+
 toLLVMType :: Type -> A.Type
 toLLVMType (TypeInt s) = A.IntegerType $ fromIntegral s
 toLLVMType (TypeUInt s) = A.IntegerType $ fromIntegral s