Add AST checks

1 files changed, 15 insertions, 11 deletions

diff --git a/Sil/Compile.icl b/Sil/Compile.icl
index 7857ff2..afdfe90 100644
--- a/Sil/Compile.icl
+++ b/Sil/Compile.icl
@@ -36,7 +36,8 @@ compile prog = case evalRWST (gen prog) () zero of
 :: Address :== Int
 
 :: FunctionSymbol =
-	{ fs_arity :: Int
+	{ fs_arity   :: Int
+	, fs_rettype :: Type
 	}
 
 :: CompileState =
@@ -104,7 +105,11 @@ reserveVar :: Int Name -> Gen Int
 reserveVar i n = modify (\cs -> {cs & addresses='M'.put n i cs.addresses}) *> pure (i+1)
 
 addFunction :: Function -> Gen ()
-addFunction f = modify (\cs -> {cs & symbols='M'.put f.f_name {fs_arity=length f.f_args} cs.symbols})
+addFunction f = modify (\cs -> {cs & symbols='M'.put f.f_name fs cs.symbols})
+where
+	fs = { fs_arity   = length f.f_args
+	     , fs_rettype = f.f_type
+	     }
 
 cleanup :: Gen ()
 cleanup = gets peekReturn >>= tell
@@ -142,18 +147,19 @@ where
 		modify (newReturn cleanup`) *>
 		gen f.f_code *>
 		cleanup *>
-		shrinkStack (args - 1) *>
+		modify (\cs -> {cs & stackoffset=0}) *>
 		tell ['ABC'.Rtn] *>
 		modify popReturn
 	where
 		cleanup` = case f.f_args of
-			[] -> [ 'ABC'.Annotation $ 'ABC'.DAnnot 1 []
+			[] -> [ 'ABC'.Annotation $ 'ABC'.DAnnot retSize []
 			      ]
-			_  -> [ 'ABC'.Comment "Cleanup"
-			      , 'ABC'.Update_a 0 args
-			      , 'ABC'.Pop_a args
-			      , 'ABC'.Annotation $ 'ABC'.DAnnot 1 []
+			_  -> [ 'ABC'.Comment "Cleanup"] ++
+			      [ 'ABC'.Update_a i (args+i) \\ i <- [0..retSize-1] ] ++
+			      [ 'ABC'.Pop_a args
+			      , 'ABC'.Annotation $ 'ABC'.DAnnot retSize []
 			      ]
+		retSize = typeSize f.f_type
 		args = length f.f_args
 		locals = length f.f_code.cb_init
 
@@ -171,9 +177,7 @@ where
 	where
 		cleanup` = case cb.cb_init of
 			[] -> []
-			_  -> [ 'ABC'.Update_a 0 locals
-			      , 'ABC'.Pop_a locals
-			      ]
+			_  -> [ 'ABC'.Pop_a locals ]
 		locals = length cb.cb_init
 
 instance gen Initialisation