Handle returns (this segfaults llvm-general...)

3 files changed, 32 insertions, 14 deletions

diff --git a/codegen.hs b/codegen.hs
index 3f35d43..34a757c 100644
--- a/codegen.hs
+++ b/codegen.hs
@@ -26,6 +26,7 @@ type LLName = String
 data GenState
     = GenState {currentBlock :: Maybe LLName
                ,allBlocks :: Map.Map LLName A.BasicBlock
+               ,currentFunction :: Declaration
                ,nextId :: Integer
                ,definitions :: [A.Definition]
                ,variables :: Map.Map Name (Type, LLName)
@@ -36,6 +37,7 @@ initialGenState :: GenState
 initialGenState
     = GenState {currentBlock = Nothing
                ,allBlocks = Map.empty
+               ,currentFunction = undefined
                ,nextId = 1
                ,definitions = []
                ,variables = Map.empty
@@ -90,6 +92,10 @@ setTerminator term = do
     let replace (A.BasicBlock n il _) = A.BasicBlock n il (A.Do term)
     state $ \s -> ((), s {allBlocks = Map.adjust replace (fromJust (currentBlock s)) (allBlocks s)})
 
+setCurrentFunction :: Declaration -> CGMonad ()
+setCurrentFunction dec = do
+    state $ \s -> ((), s {currentFunction = dec})
+
 setVar :: Name -> LLName -> Type -> CGMonad ()
 setVar name label t = do
     state $ \s -> ((), s {variables = Map.insert name (t, label) $ variables s})
@@ -180,7 +186,8 @@ genFunctions :: Program -> CGMonad [A.Definition]
 genFunctions (Program decs) = mapM gen $ filter isDecFunction decs
   where
     gen :: Declaration -> CGMonad A.Definition
-    gen (DecFunction rettype name args body) = do
+    gen dec@(DecFunction rettype name args body) = do
+        setCurrentFunction dec
         firstbb <- genBlock' body
         cleanupTrampolines
         blockmap <- liftM allBlocks get
@@ -234,16 +241,23 @@ genSingle (StVarDeclaration t n Nothing) following = do
 genSingle (StVarDeclaration _ _ (Just _)) _ = undefined
 genSingle (StAssignment name expr) following = do
     bb <- newBlockJump following
-    label <- genExpression expr
-    let oper = A.LocalReference (toLLVMType (fromJust (exTypeOf expr))) (A.Name label)
+    oper <- genExpression expr
     (dsttype, _) <- lookupVar name
     oper' <- castOperand oper dsttype
     ref <- variableStoreOperand name
     void $ addInstr $ A.Store False ref oper' Nothing 0 []
     return bb
+genSingle (StReturn expr) _ = do
+    bb <- newBlock
+    oper <- genExpression expr
+    rettype <- liftM (typeOf . currentFunction) get
+    oper' <- castOperand oper rettype
+    setTerminator $ A.Ret (Just oper') []
+    return bb
 genSingle _ _ = undefined
 
-genExpression :: Expression -> CGMonad LLName
+genExpression :: Expression -> CGMonad A.Operand
+genExpression (ExLit lit (Just t)) = literalToOperand lit t
 -- genExpression (ExLit (LitInt i) (Just t@(TypeInt sz))) = do
 --     aname <- getNewName "t"
 --     void $ addNamedInstr $ A.Name aname A.:= A.Alloca (toLLVMType t) Nothing 0 []
@@ -257,32 +271,32 @@ genExpression (ExBinOp bo e1 e2 (Just t)) = do
         Plus -> do
             e1op' <- castOperand e1op t
             e2op' <- castOperand e2op t
-            case t of
+            label <- case t of
                 (TypeInt _) -> addInstr $ A.Add False False e1op' e2op' []
                 (TypeUInt _) -> addInstr $ A.Add False False e1op' e2op' []
                 TypeFloat -> addInstr $ A.FAdd A.NoFastMathFlags e1op' e2op' []
                 TypeDouble -> addInstr $ A.FAdd A.NoFastMathFlags e1op' e2op' []
                 (TypePtr _) -> addInstr $ A.Add False False e1op' e2op' []
                 (TypeName _) -> undefined
+            return $ A.LocalReference (toLLVMType t) (A.Name label)
         Minus -> do
             e1op' <- castOperand e1op t
             e2op' <- castOperand e2op t
-            case t of
+            label <- case t of
                 (TypeInt _) -> addInstr $ A.Sub False False e1op' e2op' []
                 (TypeUInt _) -> addInstr $ A.Sub False False e1op' e2op' []
                 TypeFloat -> addInstr $ A.FSub A.NoFastMathFlags e1op' e2op' []
                 TypeDouble -> addInstr $ A.FSub A.NoFastMathFlags e1op' e2op' []
                 (TypePtr _) -> addInstr $ A.Sub False False e1op' e2op' []
                 (TypeName _) -> undefined
-        _ -> undefined
-genExpression _ = undefined
+            return $ A.LocalReference (toLLVMType t) (A.Name label)
+        _ -> throwError $ "Binary operator " ++ pshow bo ++ " not implemented"
+genExpression ex = throwError $ "Expression '" ++ pshow ex ++ "' not implemented"
 
 genExprArgument :: Expression -> CGMonad A.Operand
 genExprArgument expr = case expr of
         (ExLit lit (Just t)) -> literalToOperand lit t
-        _ -> do
-                name <- genExpression expr
-                return $ A.LocalReference (toLLVMType (fromJust (exTypeOf expr))) (A.Name name)
+        _ -> genExpression expr
 
 literalToOperand :: Literal -> Type -> CGMonad A.Operand
 literalToOperand (LitInt i) (TypeInt sz) = return $ A.ConstantOperand (A.C.Int (fromIntegral sz) i)
@@ -290,7 +304,7 @@ literalToOperand (LitVar n) t = do
     oper <- variableOperand n
     oper' <- castOperand oper t
     return oper'
-literalToOperand _ _ = undefined
+literalToOperand lit _ = throwError $ "Literal '" ++ pshow lit ++ "' not implemented"
 
 castOperand :: A.Operand -> Type -> CGMonad A.Operand
 castOperand orig@(A.LocalReference (A.IntegerType s1) _) t2@(TypeInt s2)
@@ -311,7 +325,7 @@ castOperand orig@(A.ConstantOperand (A.C.GlobalReference (A.IntegerType s1) _))
         return $ A.LocalReference (toLLVMType t2) (A.Name label)
     | fromIntegral s1 > s2 = throwError $ "Cannot implicitly cast '" ++ pshow (TypeInt (fromIntegral s1))
                                           ++ "' to '" ++ pshow t2 ++ "'"
-castOperand _ _ = undefined
+castOperand orig t2 = throwError $ "Cast from '" ++ show orig ++ "' to type '" ++ pshow t2 ++ "' not implemented"
 
 
 cleanupTrampolines :: CGMonad ()
diff --git a/main.hs b/main.hs
index ec5b050..3f3a81b 100644
--- a/main.hs
+++ b/main.hs
@@ -53,8 +53,11 @@ main = do
     -- putStrLn "Module:"
     -- print llvmMod
 
+    putStrLn "Calling withContext:"
     General.withContext $ \context -> do
+        putStrLn "Calling withModuleFromAST:"
         assert $ General.withModuleFromAST context llvmMod $ \genmod -> do
+            putStrLn "Calling moduleLLVMAssembly:"
             llvmasm <- General.moduleLLVMAssembly genmod
             putStr llvmasm
             putStrLn ""
diff --git a/simple.nl b/simple.nl
index 050935b..2accc6a 100644
--- a/simple.nl
+++ b/simple.nl
@@ -6,6 +6,7 @@ int g_var;
 int main(i32 argc, ptr(ptr(i8)) argv) {
 	//int i = g_var;
 	int i;
-	i = g_var + 2 + 3;
+	//i = g_var + 2 + 3;
 	int a = i + 2;
+	return a;
 }