implementation module gen

import StdMisc
import StdList
import StdOverloaded
import StdString
from StdFunc import id, const
import StdTuple
import StdEnum

import Data.Func
import qualified Data.Map as Map
import Data.List
import Data.Either
import Data.Tuple
import Data.Functor
import Data.Monoid
import Data.Maybe
import Control.Applicative
import Control.Monad
import Control.Monad.Trans
from Text import class Text(concat), instance Text String

import AST
import RWST

TRUE :== -1
FALSE :== 0
:: Instr    = Instr String [Arg] String
            | Lab Label
:: Label :== String
:: Arg = L Label | Lit Int | Raw String
:: SSMProgram :== [Instr]
:: GenError = Error String
:: Addressbook :== 'Map'.Map String Address
:: Address = LAB String Int | ADDR Int
:: Gen a :== RWST () SSMProgram (Addressbook, [Label]) (Either GenError) a

labelStream :: [Label]
labelStream = ["lbl_" +++ toString i\\i<-[1..]]

defaultAddressBook :: Addressbook 
defaultAddressBook = extend "1printint" (LAB "1printint" 1)
	$ extend "1printchar" (LAB "1printchar" 1)
	$ extend "1readchar" (LAB "1readchar" 0)
	$ extend "1readint" (LAB "1readint" 0)
	$ extend "isEmpty" (LAB "isempty" 1)
	'Map'.newMap

gen :: AST -> Either String String
gen (AST fds) = case evalRWST prog () (defaultAddressBook, labelStream) of
            Left (Error e) = Left e
            Right (_, p) = Right $ toString p
    where
        prog = tell [
			Instr "bsr" [L "main"] "",
			Instr "halt" [] ""
			] >>| tell programContext
            >>| mapM_ g fds

programContext :: SSMProgram
programContext =    [Lab "1printint"
                    ,Instr "link" [Lit 0] ""
                    ,Instr "ldl" [Lit -2] "load first argument"
					,Instr "trap" [Lit 0] "print int"
                    ,Instr "unlink" [] ""
                    ,Instr "ret" [] ""
					,Lab "1printchar"
                    ,Instr "link" [Lit 0] ""
                    ,Instr "ldl" [Lit -2] "load first argument"
					,Instr "trap" [Lit 1] "print char"
                    ,Instr "unlink" [] ""
                    ,Instr "ret" [] ""
                    ,Lab "1readint"
                    ,Instr "link" [Lit 0] ""
					,Instr "trap" [Lit 10] "read int"
                    ,Instr "str" [Raw "RR"] ""
                    ,Instr "unlink" [] ""
                    ,Instr "ret" [] ""
                    ,Lab "1readchar"
                    ,Instr "link" [Lit 0] ""
					,Instr "trap" [Lit 11] "read char"
                    ,Instr "str" [Raw "RR"] ""
                    ,Instr "unlink" [] ""
                    ,Instr "ret" [] ""
                    ,Lab "isempty"
                    ,Instr "link" [Lit 0] ""
                    ,Instr "ldl" [Lit -2] "load prt to list"
                    ,Instr "lda" [Lit 0] "derefrence ptr"
                    ,Instr "ldc" [Lit 0] ""
                    ,Instr "eq" [] "test for null pointer"
                    ,Instr "str" [Raw "RR"] ""
                    ,Instr "unlink" [] ""
                    ,Instr "ret" [] ""
                    ]

//helper functions for the gen monad
getAdressbook :: Gen Addressbook
getAdressbook = gets fst

updateAdressbook :: (Addressbook -> Addressbook) -> Gen Addressbook
updateAdressbook f = modify (appFst f) >>| getAdressbook

extend :: String Address Addressbook -> Addressbook
extend k pl g = 'Map'.put k pl g 

fresh :: Gen Label
fresh = gets snd >>= \vars-> 
        modify (appSnd $ const $ tail vars) >>| 
        pure (head vars)

class g a :: a -> Gen ()

instance g Op1 where
	g UnNegation = tell [Instr "not" [] ""]
	g UnMinus = tell [Instr "neg" [] ""]

instance g Op2 where
	g o = tell [Instr s [] ""]
		where
			s = case o of
    		    BiPlus = "add"
				BiMinus = "sub"
    		    BiTimes = "mul"
    		    BiDivide = "div"
    		    BiMod = "mod"
    		    BiEquals = "eq"
    		    BiLesser = "lt"
    		    BiGreater = "gt"
    		    BiLesserEq = "le"
    		    BiGreaterEq = "ge"
    		    BiUnEqual = "ne"
    		    BiAnd =  "and"
    		    BiOr = "or"
    		    BiCons = abort "Shit, Cons, how to deal with this?"

instance g FieldSelector where
	g FieldFst = tell [Instr "lda" [Lit 0] "fst"]
	g FieldSnd = tell [Instr "lda" [Lit 1] "snd"]
	g FieldHd = tell [Instr "lda" [Lit -1] "hd"]
	g FieldTl = tell [Instr "lda" [Lit 0] "tl"]

