module sjit

import StdEnv
import StdGeneric
import StdMaybe
from Data.Func import mapSt, $
from Data.Map import :: Map(..), get, put, newMap, fromList

import code from "sjit_c."

:: Expr
	= Int !Int
	| Var !String
	| Abstr ![String] !Expr
	| App !String ![Expr]

:: Function =
	{ fun_name :: !String
	, fun_expr :: !Expr
	}

:: Instr
	= PushRef !Int
	| PushI !Int
	| Put !Int
	| Pop !Int

	| Call !Int
	| Ret
	| Halt

	| IAddRet
	| IMulRet
	| ISubRet
	| IDivRet

:: Program :== {!Instr}

:: CompileState =
	{ vars :: !Map String Int
	, funs :: !Map String Int
	, sp   :: !Int
	, pc   :: !Int
	}

compile :: ![Function] -> Program
compile funs
# (len_bs, bs_funs) = bootstrap_funs
# (is,cs) = mapSt fun funs {vars=newMap, funs=fromList bs_funs, sp=0, pc=len_bs}
= case get "main" cs.funs of
	Nothing -> abort "no main function\n"
	Just m
		# bs = bootstrap m
		-> {i \\ i <- flatten [is \\ (_,is) <- bs] ++ flatten is}
where
	fun :: !Function !CompileState -> (![Instr], !CompileState)
	fun f cs
	# cs & funs = put f.fun_name cs.pc cs.funs
	# (is,cs) = expr f.fun_expr cs
	= (reverse [Ret:is], {cs & pc=cs.pc+1})

	expr :: !Expr !CompileState -> (![Instr], !CompileState)
	expr (Int i) cs = ([PushI i], {cs & sp=cs.sp+1, pc=cs.pc+1})
	expr (Var v) cs = case get v cs.vars of
		Just i  -> ([PushRef (i-cs.sp)], {cs & sp=cs.sp+1, pc=cs.pc+1})
		Nothing -> abort "undefined variable\n"
	expr (App f args) cs
	# (iss,cs) = mapSt expr args cs
	= case get f cs.funs of
		Just f  -> ([Pop (length args-1):Call f:flatten iss], {cs & sp=cs.sp+1, pc=cs.pc+2})
		Nothing -> abort "undefined function\n"
	expr (Abstr vs e) cs
	# cs & vars = foldr (uncurry put) cs.vars [(v,sp) \\ v <- vs & sp <- [cs.sp+1..]]
	# (is,cs) = expr e cs
	= ([Put (max 1 (length vs)+1):is], {cs & sp=cs.sp-1, pc=cs.pc+1})

	bootstrap_funs :: (!Int, ![(String, Int)])
	bootstrap_funs = iter 0 (bootstrap 0)
	where
		iter :: !Int ![(String, [Instr])] -> (!Int, ![(String, Int)])
		iter pc [] = (pc, [])
		iter pc [(name,is):rest]
		# fun = (name,pc)
		# (pc,funs) = iter (pc+length is) rest
		= (pc,[fun:funs])

	bootstrap :: !Int -> [(String, [Instr])]
	bootstrap main =
		[ ("_", [PushI 0,Call main,Halt])
		, ("+", [IAddRet])
		, ("*", [IMulRet])
		, ("-", [ISubRet])
		, ("/", [IDivRet])
		]

exec :: !Program -> Int
exec prog = exec 0 []
where
	sz = size prog

	exec :: !Int ![Int] -> Int
	exec i stack
	| i < 0 || i >= sz = abort "out of bounds\n"
	| otherwise = case prog.[i] of
		PushI n   -> exec (i+1) [n:stack]
		PushRef r -> exec (i+1) [stack!!r:stack]
		Put n     -> case stack of
			[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]
		Ret -> case stack of
			[ret:stack] -> exec ret stack
		Halt -> case stack of
			[r] -> r
			_   -> abort (toString (length stack) +++ " values left on stack\n")

		IAddRet -> case stack of
			[ret:a:b:stack] -> exec ret [a:a+b:stack]
		IMulRet -> case stack of
			[ret:a:b:stack] -> exec ret [a:a*b:stack]
		ISubRet -> case stack of
			[ret:a:b:stack] -> exec ret [a:a-b:stack]
		IDivRet -> case stack of
			[ret:a:b:stack] -> exec ret [a:a/b:stack]

generic gEncodedSize a :: !a -> Int
gEncodedSize{|Int|} _ = 1
gEncodedSize{|{!}|} fx xs = 1 + sum [fx x \\ x <-: xs]
gEncodedSize{|UNIT|} _ = 0
gEncodedSize{|PAIR|} fx fy (PAIR x y) = fx x + fy y
gEncodedSize{|EITHER|} fl _ (LEFT l)  = fl l
gEncodedSize{|EITHER|} _ fr (RIGHT r) = fr r
gEncodedSize{|CONS|} fx (CONS x) = fx x + 1
gEncodedSize{|OBJECT|} fx (OBJECT x) = fx x

derive gEncodedSize Instr

generic gEncode a :: !a !Int !*{#Int} -> (!Int, !*{#Int})
gEncode{|Int|} n i arr = (i+1, {arr & [i]=n})
gEncode{|{!}|} fx xs i arr = walk 0 (i+1) {arr & [i]=sz}
where
	sz = size xs

	walk ai i arr
	| ai >= sz = (i,arr)
	# (i,arr) = fx xs.[ai] i arr
	= walk (ai+1) i arr
gEncode{|UNIT|} _ i arr = (i,arr)
gEncode{|PAIR|} fx fy (PAIR x y) i arr
	# (i,arr) = fx x i arr
	= fy y i arr
gEncode{|EITHER|} fl _ (LEFT l)  i arr = fl l i arr
gEncode{|EITHER|} _ fr (RIGHT r) i arr = fr r i arr
gEncode{|CONS of {gcd_index}|} fx (CONS x) i arr = fx x (i+1) {arr & [i]=gcd_index}
gEncode{|OBJECT|} fx (OBJECT x) i arr = fx x i arr

derive gEncode Instr

encode :: !a -> *{#Int} | gEncodedSize{|*|}, gEncode{|*|} a
encode x
# (_,arr) = gEncode{|*|} x 0 (createArray (gEncodedSize{|*|} x) -1)
= arr

jit :: !Program -> Int
jit prog = jit (encode prog)
where
	jit :: !*{#Int} -> Int
	jit _ = code {
		ccall jit "A:I"
	}

Start = (exec prog, jit prog)

prog =: compile
	[ {fun_name="id",    fun_expr=Abstr ["x"] (Var "x")}
	, {fun_name="const", fun_expr=Abstr ["x","y"] (Var "x")}
	, {fun_name="main",  fun_expr=Abstr [] (App "+" [App "const" [Int 37, Int 10], App "const" [Int 5, Int 10]])}
	]