Fix functions with no arguments

1 files changed, 4 insertions, 4 deletions

diff --git a/codegen.hs b/codegen.hs
index b7af7a1..23db784 100644
--- a/codegen.hs
+++ b/codegen.hs
@@ -284,7 +284,9 @@ genFunctionBlock bl args = do
     bodybb <- genBlock bl termbb
     changeBlock firstbb
     setTerminator $ A.Br (A.Name bodybb) []
-    return firstbb
+    if length args > 0
+        then return firstbb
+        else return bodybb
 
 genBlock :: Block
          -> LLName  -- name of BasicBlock following this Block
@@ -562,9 +564,7 @@ commonTypeM t1 t2 = maybe err return $ commonType t1 t2
 
 cleanupTrampolines :: CGMonad ()
 cleanupTrampolines = do
-    st <- get
-    let newblocks = go (allBlocks st)
-    put $ st {allBlocks = newblocks}
+    state $ \s -> ((), s {allBlocks = go (allBlocks s)})
   where
     go :: Map.Map LLName A.BasicBlock -> Map.Map LLName A.BasicBlock
     go bbs = folder bbs (Map.toList bbs)