instance g Expr where
    g (IntExpr _ i) = tell [Instr "ldc" [Lit i] ""]
    g (CharExpr _ c) = tell [Instr "ldc" [Lit (toInt c)] ""]
	g (BoolExpr _ b) = tell [Instr "ldc" [Lit (if b TRUE FALSE)] ""]
    g (EmptyListExpr _) = tell [Instr "ldc" [Lit 0] ""]
		>>| tell [Instr "sth" [] ""]
    g (Op1Expr _ o e) = g e >>| g o
    g (Op2Expr _ e1 BiCons e2) = g e2 >>| g e1
		>>| tell [Instr "sth" [] ""]
		>>| tell [Instr "ajs" [Lit -1] ""]
		>>| tell [Instr "sth" [] ""]
    g (Op2Expr _ e1 op e2) = g e1 >>| g e2 >>| g op
    g (TupleExpr _ (e1,e2)) = g e1
		>>| tell [Instr "sth" [] ""]
		>>| g e2
		>>| tell [Instr "sth" [] ""]
		>>| tell [Instr "ajs" [Lit -1] ""]
    g (VarExpr _ (VarDef k fs)) = getAdressbook >>= \ab->case 'Map'.get k ab of
		Just (ADDR t) = tell [Instr "ldl" [Lit t] ""] >>| mapM_ g fs >>| pure ()
		_ = liftT (Left $ Error "Higher order functions not implemented")
    g (FunExpr _ k es fs) = 
		mapM_ g es
		>>| jump "bsr" k
        >>| mapM_ (const $ tell [Instr "ajs" [Lit -1] ""]) es //clean up args
		>>| tell [Instr "ldr" [Raw "RR"] ""]

jump :: String String -> Gen ()
jump instr k = getAdressbook >>= \ab->case 'Map'.get k ab of
	Nothing = liftT (Left $ Error $ concat ["PANIC: ", k, " not found as function"])
	Just (LAB t _) = tell [Instr instr [L t] (k +++"()")]
	Just (ADDR t) = liftT (Left $ Error $ "PANIC: jump should go to label")

instance g Stmt where
    g (IfStmt cond th el) = 
        fresh >>= \elseLabel->
        fresh >>= \endLabel->
        g cond >>|
        tell [Instr "brf" [L elseLabel] "branch else"] >>|
        mapM_ g th >>|
        tell [Instr "bra" [L endLabel] "branch end if"] >>|
        tell [Lab elseLabel] >>|
        mapM_ g el  >>|
        tell [Lab endLabel]
    g (WhileStmt cond th) = 
        fresh >>= \startLabel->
        fresh >>= \endLabel ->
        tell [Lab startLabel] >>|
        g cond >>|
        tell [Instr "brf" [L endLabel] "branch end while"] >>|
        mapM_ g th >>|
        tell [Instr "bra" [L startLabel] "branch start while"] >>|
        tell [Lab endLabel]
    g (AssStmt (VarDef k fs) e) = 
        g e >>| getAdressbook >>= \ab->case 'Map'.get k ab of
            Nothing = liftT (Left $ Error $ concat ["PANIC: ", k, " not found as var"])
            Just (LAB t _) = liftT (Left $ Error $ "PANIC: cannot assign to function")
            Just (ADDR t) = tell [Instr "stl" [Lit t] ""]
    g (FunStmt k es fs) = mapM_ g es 
        >>| jump "bsr" k
        >>| mapM_ (const $ tell [Instr "ajs" [Lit -1] ""]) es //clean up args
        >>| mapM_ g fs
		>>| pure ()
    g (ReturnStmt Nothing) = tell [Instr "unlink" [] ""]
		>>| tell [Instr "ret" [] ""]
    g (ReturnStmt (Just e)) = g e
		>>| tell [Instr "str" [Raw "RR"] ""] 
		>>| g (ReturnStmt Nothing)

foldVarDecl :: Int VarDecl -> Gen Int
foldVarDecl x (VarDecl _ _ k e) = g e 
    >>| annote x k
	>>| updateAdressbook (extend k (ADDR x)) >>| pure (x + 1)

addVars :: [String] -> (Addressbook -> Addressbook)
addVars [] = id
addVars [x:xs] = \ab->extend x (ADDR (-2 - (length xs))) (addVars xs ab)

instance g FunDecl where
    g (FunDecl _ k args _ vds stms) = 
        //varDecls can call the enclosing function, so first reserve a label for it 
        updateAdressbook (extend k (LAB k (length args))) >>|
        getAdressbook >>= \oldMap ->
		updateAdressbook (addVars args) >>|
        tell [Lab k] >>|
		tell [Instr "link" [Lit 0] ""] >>|
        //add the vars
        foldM foldVarDecl 1 vds  >>|
        //and the statements
        mapM_ g stms >>|
        updateAdressbook (const oldMap) >>| pure ()

annote :: Int String -> Gen ()
annote pos key = 
    tell [Instr "annote" [Raw "MP", Lit pos, Lit pos, Raw "green", Raw key] ""]

class print a :: a -> [String]

instance print Instr where
    print (Lab l) = [l, ":", "\n"]
    print (Instr i args com) = ["\t", i] ++ print args ++ [" ;", com, "\n"]

instance print [Arg] where
    print args = (map toString args)

instance toString Arg where
    toString (L l) = l
    toString (Lit int) = toString int
    toString (Raw s) = s

instance toString SSMProgram where
    toString p = concat $ intersperse " " $ map (\i-> concat $ intersperse " " $ print i) p