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

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

defaultAddressBook :: [FunDecl] -> Addressbook 
defaultAddressBook fd = extend "1printint" (LAB "1printint" 1 0)
	$ extend "1printchar" (LAB "1printchar" 1 1)
	$ extend "read" (LAB "read" 0 2)
	$ extend "1readint" (LAB "1readint" 0 3)
	$ extend "isEmpty" (LAB "isempty" 1 4)
	$ addFuncs fd 5
	where
		addFuncs [] _ = 'Map'.newMap
		addFuncs [(FunDecl _ k args _ _ _):xs] n =
			extend k (LAB k (length args) n) $ addFuncs xs (n+1)

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

programContext :: [FunDecl] -> SSMProgram
programContext x = [Lab "1func"
				:fS ["1printint" ,"1printchar"
					,"read" ,"1readint"
					,"isEmpty":map (\(FunDecl _ k _ _ _ _)->k) x] 0] ++ context
	where

		fS :: [String] Int -> SSMProgram
		fS [] _ = []
		fS [k:xs] n = [
			 Instr "lds" [Lit 0] ""
			,Instr "ldc" [Lit n] $ "branch to: " +++ k
			,Instr "eq" [] ""
			,Instr "ajs" [Lit -1] ""
			,Instr "bsr" [L k] ""
			:fS xs $ n+1]
		context :: SSMProgram
		context =   [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 "read"
                    ,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" [] ""
                    ,Lab "read"
                    ]

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 -1] "fst"]
	g FieldSnd = tell [Instr "lda" [Lit 0] "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
		>>| g e2
		>>| tell [Instr "stmh" [Lit 2] ""]
    g (VarExpr _ (VarDef k fs)) = getAdressbook >>= \ab->case 'Map'.get k ab of
		Just (ADDR t arity) = tell [Instr "ldl" [Lit t] ""] >>| mapM_ g fs >>| pure ()
		Just (LAB l _ fn) = tell 
			[Instr "ldc" [Lit fn] ""
			,Instr "sth" [] ""
			,Instr "ldc" [Lit 0] ""
			,Instr "sth" [] ""]
    g (FunExpr _ k es fs) = getAdressbook >>= \ab->case 'Map'.get k ab of
		//Identifier points to function
		Just (LAB l arity fn) = if (arity <> (length es))
			(	tell
					[Instr "ldc" [Lit fn] "Store function number"
					,Instr "sth" [] ""
					,Instr "ldc" [Lit $ length es] "Store arity"
					,Instr "sth" [] ""
					,Instr "ajs" [Lit -1] ""]
				>>| mapM_ g es
				>>| if (isEmpty es) (pure ()) (tell 
					[Instr "stmh" [Lit $ length es] "Store arguments"
					,Instr "ajs" [Lit -1] ""]))
			(	mapM_ g es
				>>| jump "bsr" k
        		>>| tell 
					[Instr "ajs" [Lit $ ~(length es)] "Clean arguments"
					,Instr "ldr" [Raw "RR"] ""])
		//Identifier points to variable, thus higher order function
		Just (ADDR t arity) = if (arity <> (length es))
			( liftT (Left $ Error "Variable is function and still not complete"))
			( liftT (Left $ Error "Variable is function and can be executed"))
		Nothing = liftT (Left $ Error "Undefined function!!!")

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 arity) = liftT (Left $ Error "Address as jump???")

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 ar) = tell [Instr "stl" [Lit t] ""]
    g (FunStmt k es fs) = mapM_ g es 
        >>| jump "bsr" k
        >>| tell [Instr "ajs" [Lit (~(length 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 _ mt k e) = g e 
    >>| annote x k
	>>| updateAdressbook (extend k (ADDR x $ arity $ fromJust mt))
	>>| pure (x + 1)

arity :: Type -> Int
arity (_ ->> x) = 1 + arity x
arity _ = 0

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

instance g FunDecl where
    g (FunDecl _ k args mt vds stms) = 
        //varDecls can call the enclosing function, so first reserve a label for it 
        getAdressbook >>= \oldMap ->
		updateAdressbook (addVars (fromJust mt) args) >>|
        tell [Lab k] >>|
		tell [Instr "link" [Lit 0] ""] >>|
        //add the vars
        foldM foldVarDecl 1 vds  >>|
        //and the statements
        mapM_ g stms >>|
		//Ugly hack to always return
		g (ReturnStmt Nothing) >>|
        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