implementation module Snug.Compile

import StdEnv

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
from Text import concat4

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

:: CompileM a :== MaybeError String a

:: Locals :== Map SymbolIdent LocalLocation

:: 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 <$> mapM (compileDefinition ns globals) defs
where
	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
		]
	}

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

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

compileDefinition :: !Namespace !Globals !Definition -> CompileM [Line]
compileDefinition _ _ (TypeDef _ _) = pure
	[]
compileDefinition ns globals (DataDef _ _ constructors) =
	(++) [StartSection "data"] <$>
	flatten <$> mapM (compileConstructor ns globals) constructors
compileDefinition ns globals (FunDef id args ret expr) =
	(++)
		[ StartSection "text"
		, Global 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 label
		] <$>
	map Instr <$> compileExpr ns globals locals expr
where
	label = functionLabel ns NodeEntry id
	locals = 'Data.Map'.fromList
		[ (id, FrontPtrArg offset)
		\\ (id,_) <- args
		& offset <- [0..]
		]

compileConstructor :: !Namespace !Globals !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 0 // number of arguments still to be curried in
	]
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
				| otherwise -> fail "arity mismatch in application" // TODO implement
		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 _ _ _ _ = // TODO
	pushBasicValue (BVInt 0) >>|
	buildCons (constructorLabel "" "INT") 1
//	| Symbol !SymbolIdent
//	| Constructor !ConstructorIdent
//	| Case !Expression ![CaseAlternative]
//	| ExpApp !Expression !Expression