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 "isEmpty" (LAB "isEmpty" 1 3)
	$ addFuncs fd 4
	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"
					,"isEmpty":map (\(FunDecl _ k _ _ _ _)->k) x] 0] ++ context
	where

		fS :: [String] Int -> SSMProgram
		fS [] _ = []
		fS [k:xs] n = [
			 Lab $ "1next" +++ toString n
			,Instr "ldr" [Raw "R5"] ""
			,Instr "ldc" [Lit n] $ "branch to: " +++ k
			,Instr "eq" [] ""
			,if (isEmpty xs)
				(Instr "nop" [] "")
				(Instr "brf" [L $ "1next" +++ (toString $ n + 1)] "")
			,Instr "bra" [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" [] ""
			,Instr "ajs" [Lit -1] ""]
		Nothing = liftT $ Left $ Error "PANIC: unresolver variable expr"
    g (LambdaExpr _ _ _) = liftT $ Left $ Error "PANIC: Lambdas should be unfolded"
    g (FunExpr _ k es fs) = funnyStuff k es fs

funnyStuff :: String [Expr] [FieldSelector] -> Gen ()
funnyStuff k es fs = getAdressbook >>= \ab->case 'Map'.get k ab of
	//Identifier points to function
	Just (LAB l arity fn) = if (arity <> (length es))
		//Function is not complete
		(	mapM_ g es
			>>| tell
				[Instr "ldc" [Lit fn] "Store function number"
				,Instr "sth" [] ""
				,Instr "str" [Raw "R7"] ""
				,Instr "ldc" [Lit $ length es] "Store arity"
				,Instr "sth" [] ""
				,Instr "ajs" [Lit -1] ""
				]
			>>| if (isEmpty es) (pure ()) (tell 
				[Instr "stmh" [Lit $ length es] "Store arguments"
				,Instr "ajs" [Lit -1] ""
				,Instr "ldr" [Raw "R7"] ""]))
		//Function is complete
		(	mapM_ g es
			>>| getAdressbook >>= \ab->(case 'Map'.get k ab of
				Nothing = liftT (Left $ Error $ concat ["PANIC: ", k, " not found as function"])
				Just (LAB t _ _) = tell [Instr "bsr" [L t] (k +++"()")]
				Just (ADDR t arity) = liftT (Left $ Error "NO ADDRESS JUMPING FFS")
			)
    		>>| 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))
		//Function is still  not complete
		( fresh >>= \finish->fresh >>= \start->tell [
			//Store function number
			 Instr "ldl" [Lit t] "STARTING HIGHER ORDER UPDATE"
			,Instr "ldh" [Lit 0] "get function number"
			,Instr "sth" [] "Store"
			//Store function arity
			,Instr "ldl" [Lit t] "get pointer again"
			,Instr "ldh" [Lit 1] "get function arity"
			,Instr "ldc" [Lit $ length es] "add argument number"
			,Instr "add" [] "add"
			,Instr "sth" [] "Store"
			,Instr "ajs" [Lit -1] "Adjust pointer"
			//load the arguments
			,Instr "ldl" [Lit t] ""
			,Instr "ldh" [Lit 1] "Load available arguments"
			,Instr "str" [Raw "R5"] "Store available args in register"
			,Instr "ldc" [Lit 0] "Store offset"
			,Instr "str" [Raw "R6"] "Store offset in register"
			,Lab start
			,Instr "ldr" [Raw "R5"] ""
			,Instr "ldc" [Lit 0] ""
			,Instr "eq" [] ""
			,Instr "brt" [L finish] "Done pushing arg, bye"
			//Load heapadress
			,Instr "ldl" [Lit t] ""
			,Instr "ldr" [Raw "R6"] ""
			,Instr "add" [] "Corrected heapaddress"
			,Instr "ldh" [Lit 2] "Load argument"
			,Instr "sth" [] "And store it immediatly after"
			//Decrease available arguments
			,Instr "ldr" [Raw "R5"] ""
			,Instr "ldc" [Lit 1] ""
			,Instr "sub" [] ""
			,Instr "str" [Raw "R5"] ""
			//Increase available arguments
			,Instr "ldr" [Raw "R6"] ""
			,Instr "ldc" [Lit 1] ""
			,Instr "add" [] ""
			,Instr "str" [Raw "R6"] ""
			,Instr "bra" [L start] ""
			,Lab finish
			]
		  >>| mapM_ g es
		  >>| tell 
				[Instr "stmh" [Lit $ length es] "Store extra args"
				,Instr "ajs" [Lit -1] ""]
		)
		//Function is complete
		( fresh >>= \finish->fresh >>= \start->tell [
			 Instr "ldl" [Lit t] "STARTING HIGHER ORDER CALL"
			,Instr "ldh" [Lit 1] "Load available arguments"
			,Instr "str" [Raw "R5"] "Store available args in register"
			,Instr "ldc" [Lit 0] "Store offset"
			,Instr "str" [Raw "R6"] "Store offset in register"
		
			,Lab start
			,Instr "ldr" [Raw "R5"] ""
			,Instr "ldc" [Lit 0] ""
			,Instr "eq" [] ""
			,Instr "brt" [L finish] "Done pushing arg, bye"
			//Load heapadress
			,Instr "ldl" [Lit t] ""
			,Instr "ldr" [Raw "R6"] ""
			,Instr "add" [] "Corrected heapaddress"
			,Instr "ldh" [Lit 2] "Load argument"
			//Decrease available arguments
			,Instr "ldr" [Raw "R5"] ""
			,Instr "ldc" [Lit 1] ""
			,Instr "sub" [] ""
			,Instr "str" [Raw "R5"] ""
			//Increase available arguments
			,Instr "ldr" [Raw "R6"] ""
			,Instr "ldc" [Lit 1] ""
			,Instr "add" [] ""
			,Instr "str" [Raw "R6"] ""
			,Instr "bra" [L start] ""
			,Lab finish
			]
		  >>| mapM_ g es
		  >>| tell
			[Instr "ldl" [Lit t] ""
			,Instr "ldh" [Lit 0] "Get function number"
			,Instr "str" [Raw "R5"] ""
			,Instr "bsr" [L "1func"] "HIGHER ORDER END"
			,Instr "ldl" [Lit t] ""
			,Instr "ldh" [Lit 1] ""
			,Instr "neg" [] ""
			,Instr "ldr" [Raw "SP"] ""
			,Instr "add" [] ""
			,Instr "ldc" [Lit $ length es + 1] ""
			,Instr "sub" [] ""
			,Instr "str" [Raw "SP"] ""
			,Instr "ldr" [Raw "RR"] ""
			]
		)
	Nothing = liftT (Left $ Error $ "PANIC: Undefined function: " +++ k)

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) = case fs of 
                []  = tell [Instr "stl" [Lit t] ""]
                _   = tell [Instr "ldl" [Lit t] ""]
                        >>| followFs fs
                        >>| tell [Instr "sta" [Lit 0] ""]
    g (FunStmt k es fs) = funnyStuff k es fs
    g (ReturnStmt Nothing) = tell [Instr "unlink" [] ""]
		>>| tell [Instr "ret" [] ""]
    g (ReturnStmt (Just e)) = g e
		>>| tell [Instr "str" [Raw "RR"] ""] 
		>>| g (ReturnStmt Nothing)

//expects the heap address to the var to slect on to be on the stack
//and leaves the heap adress to write to on the stack
followFs :: [FieldSelector] -> Gen()
followFs [] = tell []
followFs [FieldHd:fs]   = tell [Instr "ldc" [Lit 1] "select hd"
                             ,Instr "sub" [] "select hd"]
                            >>| followFs fs
followFs [FieldTl]      = tell []
followFs [FieldTl:fs]   = tell [Instr "lda" [Lit 0] ""]
                            >>| followFs fs
followFs [FieldFst:fs]  = tell [Instr "ldc" [Lit 1] "select fst"
                                ,Instr "sub" [] "select fst"]
                            >>| followFs fs
followFs [FieldSnd]     = tell []
followFs [FieldSnd:fs]  = tell [Instr "lda" [Lit 0] ""]
                            >>| followFs fs

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 ts xs ab)
addVars t [x] = \ab->
	extend x (ADDR -2 0) 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 "black", 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