Add if construct; fib example

5 files changed, 78 insertions, 34 deletions

diff --git a/Sjit/Compile.dcl b/Sjit/Compile.dcl
index 7cb74e3..fd4b2ee 100644
--- a/Sjit/Compile.dcl
+++ b/Sjit/Compile.dcl
@@ -12,6 +12,8 @@ from Sjit.Syntax import :: Function
 	| Pop !Int
 
 	| Call !Int
+	| Jmp !Int
+	| JmpTrue !Int
 	| Ret
 	| Halt
 
@@ -20,16 +22,19 @@ from Sjit.Syntax import :: Function
 	| ISubRet
 	| IDivRet
 
+	| PlaceHolder !Int !Int // only used during compilation
+
 :: Program :== {!Instr}
 
 :: CompileState =
-	{ vars      :: !Map String Int
-	, funs      :: !Map String Int
-	, sp        :: !Int
-	, pc        :: !Int
-	, blocks    :: ![!Program!]
-	, new_block :: ![!Instr!]
-	, jitst     :: !JITState
+	{ vars        :: !Map String Int
+	, funs        :: !Map String Int
+	, sp          :: !Int
+	, pc          :: !Int
+	, blocks      :: ![!Program!]
+	, new_block   :: ![!Instr!]
+	, placeholder :: !Int
+	, jitst       :: !JITState
 	}
 
 :: JITState =
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
diff --git a/Sjit/Run.icl b/Sjit/Run.icl
index b5858ec..e3e623f 100644
--- a/Sjit/Run.icl
+++ b/Sjit/Run.icl
@@ -23,6 +23,10 @@ where
 			[val:stack] -> exec (i+1) (take (n-1) stack ++ [val:drop n stack])
 		Pop n     -> exec (i+1) (drop n stack)
 		Call f    -> exec f [i+1:stack]
+		Jmp f     -> exec f stack
+		JmpTrue f -> case stack of
+			[0:stack] -> exec (i+1) stack
+			[_:stack] -> exec f stack
 		Ret -> case stack of
 			[ret:stack] -> exec ret stack
 		Halt -> case stack of
diff --git a/Sjit/Syntax.dcl b/Sjit/Syntax.dcl
index 6a9f056..64895f6 100644
--- a/Sjit/Syntax.dcl
+++ b/Sjit/Syntax.dcl
@@ -7,6 +7,7 @@ from Data.Either import :: Either
 	| Bool !Bool
 	| Var !String
 	| App !String ![Expr]
+	| If !Expr !Expr !Expr
 
 :: Function =
 	{ fun_name :: !String
diff --git a/Sjit/Syntax.icl b/Sjit/Syntax.icl
index 7412d7c..90aeb39 100644
--- a/Sjit/Syntax.icl
+++ b/Sjit/Syntax.icl
@@ -16,6 +16,8 @@ import Text.Parsers.Simple.Core
 	| TTrue
 	| TFalse
 
+	| TIf
+
 	| TEq
 	| TComma
 
@@ -32,6 +34,7 @@ where
 		TInt n      -> toString n
 		TTrue       -> "True"
 		TFalse      -> "False"
+		TIf         -> "if"
 		TEq         -> "="
 		TComma      -> ","
 		TParenOpen  -> "("
@@ -51,6 +54,7 @@ where
 			# tk = case n of
 				"True"  -> TTrue
 				"False" -> TFalse
+				"if"    -> TIf
 				n       -> TIdent n
 			-> lex [tk:tks] i e s
 
@@ -112,7 +116,8 @@ where
 
 	noInfix :: Parser Token Expr
 	noInfix =
-		liftM2 App ident (pToken TParenOpen *> pSepBy expr (pToken TComma) <* pToken TParenClose)
+		    liftM2 App ident (pToken TParenOpen *> pSepBy expr (pToken TComma) <* pToken TParenClose)
+		<|> liftM3 If (pToken TIf *> expr) expr expr
 		<|> Var <$> ident
 		<|> Int <$> int
 		<|> Bool <$> bool