implementation module Snug.Compile

import Control.Monad
import Control.Monad.Fail
import Control.Monad.State
import Control.Monad.Trans
import Data.Error
import Data.Func
import Data.Functor
import Data.List
import qualified Data.Map
from Data.Map import :: Map
import qualified Data.Set
import Data.Tuple
import StdEnv
from Text import concat3, concat4

import MIPS.MIPS32
import Snug.Compile.ABI
import Snug.Compile.Simulate
import Snug.Compile.Typing
import Snug.Syntax

:: CompileState =
	{ fresh_ident :: !Int
	, globals     :: !Globals
	}

:: LocalLocation = FrontPtrArg !Int

instance == (Namespaced id) | == id
where
	(==) x y = x.id == y.id && x.ns == y.ns

instance < (Namespaced id) | < id
where
	(<) x y
		| x.id < y.id = True
		| x.id > y.id = False
		| otherwise = x.ns < y.ns

compile :: !Namespace ![Definition] -> MaybeError String [Line]
compile ns defs =
	flatten <$> evalStateT
		(
			mapM (liftCases ns) defs <$&> flatten >>= \defs ->
			mapM (compileDefinition ns) defs
		)
		init
where
	init =
		{ fresh_ident = 0
		, globals = combineGlobals [builtin, gatherGlobals ns defs]
		}

	builtin =
		{ constructors = 'Data.Map'.fromList
			[ ({ns="", id="INT"}, ConstructorDef "INT" [])
			]
		, functions = 'Data.Map'.newMap
		}

	combineGlobals :: ![Globals] -> Globals
	combineGlobals sets =
		{ constructors = 'Data.Map'.unions [g.constructors \\ g <- sets]
		, functions = 'Data.Map'.unions [g.functions \\ g <- sets]
		}

	gatherGlobals :: !Namespace ![Definition] -> Globals
	gatherGlobals ns defs =
		{ constructors = 'Data.Map'.fromList
			[ ({ns=ns, id=id}, cons)
			\\ DataDef _ _ conses <- defs
			, cons=:(ConstructorDef id _) <- conses
			]
		, functions = 'Data.Map'.fromList
			[ ({ns=ns, id=id}, {arity=length args, type=foldr TyApp ret (map snd (reverse args))})
			\\ FunDef id args ret _ <- defs
			]
		}

addGlobalFunction :: !(Namespaced SymbolIdent) !FunctionInfo -> CompileM ()
addGlobalFunction sym fi = modify \s -> {s & globals.functions='Data.Map'.put sym fi s.globals.functions}

/**
 * This pass ensures that `Case` only appears as the toplevel expression of a
 * `FunDef`, and that the expression that is matched on is a `Symbol` (namely,
 * one of the arguments to the function). It does this by creating new
 * `FunDef`s for cases deeper in the rhs. For now we assume that the resulting
 * `Case` expressions always make the expression that is matched on strict
 * (which is strictly speaking not correct for expressions like `case x of _ ->
 * ...`).
 *
 * TODO: the remaining work for this changeset consists of:
 * - Generating the correct matching code in `compileCase`
 * - Adding a fallthrough case in the base case of `compileCase`
 * - Correctly determining the type of generated `FunDef`s in `liftCases`.
 *   Types are currently not used, so an alternative is to do type checking
 *   before `liftCases` and do `liftCases` at an untyped stage. But since we
 *   need type inference in the end anyway it is not so clear what the benefit
 *   of a separate untyped stage would be.
 */
liftCases :: !Namespace !Definition -> CompileM [Definition]
liftCases ns (FunDef name args ret expr) =
	liftCasesFromExpr True expr <$&> \(defs,expr) -> [FunDef name args ret expr:defs]
where
	liftCasesFromExpr toplevel e = case e of
		BasicValue _ ->
			pure ([], e)
		Symbol _ ->
			pure ([], e)
		Constructor _ ->
			pure ([], e)
		Case e alts ->
			mapM liftCasesFromAlt alts <$&> unzip <$&> appFst flatten >>= \(defs,alts) ->
			liftCase e alts <$&> appFst ((++) defs)
		ExpApp e1 e2 ->
			liftCasesFromExpr False e1 >>= \(ds1,e1) ->
			liftCasesFromExpr False e2 >>= \(ds2,e2) ->
			pure (ds1 ++ ds2, ExpApp e1 e2)

	liftCasesFromAlt (CaseAlternative pat expr) =
		liftCasesFromExpr False expr <$&> appSnd (CaseAlternative pat)

	liftCase e alts =
		freshLabel "case" >>= \funName ->
		addGlobalFunction {ns=ns, id=funName} {arity=length argsToLift + 1, type=undef} $> // TODO type
		( [FunDef funName (argsToLift ++ [("_casearg", undef)]) undef (Case (Symbol "_casearg") alts)] // TODO types
		, foldl ExpApp (Symbol funName) ([Symbol sym \\ (sym,_) <- argsToLift] ++ [e])
		)
	where
		usedSyms = usedSymbols (e, alts)
		argsToLift = [arg \\ arg=:(sym,_) <- args | 'Data.Set'.member sym usedSyms]
