implementation module Sil.Compile

import StdList
import StdString

import Control.Applicative
import Control.Monad
import Control.Monad.RWST
import Control.Monad.Trans
import Data.Error
from Data.Func import $
import qualified Data.Map as M
import Data.Maybe
import Data.Monoid
from Text import <+

import qualified ABC.Assembler as ABC

import Sil.Syntax
import Sil.Util

instance toString CompileError
where
	toString (UndefinedName n) = "Undefined name '" <+ n <+ "'."
	toString VariableLabel     = "Variable stored at label."
	toString FunctionOnStack   = "Function stored on the stack."
	toString UnknownError      = "Unknown error."

compile :: Program -> MaybeError CompileError 'ABC'.Assembler
compile prog = case evalRWST (gen prog) () zero of
	Error e  -> Error e
	Ok (_,p) -> Ok p

:: Address :== Int

:: FunctionSymbol =
	{ fs_arity :: Int
	}

:: CompileState =
	{ labels    :: ['ABC'.Label]
	, addresses :: 'M'.Map Name Address
	, symbols   :: 'M'.Map Name FunctionSymbol
	, returns   :: ['ABC'.Assembler]
	}

instance zero CompileState
where
	zero =
		{ labels    = ["_l" <+ i \\ i <- [0..]]
		, addresses = 'M'.newMap
		, symbols   = 'M'.newMap
		, returns   = []
		}

labels :: CompileState -> ['ABC'.Label]
labels cs = cs.labels

addresses :: CompileState -> 'M'.Map Name Address
addresses cs = cs.addresses

symbols :: CompileState -> 'M'.Map Name FunctionSymbol
symbols cs = cs.symbols

newReturn :: 'ABC'.Assembler CompileState -> CompileState
newReturn ret cs = {cs & returns=[ret:cs.returns]}

addToReturn :: 'ABC'.Assembler CompileState -> CompileState
addToReturn ret cs=:{returns=[r:rs]} = {cs & returns=[ret ++ r:rs]}

removeFromReturn :: Int CompileState -> CompileState
removeFromReturn i cs=:{returns=[r:rs]} = {cs & returns=[drop i r:rs]}

popReturn :: CompileState -> CompileState
popReturn cs = {cs & returns=tl cs.returns}

peekReturn :: CompileState -> 'ABC'.Assembler
peekReturn cs = hd cs.returns

:: Gen a :== RWST () 'ABC'.Assembler CompileState (MaybeError CompileError) a

fresh :: Gen 'ABC'.Label
fresh = gets labels
	>>= \labs -> modify (\cs -> {cs & labels=tl labs})
	*> pure (hd labs)

reserveVar :: Int Name -> Gen Int
reserveVar i n = modify (\cs -> {cs & addresses='M'.put n i cs.addresses}) *> pure (i+1)

addFunction :: Function -> Gen ()
addFunction f = modify (\cs -> {cs & symbols='M'.put f.f_name {fs_arity=length f.f_args} cs.symbols})

cleanup :: Gen ()
cleanup = gets peekReturn >>= tell

class gen a :: a -> Gen ()

instance gen Program
where
	gen p =
		tell [ 'ABC'.Annotation $ 'ABC'.RawAnnot ["comp", "920", "01011101001"]
		     , 'ABC'.Annotation $ 'ABC'.RawAnnot ["start", "__sil_boot"]
		     , 'ABC'.Annotation $ 'ABC'.RawAnnot ["endinfo"]
		     , 'ABC'.Annotation $ 'ABC'.RawAnnot ["module", "m_sil_compiled", "\"sil_compiled\""]
		     , 'ABC'.Label "__sil_boot"
		     , 'ABC'.Create
		     , 'ABC'.Fill "_" 0 "main" 0
		     , 'ABC'.Jmp "_driver"
		     ] *>
		mapM_ addFunction p.p_funs *>
		mapM_ gen p.p_funs

instance gen Function
where
	gen f =
		tell [ 'ABC'.Annotation $ 'ABC'.OAnnot args []
		     , 'ABC'.Label f.f_name
		     ] *>
		foldM reserveVar locals [a.arg_name \\ a <- reverse f.f_args] *>
		modify (newReturn cleanup) *>
		gen f.f_code *>
		modify popReturn
	where
		cleanup = case f.f_args of
			[] -> [ 'ABC'.Annotation $ 'ABC'.DAnnot 1 []
			      ]
			_  -> [ 'ABC'.Comment "Cleanup"
			      , 'ABC'.Update_a 0 args
			      , 'ABC'.Pop_a args
			      , 'ABC'.Annotation $ 'ABC'.DAnnot 1 []
			      ]
		args = length f.f_args
		locals = length f.f_code.cb_init

instance gen CodeBlock
where
	gen cb = foldM reserveVar 0 [i.init_name \\ i <- cb.cb_init] *>
		mapM_ gen cb.cb_init *>
		modify (addToReturn cleanup`) *>
		mapM_ gen cb.cb_content *>
		cleanup *>
		tell ['ABC'.Rtn] *>
		modify (removeFromReturn $ length cleanup`)
	where
		cleanup` = case cb.cb_init of
			[] -> []
			_  -> [ 'ABC'.Update_a 0 locals
			      , 'ABC'.Pop_a locals
			      ]
		locals = length cb.cb_init

instance gen Initialisation
where
	gen init = comment ("Initialise " <+ init.init_name) *> tell ['ABC'.Create]

instance gen Statement
where
	gen st=:(Declaration n app) = gets addresses >>= \addrs -> case 'M'.get n addrs of
		Just i -> comment (toString st) *> gen app *>
		          tell ['ABC'.Update_a 0 $ i+1, 'ABC'.Pop_a 1]
		_      -> liftT $ Error $ UndefinedName n
	gen (Application app)   = comment "Application" *> gen app
	gen (Return (Just app)) = comment "Return" *> gen app *> cleanup *> tell ['ABC'.Rtn]
	gen (Return Nothing)    = comment "Return" *> cleanup *> tell ['ABC'.Rtn]
	gen (MachineStm s)      = tell ['ABC'.Raw s]

instance gen Application
where
	gen (Name n) = gets addresses >>= \addrs -> case 'M'.get n addrs of
		Just i -> tell ['ABC'.Push_a i]
		_      -> liftT $ Error $ UndefinedName n
	gen (Literal (BLit b)) = tell ['ABC'.Create, 'ABC'.FillB b 0]
	gen (Literal (ILit i)) = tell ['ABC'.Create, 'ABC'.FillI i 0]
	gen (App n args) = gets addresses >>= \addrs -> case 'M'.get n addrs of
		Just i -> liftT $ Error FunctionOnStack
		_      -> gets symbols >>= \syms -> case 'M'.get n syms of
			Just fs ->
				comment "Retrieve arguments" *> mapM gen args *>
				comment "Apply function" *>
				tell [ 'ABC'.Annotation $ 'ABC'.DAnnot fs.fs_arity []
				     , 'ABC'.Jsr n
				     , 'ABC'.Annotation $ 'ABC'.OAnnot 1 []
				     ]
			_ -> liftT $ Error $ UndefinedName n

comment :: String -> Gen ()
comment s = tell ['ABC'.Comment s]