implementation module Sjit.Compile

import StdEnv
import StdGeneric
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

import code from "sjit_c."

appendProgram :: !Bool !Program !JITState -> JITState
appendProgram is_main prog jitst
# new_code_ptr = append
	jitst.code_start jitst.code_len jitst.code_ptr
	jitst.mapping
	jitst.n_instr
	(encode prog)
	is_main
=
	{ jitst
	& code_ptr = new_code_ptr
	, n_instr  = jitst.n_instr + size prog
	}
where
	append :: !Int !Int !Int !Int !Int !{#Int} !Bool -> Int
	append _ _ _ _ _ _ _ = code {
		ccall jit_append "pIppIAI:p"
	}

bootstrap :: (!Program, !CompileState)
bootstrap
# (len_bs, bs_funs) = bootstrap_funs
# is = {i \\ i <- flatten [is \\ (_,is) <- header]}
=
	( is,
	{ 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)])
	bootstrap_funs = iter 0 header
	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])

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

	initJITState :: !Int -> JITState
	initJITState maxlen
	# (code_start,mapping) = init maxlen (maxlen*10)
	=
		{ n_instr    = 0
		, code_start = code_start
		, code_len   = maxlen*10
		, code_ptr   = code_start
		, mapping    = mapping
		}
	where
		init :: !Int !Int -> (!Int, !Int)
		init _ _ = code {
			ccall init_jit "II:Vpp"
		}

class toInstrs a :: !a -> [Instr]
instance toInstrs Instr where toInstrs i = [i]
instance toInstrs [a] | toInstrs a where toInstrs xs = [i \\ x <- xs, i <- toInstrs x]

gen :: !newis !CompileState -> m CompileState | Monad m & toInstrs newis
gen newis cs = pure (foldr add cs (toInstrs newis))
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 _          ->   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 <- [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
		| 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
			, blocks    = cs.blocks ++| [!is!]
			, new_block = [!!]
			, jitst     = appendProgram (f.fun_name == "main") is cs.jitst
			}
where
	expr :: !Expr !CompileState -> Either String CompileState
	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
		Nothing -> Left ("undefined variable '" +++ v +++ "'")
	expr (App f args) cs
	# args = if (args=:[]) [Int 0] args
	= foldM (flip expr) cs (reverse args) >>= \cs -> case get f cs.funs of
		Nothing -> Left ("undefined function '" +++ toString f +++ "'")
		Just f -> case length args of
			1 -> gen [Call f] cs
			n -> gen [Pop (n-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
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

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 gEncodedSize Instr
derive gEncode Instr

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