liftCases ns (TestDef name ty expr expected) =
	freshLabel "test" >>= \funName ->
	liftCases ns (FunDef funName [] ty expr) <$&> \defs ->
	[TestDef name ty (Symbol funName) expected : defs]
liftCases _ def =
	pure [def]

freshLabel :: !String -> CompileM Label
freshLabel prefix = state \st -> (concat3 "_l" prefix (toString st.fresh_ident), {st & fresh_ident=st.fresh_ident+1})

lookupConstructor :: !Namespace !ConstructorIdent !Globals -> MaybeError String ConstructorDef
lookupConstructor ns id globals = mb2error
	(concat4 "Unknown constructor " ns "." id)
	('Data.Map'.get {ns=ns, id=id} globals.constructors)

lookupFunction :: !Namespace !SymbolIdent !Globals -> MaybeError String FunctionInfo
lookupFunction ns id globals = mb2error
	(concat4 "Unknown symbol " ns "." id)
	('Data.Map'.get {ns=ns, id=id} globals.functions)

compileDefinition :: !Namespace !Definition -> CompileM [Line]
compileDefinition _ (TypeDef _ _) = pure
	[]
compileDefinition ns (DataDef _ _ constructors) =
	(++) [StartSection "data"] <$>
	flatten <$> mapM (compileConstructor ns) constructors
compileDefinition ns (FunDef id args ret expr) =
	gets (\s -> s.globals) >>= \globals ->
	(++)
		(if (isEmpty args) [] (
			[ StartSection "data"
			, Align 2
			] ++ flatten
			[[ Label (closure_label i)
			// TODO: Ideally we would use the following here:
			//, RawByte i // pointer arity
			//, RawByte 0 // basic value arity
			//, RawByte (length args-i-1) // number of arguments that still have to be curried in minus 1
			//, RawByte 0 // reserved
			// But since SPIM does not allow .byte in the text section, we use:
			, RawWord
				(i bitor // pointer arity
				((length args-i-1) << 16)) // number of arguments that still have to be curried in minus 1
			] \\ i <- [0..length args-1]
			] ++
			[ RawWordLabel n_label
			])) <$>
	(++)
		[ StartSection "text"
		, Global n_label
		// TODO: Ideally we would use the following here:
		//, Align 1
		//, RawByte (sum [2^i \\ i <- [0..] & _ <- args]) // all strict for now, TODO change
		//, RawByte (length args) // arity
		// But since SPIM does not allow .byte in the text section, we use:
		, Align 2
		, RawWord
			(sum [2^i \\ i <- [0..] & _ <- args] bitor // all strict for now, TODO change
			(length args << 8)) // arity
		// instead... (end modification)
		, Label n_label
		] <$>
	case expr of
		// due to liftCases, all Case expressions are on the top level and have
		// a Symbol as expression which is the last argument to the function
		Case (Symbol local) alts ->
			compileCase ns globals locals local alts
		_ ->
			map Instr <$> compileExpr ns globals locals expr
where
	closure_label i = closureLabel ns id i
	n_label = functionLabel ns NodeEntry id
	locals = 'Data.Map'.fromList
		[ (id, FrontPtrArg offset)
		\\ (id,_) <- args
		& offset <- [0..]
		]

compileCase :: !Namespace !Globals !Locals !SymbolIdent ![CaseAlternative] -> CompileM [Line]
compileCase _ _ _ _ [] =
	pure [] // TODO: add catch for partial cases
compileCase ns globals locals exprSymbol [CaseAlternative pat rhs:rest] =
	liftM2 (++) caseAlt (compileCase ns globals locals exprSymbol rest)
