implementation module gen


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

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

//Instruction is an instruction, with possible arguments and a possible comment
//Or is a label
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
:: GenMap :== 'Map'.Map String LoadPlace
//completely change to either Stack, Heap, Register?
:: LoadPlace    = LDA Int | LDC Int | LDH Int | LDL Int 
                | LDR Int | LDS Int
                | FUNC Label
:: Gen a :== RWST Label SSMProgram (GenMap, [Label]) (Either GenError) a

labelStream :: [Label]
labelStream = map (\i-> concat ["lbl_", toString i]) [1..]

gen :: AST -> String
gen (AST fds) = case evalRWST prog "" ('Map'.newMap, labelStream) of
            Left (Error e) = e
            Right (_, p) = toString p
    where
        prog = tell [Instr "bra" [L "main"] ""] >>| mapM_ g fds
//gen _ = prog
//    where
//        expr = (Op2Expr zero (Op1Expr zero UnMinus (IntExpr zero 4)) BiPlus (IntExpr zero 7))
//        expr2 = (FunExpr zero "test" [IntExpr zero 4] [])
//        stmt = (IfStmt (BoolExpr zero True) [] [])
//        prog = case evalRWST (g stmt) 0 ('Map'.newMap, labelStream) of
//            Left (Error e) = abort e
//            Right (_, prog) = toString prog
//gen _ = toString    [Lab "Test"
//                    ,Instr "ldc" [Lit 1] "Eerste instructie"
//                    ,Instr "ldc" [Lit 2] "Tweede instructie"]

//TODO:
//For now in the generation we assume all vars fit on the stack...

//helper functions for the gen monad
genMap :: Gen GenMap
genMap = gets fst
changeGenMap :: (GenMap -> GenMap) -> Gen GenMap
changeGenMap f = modify (appFst f) >>| genMap

extend :: String LoadPlace GenMap -> GenMap
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 Expr where
    g (VarExpr _ (VarDef k fs)) = load k >>= \instr-> tell [instr] //note: pure is pure for list, i.e. []
    g (IntExpr _ i) = loadP (LDC i) >>= \instr-> tell [instr]
    g (CharExpr _ c) = abort "How to deal with chars?"
    g (BoolExpr _ True) = loadP (LDC TRUE) >>= \instr-> tell [instr]
    g (BoolExpr _ False) = loadP (LDC FALSE) >>= \instr-> tell [instr]
    g (Op2Expr _ e1 op e2) = g e1 >>| g e2 >>| tell [Instr (op2ins op) [] ""]
    g (Op1Expr _ UnNegation e) = g e >>| tell [Instr "not" [] ""]
    g (Op1Expr _ UnMinus e) = g e >>| tell [Instr "neg" [] ""]
    g (EmptyListExpr _) = abort "Shit, empty list expr"
    g (TupleExpr p (e1,e2)) = abort "How to deal with tuples?"
    g (FunExpr _ k es fs) = 
        mapM g es >>| //put all arguments on the stack (todo: fix argument handling!)
        jump "bsr" k >>= \instr-> 
        tell [instr] >>| //actually branch to function
        tell [Instr "ldr" [Raw "RR"] ""] //push return value on stack, todo: check for VOID

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 >>|
        abort "Shit, an assignment, figure out something with storing vars or something"
        //vars will be on stack in locals (possible pointers to heap)
    g (FunStmt _ _) = abort "CodeGen, FunStmt unused" //not used
    g (ReturnStmt Nothing)  = tell [Instr "ret" [] ""] //NOTE! Assumes only return address on stack, safe?
    g (ReturnStmt (Just e)) = 
        g e >>|
        tell [Instr "str" [Raw "RR"] ""] >>|
        g (ReturnStmt Nothing)

instance g VarDecl where
    g (VarDecl _ Nothing _ _) = liftT (Left $ Error "PANIC: untyped vardecl")
    g (VarDecl _ (Just t) k e) = 
        (\l->k+++"_"+++l) <$> fresh >>= \lbl->
        changeGenMap (extend k (FUNC lbl)) >>|
        tell [Lab lbl] >>|
        g e >>| 
        tell [Instr "str" [Raw "RR"] ""] >>|
        tell [Instr "ret" [] ""]


instance g FunDecl where
    g (FunDecl _ k _ _ vds stms) = 
        //varDecls can call the enclosing function, so first reserve a label for it 
        (\l-> if (k=="main") "main" (l+++"_"+++k)) <$> fresh >>= \lbl->
        changeGenMap (extend k (FUNC lbl)) >>|
        //then generate functions for the VarDecls
        genMap >>= \oldMap ->
        mapM_ g vds  >>|
        //then the main function 
        tell [Lab lbl] >>|
        mapM_ g stms >>|
        changeGenMap (const oldMap) >>| pure ()

op2ins :: Op2 -> String
op2ins op = case op 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?"

load :: String -> Gen Instr
load k = genMap >>= \g-> case 'Map'.member k g of
    False = liftT (Left $ Error $ concat ["PANIC: ", k, " not found in variable mapping"])
    True = loadP $ 'Map'.find k g

loadP :: LoadPlace -> Gen Instr
loadP pl = dec pl >>= \(instr, arg)-> pure $ Instr instr [arg] ""
where
    dec (LDA i) = pure ("lda", Lit i)
    dec (LDC i) = pure ("ldc", Lit i)
    dec (LDH i) = pure ("ldh", Lit i)
    dec (LDL i) = pure ("ldl", Lit i)
    dec (LDR i) = pure ("ldr", Lit i)
    dec (LDS i) = pure ("lds", Lit i)
    dec _       = liftT (Left $ Error "PANIC: trying to load non adres")

//Instruction (String), key of function to jump to    
jump :: String String -> Gen Instr
jump instr k = genMap >>= \g-> case 'Map'.member k g of
    False = liftT (Left $ Error $ concat ["PANIC: ", k, " not found as function"])
    True = dec ('Map'.find k g) >>= \lbl-> pure $ Instr instr [lbl] (k +++"()") 
where
    dec (FUNC l) = pure (L l)
    dec _ = liftT (Left $ Error "PANIC: trying to jump to non label")

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