implementation module Snug.Compile.Simulate

import StdEnv

import Control.Monad
import Control.Monad.Fail
import Control.Monad.Identity
import Control.Monad.State
import Data.Error
import Data.Functor

import MIPS.MIPS32
import Snug.Compile.ABI
import Snug.Syntax

:: SimulationState =
	{ instrs    :: ![[Instruction]]
	, hp_offset :: !Int
	, stack     :: ![StackValue]
	}

simulate :: ![StackValue] !(Simulator a) -> MaybeError String [Instruction]
simulate stack sim =
	execStateT sim initial >>= \state ->
	if (length state.stack == length stack)
		(pure (flatten (reverse state.instrs)))
		(fail "stack size changed")
where
	initial =
		{ instrs = []
		, hp_offset = 0
		, stack = stack
		}

stackSize :: Simulator Int
stackSize = gets \s -> length s.stack

add :: ![Instruction] -> Simulator ()
add is = modify \s -> {s & instrs=[is:s.instrs]}

push :: StackValue -> Simulator ()
push sv = modify \s -> {s & stack=[sv:s.stack]}

pop :: Simulator StackValue
pop =
	getState >>= \{stack} ->
	modify (\s -> {s & stack=tl stack}) $> hd stack

write_heap :: !StackValue -> Simulator StackValue
write_heap sv =
	getState >>= \{hp_offset} ->
	storeStackValue sv hp_offset HeapPtr >>|
	modify (\s -> {s & hp_offset=hp_offset+4}) $>
	SVRegOffset HeapPtr hp_offset

storeStackValue :: !StackValue !Offset !Register -> Simulator ()
storeStackValue (SVImmediate (Immediate i)) doffset dreg
	| i == 0 = add
		[ StoreWord R0 doffset dreg
		]
	| 0 <= i && i < 0x10000 = add
		[ OrImmediate TempImm R0 (Immediate i)
		, StoreWord TempImm doffset dreg
		]
	| 0 > i && i >= -0x8000 = add
		[ AddImmediate Signed TempImm R0 (Immediate i)
		, StoreWord TempImm doffset dreg
		]
	| otherwise = add
		[ LoadUpperImmediate TempImm (Immediate (i >> 16))
		, OrImmediate TempImm TempImm (Immediate (i bitand 0xffff))
		, StoreWord TempImm doffset dreg
		]
storeStackValue (SVImmediate imm=:(Address _ _)) doffset dreg = add
	[ LoadAddress TempImm imm
	, StoreWord TempImm doffset dreg
	]
storeStackValue (SVRegOffset reg 0) doffset dreg = add
	[ StoreWord reg doffset dreg
	]
storeStackValue (SVRegOffset reg offset) doffset dreg = add
	[ AddImmediate Signed TempImm reg (Immediate offset)
	, StoreWord TempImm doffset dreg
	]
storeStackValue (SVIndirect offset reg) doffset dreg = add
	[ LoadWord TempImm offset reg
	, StoreWord TempImm doffset dreg
	]

buildCons :: !Label !Int -> Simulator ()
buildCons cons nargs =
	write_heap (SVImmediate (Address 0 cons)) >>= \addr ->
	mapM_ (\_ -> pop >>= write_heap) [1..nargs] >>|
	push addr

buildThunk :: !Label !Int -> Simulator ()
buildThunk cons 0 =
	write_heap (SVImmediate (Address 0 cons)) >>= \addr ->
	modify (\s -> {s & hp_offset=s.hp_offset+4}) >>| // reserve space to overwrite with indir
	push addr
buildThunk cons nargs =
	write_heap (SVImmediate (Address 0 cons)) >>= \addr ->
	mapM_ (\_ -> pop >>= write_heap) [1..nargs] >>|
	push addr

pushBasicValue :: !BasicValue -> Simulator ()
pushBasicValue val = push (SVImmediate (Immediate imm))
where
	imm = case val of
		BVInt i -> i
		BVChar c -> toInt c

pushArg :: !Int !Int -> Simulator ()
pushArg i j =
	getState >>= \{stack} ->
	case stack !! i of
		SVRegOffset reg offset ->
			push (SVIndirect (offset + (j+1)*4) reg)
		_ ->
			fail "unexpected reference in pushArg\n"

indirectAndEval :: Simulator ()
indirectAndEval =
	pop >>= \sv ->
	case sv of
		SVRegOffset HeapPtr offset ->
			add
				// Build indirection
				[ LoadAddress TempImm (Address 0 (functionLabel "" NodeEntry "indir"))
				, StoreWord TempImm 0 FrontEvalPtr
				] >>|
			storeStackValue sv 4 FrontEvalPtr >>|
			getState >>= \{hp_offset} ->
			add
				// Evaluate
				[ AddImmediate Signed FrontEvalPtr HeapPtr (Immediate offset)
				, Jump NoLink (Direct (Address 0 "eval"))
				, AddImmediate Signed HeapPtr HeapPtr (Immediate hp_offset)
				]
		SVIndirect 4 FrontEvalPtr ->
			// We only need to overwrite the descriptor with an indirection
			getState >>= \{hp_offset} ->
			add
				[ LoadAddress TempImm (Address 0 (functionLabel "" NodeEntry "indir"))
				, StoreWord TempImm 0 FrontEvalPtr
				, Jump NoLink (Direct (Address 0 "eval"))
				, AddImmediate Signed HeapPtr HeapPtr (Immediate hp_offset)
				]
		_ ->
			fail "unexpected top of stack in indirect\n"