Add if construct; fib example

1 files changed, 55 insertions, 26 deletions

diff --git a/Sjit/Compile.icl b/Sjit/Compile.icl
index 4f10096..55b8ff3 100644
--- a/Sjit/Compile.icl
+++ b/Sjit/Compile.icl
@@ -2,14 +2,15 @@ implementation module Sjit.Compile
 
 import StdEnv
 import StdGeneric
-import StdMaybe
 import StdOverloadedList
 
 import Control.Applicative
 import Control.Monad
 import Data.Either
 from Data.Func import mapSt, $
+import Data.Functor
 from Data.Map import :: Map(..), get, put, newMap, fromList
+import Data.Maybe
 
 import Sjit.Syntax
 
@@ -40,13 +41,14 @@ bootstrap
 # is = {i \\ i <- flatten [is \\ (_,is) <- header]}
 =
 	( is,
-	{ vars      = newMap
-	, funs      = fromList bs_funs
-	, sp        = 0
-	, pc        = len_bs
-	, blocks    = [!is!]
-	, new_block = [!!]
-	, jitst     = appendProgram False is (initJITState 1000)
+	{ vars        = newMap
+	, funs        = fromList bs_funs
+	, sp          = 0
+	, pc          = len_bs
+	, blocks      = [!is!]
+	, new_block   = [!!]
+	, placeholder = 0
+	, jitst       = appendProgram False is (initJITState 1000)
 	})
 where
 	bootstrap_funs :: (!Int, ![(String, Int)])
@@ -94,32 +96,51 @@ where
 	add i cs = {cs & new_block=[!i:cs.new_block!], sp=sp, pc=cs.pc+1}
 	where
 		sp = cs.sp + case i of
-			PushRef _ -> 1
-			PushI _ -> 1
-			Put _ -> -1
-			Pop n -> 0-n
-			Call _ -> 1
-			JmpRelTrue _ -> 0
-			Ret -> -1
-			Halt -> -2
-			IAddRet -> -1
-			IMulRet -> -1
-			ISubRet -> -1
-			IDivRet -> -1
+			PushRef _       ->   1
+			PushI _         ->   1
+			Put _           ->  -1
+			Pop n           -> 0-n
+			Call _          ->   0
+			Jmp _           ->   0
+			JmpTrue _       ->   0
+			Ret             ->  -1
+			Halt            ->  -2
+			IAddRet         ->  -1
+			IMulRet         ->  -1
+			ISubRet         ->  -1
+			IDivRet         ->  -1
+			PlaceHolder _ n ->   n
+
+reserve :: !Int !CompileState -> m (!Int, !CompileState) | Monad m
+reserve stack_effect cs=:{placeholder=p} =
+	gen (PlaceHolder p stack_effect) {cs & placeholder=p+1} >>= \cs -> pure (p,cs)
+
+fillPlaceHolder :: !Int !Instr !CompileState -> Either String CompileState
+fillPlaceHolder p newi cs = case replace cs.new_block of
+	Nothing -> Left "internal error with placeholder"
+	Just nb -> Right {cs & new_block=nb}
+where
+	replace :: ![!Instr!] -> Maybe [!Instr!]
+	replace [!PlaceHolder n _:is!] | n==p = Just [!newi:is!]
+	replace [!i:is!] = (\is -> [!i:is!]) <$> replace is
+	replace [!!] = Nothing
 
 compile :: !Function !CompileState -> Either String CompileState
 compile f cs
+# cs & sp = 0
 # cs & funs = put f.fun_name cs.pc cs.funs
 # vars = cs.vars
-# cs & vars = foldr (uncurry put) cs.vars [(v,sp) \\ v <- f.fun_args & sp <- [cs.sp+1..]]
-= case expr f.fun_expr cs of
+# cs & vars = foldr (uncurry put) cs.vars [(v,sp) \\ v <- f.fun_args & sp <- [1..]]
+# nargs = max 1 (length f.fun_args)
+= case expr f.fun_expr cs >>= gen [Ret,Put nargs] of
 	Left e -> Left e
 	Right cs
-		# is = {i \\ i <|- Reverse [!Ret:Put (max 1 (length f.fun_args)+1):cs.new_block!]}
+		| cs.sp <> -1 -> Left ("sp was " +++ toString cs.sp +++ " after compiling '" +++ f.fun_name +++ "'")
+		# is = {i \\ i <|- Reverse cs.new_block}
 		-> Right
 			{ cs
 			& vars      = vars
-			, pc        = cs.pc+2
+			, pc        = cs.pc
 			, blocks    = cs.blocks ++| [!is!]
 			, new_block = [!!]
 			, jitst     = appendProgram (f.fun_name == "main") is cs.jitst
@@ -129,13 +150,21 @@ where
 	expr (Int i) cs = gen (PushI i) cs
 	expr (Bool b) cs = gen (PushI (if b 1 0)) cs
 	expr (Var v) cs = case get v cs.vars of
-		Just i  -> gen (PushRef (i-cs.sp)) cs
+		Just i  -> gen (PushRef (i+cs.sp)) cs
 		Nothing -> Left ("undefined variable '" +++ v +++ "'")
 	expr (App f args) cs
 	# args = if (args=:[]) [Int 0] args
-	= foldM (flip expr) {cs & sp=cs.sp+1} (reverse args) >>= \cs -> case get f cs.funs of
+	= foldM (flip expr) cs (reverse args) >>= \cs -> case get f cs.funs of
 		Nothing -> Left ("undefined function '" +++ toString f +++ "'")
 		Just f  -> gen [Pop (length args-1),Call f] cs
+	expr (If b t e) cs =
+		expr b cs >>=
+		reserve -1 >>= \(jmptrue,cs=:{sp=orgsp}) ->
+		expr e cs >>=
+		reserve 0 >>= \(jmpend,cs) ->
+		fillPlaceHolder jmptrue (JmpTrue cs.pc) {cs & sp=orgsp} >>=
+		expr t >>= \cs ->
+		fillPlaceHolder jmpend (Jmp cs.pc) cs
 
 generic gEncodedSize a :: !a -> Int
 gEncodedSize{|Int|} _ = 1