where
	// NB: we can assume that the expression has been evaluated; cases are strict
	caseAlt = case pat of
		Wildcard ->
			map Instr <$>
			compileExpr ns globals locals rhs
		BasicValuePattern bv ->
			abort "compileCase: BasicValuePattern\n" // TODO
		IdentPattern sym ->
			map Instr <$>
			compileExpr ns globals locals rhs // TODO: add sym to locals, update stack offsets
		ConstructorPattern cons args ->
			freshLabel "nomatch" >>= \nomatch ->
			pure // TODO: match
				[ Instr (Jump NoLink (Direct (Address 0 nomatch)))
				, Label nomatch
				]

compileConstructor :: !Namespace !ConstructorDef -> CompileM [Line]
compileConstructor ns (ConstructorDef id args) = pure
	[ Global label
	, Align 1
	, Label label
	, RawByte (length args) // pointer arity
	, RawByte 0 // basic value arity
	//, RawByte -1 // number of arguments still to be curried in (unused for constructors)
	]
where
	label = constructorLabel ns id

compileExpr :: !Namespace !Globals !Locals !Expression -> CompileM [Instruction]
compileExpr ns globals locals expr =
	case simulate [SVRegOffset FrontEvalPtr 0] expr` of
		Error e -> fail ("Compiling an expression failed: " +++ e)
		Ok instrs -> pure instrs
where
	expr` = simulator ns globals locals expr >>| indirectAndEval

simulator :: !Namespace !Globals !Locals !Expression -> Simulator ()
simulator _ _ _ (BasicValue bv) =
	pushBasicValue bv >>|
	buildCons (constructorLabel "" (label bv)) 1
where
	label (BVInt _) = "INT"
	label (BVChar _) = "CHAR"
simulator ns globals locals (Symbol id) =
	case 'Data.Map'.get id locals of
		?Just (FrontPtrArg i) ->
			stackSize >>= \n ->
			pushArg (n-1) i
		?None ->
			liftT (lookupFunction ns id globals) >>= \info -> case info.arity of
				0 ->
					buildThunk (functionLabel ns NodeEntry id) 0
				_ ->
					fail "symbol with arity > 0" // TODO implement
simulator ns globals locals expr=:(ExpApp _ _) =
	case f of
		Symbol id -> // TODO include locals
			liftT (lookupFunction ns id globals) >>= \info
				| info.arity == length args ->
					mapM_ (simulator ns globals locals) (reverse args) >>|
					buildThunk (functionLabel ns NodeEntry id) info.arity
				| info.arity > length args ->
					mapM_ (simulator ns globals locals) (reverse args) >>|
					buildCons (closureLabel ns id (length args)) (length args)
				| info.arity < length args ->
					let
						(closure_args,extra_args) = splitAt info.arity args
						closure = foldl ExpApp f closure_args
					in
					mapM_ (simulator ns globals locals) extra_args >>|
					simulator ns globals locals closure >>|
					mapM_ (\_ -> buildThunk (functionLabel "" NodeEntry "ap") 2) extra_args
		Constructor id ->
			liftT (lookupConstructor ns id globals) >>= \(ConstructorDef _ arg_types)
				| length arg_types == length args ->
					mapM_ (simulator ns globals locals) (reverse args) >>|
					buildCons (constructorLabel ns id) (length args)
				| otherwise -> fail ("arity mismatch in application of " +++ id) // TODO implement
		_ -> // TODO
			fail "unexpected lhs of function application"
where
	(f, args) = linearizeApp expr []

	linearizeApp (ExpApp f x) xs = linearizeApp f [x:xs]
	linearizeApp e xs = (e, xs)
simulator ns globals locals (Case e alts) =
	liftT (fail "case in simulator")
/*
	simulator ns locals e >>|
	//eval >>| // TODO
	liftT freshLabel >>= \end ->
	mapM (simulateAlternative end) alts >>|
	label end
where
	simulateAlternative end (CaseAlternative pattern expr) =
		liftT freshLabel >>= \no_match ->
		//simulatePattern no_match locals pattern >>= \new_locals ->
		//simulator ns new_locals expr >>|
		jump end >>|
		label no_match
simulator _ _ _ _ = // TODO
	pushBasicValue (BVInt 0) >>|
	buildCons (constructorLabel "" "INT") 1
*/
//	= BasicValue !BasicValue
//	| Symbol !SymbolIdent
//	| Constructor !ConstructorIdent
//	| Case !Expression ![CaseAlternative]
//	| ExpApp !Expression !Expression

//	= Wildcard
//	| BasicValuePattern !BasicValue
//	| IdentPattern !SymbolIdent
//	| ConstructorPattern !ConstructorIdent ![SymbolIdent]