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 | ADDR Int
:: Gen a :== RWST () SSMProgram (Addressbook, [Label]) (Either GenError) a

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

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

//Current issues:
//All VarDecls are added as function, how to deal with assignments?
//  (And when we deal with assignments, how to deal with assignments to higher order functions?)
//Dealing with arguments
//Dealing with types that do not fit on the Stack
//  Probably completely change LoadPlace to a Type and a position relative to *something*
//  And where the type determines if this position is a pointer to the heap or an
//  unboxed value

//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 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
		Nothing = liftT (Left $ Error "PANIC: undefined variable")
		Just (ADDR t) = tell [Instr "ldl" [Lit t] ""]
		Just (LAB t) = liftT (Left $ Error "PANIC: variable and function name clash")
		//load k >>= \instr-> tell [instr] //note: pure is pure for list, i.e. []
    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) = mapM_ g es 
        >>| mapM_ (const $ tell [Instr "ajs" [Lit -1] ""]) es //clean up args
        >>| jump "bsr" k
    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 
	>>| 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)) >>|
        getAdressbook >>= \oldMap ->
		updateAdressbook (addVars args) >>|
        tell [Lab k] >>|
		tell [Instr "link" [Lit 0] ""] >>|
        //then generate functions for the VarDecls
        foldM foldVarDecl 1 vds  >>|
        //then the main function 
        mapM_ g stms >>|
        updateAdressbook (const oldMap) >>| pure ()

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