implementation module sem

import qualified Data.Map as Map

from Data.Func import $
from StdFunc import o, flip, const, id

import Control.Monad
import Control.Monad.Trans
import Data.Either
import Data.Maybe
import Data.Monoid
import Data.List 
import Data.Functor

import StdString
import StdList
import StdMisc
import StdEnum
import RWST
import GenEq

from Text import class Text(concat), instance Text String

import AST

:: Scheme = Forall [String] Type
:: Gamma :== 'Map'.Map String Scheme
:: Constraints :== [(Type, Type)]
:: Infer a :== RWST Gamma Constraints [String] (Either SemError) a
:: SemError
	= ParseError Pos String 
	| UnifyError Pos Type Type 
    | FieldSelectorError Pos Type FieldSelector 
	| OperatorError Pos Op2 Type
    | UndeclaredVariableError Pos String
    | ArgumentMisMatchError Pos String
	| SanityError Pos String
	| Error String

instance zero Gamma where
	zero = 'Map'.newMap

//runInfer :: Gamma (Infer a) -> Either SemError a
//runInfer env m = evalRWST m env variableStream

variableStream :: [String]
variableStream = map toString [1..]

sem :: AST -> Either [SemError] Constraints
sem (AST fd) = case foldM (const $ hasNoDups fd) () fd 
		>>| foldM (const isNiceMain) () fd
		>>| hasMain fd of
	Left e = Left [e]
	_ = case execRWST (constraints fd) zero variableStream of
		Left e = Left [e]
		Right (a, b) = Right b
where
		constraints :: [FunDecl] -> Infer ()
		constraints fds = mapM_ funconstraint fds >>| pure ()

		funconstraint :: FunDecl -> Infer ()
		funconstraint fd=:(FunDecl _ ident args mt vardecls stmts) = case mt of
			Nothing = abort "Cannot infer functions yet"
			Just t = inEnv (ident, (Forall [] t)) (
				mapM_ vardeclconstraint vardecls >>| pure ())

		vardeclconstraint :: VarDecl -> Infer ()
		vardeclconstraint (VarDecl p mt ident expr) = infer expr
			>>= \it->inEnv (ident, (Forall [] it)) (pure ())

		hasNoDups :: [FunDecl] FunDecl -> Either SemError ()
		hasNoDups fds (FunDecl p n _ _ _ _)
		# mbs = map (\(FunDecl p` n` _ _ _ _)->if (n == n`) (Just p`) Nothing) fds
		= case catMaybes mbs of
			[] = Left $ SanityError p "HUH THIS SHOULDN'T HAPPEN"
			[x] = pure ()
			[_:x] = Left $ SanityError p (concat 
				[n, " multiply defined at ", toString p])

		hasMain :: [FunDecl] -> Either SemError ()
		hasMain [(FunDecl _ "main" _ _ _ _):fd] = pure ()
		hasMain [_:fd] = hasMain fd
		hasMain [] = Left $ SanityError zero "no main function defined"

		isNiceMain :: FunDecl -> Either SemError ()
		isNiceMain (FunDecl p "main" as mt _ _) = case (as, mt) of
			([_:_], _) = Left $ SanityError p "main must have arity 0"
			([], t) = (case t of
				Nothing = pure ()
				Just VoidType = pure ()
				_ = Left $ SanityError p "main has to return Void")
		isNiceMain _ = pure ()

instance toString Scheme where
	toString (Forall x t) = 
		concat ["Forall ": map ((+++) "\n") x] +++ toString t

instance toString Gamma where
	toString mp = 
		concat [concat [k, ": ", toString v, "\n"]\\(k, v)<-'Map'.toList mp]

instance toString SemError where
	toString (SanityError p e) = concat [toString p, 
		"SemError: SanityError: ", e]
	toString se = "SemError: "

// ------------------------
// First step: Inference
// ------------------------

unify :: Type Type -> Infer ()
unify t1 t2 = tell [(t1, t2)]

inEnv :: (String, Scheme) (Infer a) -> Infer a
inEnv (x, sc) m = local ('Map'.put x sc) m

fresh :: Infer Type
fresh = (gets id) >>= \vars-> (put $ tail vars) >>| (pure $ IdType $ head vars)

op2Type :: Op2 -> Infer Type
op2Type op
| elem op [BiPlus, BiMinus, BiTimes, BiDivide, BiMod]
	= pure (IntType ->> IntType ->> IntType)
| elem op [BiEquals, BiUnEqual]
	= fresh >>= \t1-> fresh >>= \t2-> pure (t1 ->> t2 ->> BoolType)
| elem op [BiLesser, BiGreater, BiLesserEq, BiGreaterEq]
	= pure (IntType ->> IntType ->> BoolType)
| elem op [BiAnd, BiOr]
	= pure (BoolType ->> BoolType ->> BoolType)
| op == BiCons 
	= fresh >>= \t1-> pure (t1 ->> ListType t1 ->> ListType t1)

op1Type :: Op1 -> Infer Type
op1Type UnNegation = pure $ (BoolType ->> BoolType)
op1Type UnMinus = pure $ (IntType ->> IntType)

//instantiate :: Scheme -> Infer Type
//instantiate (Forall as t) = mapM (const fresh) as

lookupEnv :: String -> Infer Type
lookupEnv ident = asks ('Map'.get ident)
	>>= \m->case m of
		Nothing = liftT $ Left $ UndeclaredVariableError zero ident
		Just (Forall as t) = pure t //instantiate ???

class infer a :: a -> Infer Type
instance infer Expr where
	infer (VarExpr _ (VarDef ident fs)) = lookupEnv ident
    infer (Op2Expr _ e1 op e2) = 
        infer e1 >>= \t1 -> 
        infer e2 >>= \t2 ->
        fresh >>= \frsh -> 
        let given = t1 ->> (t2 ->> frsh) in
        op2Type op >>= \expected -> 
        unify expected given >>| 
        return frsh
	infer (Op1Expr _ op e) = 
        infer e >>= \t1 -> 
        fresh >>= \frsh -> 
        let given = t1 ->> frsh in
        op1Type op >>= \expected -> 
        unify expected given >>| 
        pure frsh
	infer (IntExpr _ _) = pure IntType
	infer (CharExpr _ _) = pure CharType
	infer (BoolExpr _ _) = pure BoolType
	infer (FunExpr _ f args sels) =  //todo, iets met field selectors
        lookupEnv f >>= \expected ->
        fresh >>= \frsh -> 
        mapM infer args >>= \argTypes ->
        let given = foldr (->>) frsh argTypes in
        unify expected given >>|
        pure frsh
	infer (EmptyListExpr _) = ListType <$> fresh
	infer (TupleExpr _ (e1, e2)) =
		infer e1 >>= \et1->infer e2 >>= \et2->pure $ TupleType (et1, et2)

//:: VarDef = VarDef String [FieldSelector]
//:: FieldSelector = FieldHd | FieldTl | FieldFst | FieldSnd
//:: Op1 = UnNegation | UnMinus
//:: Op2 = BiPlus | BiMinus | BiTimes | BiDivide | BiMod | BiEquals | BiLesser |
//	BiGreater | BiLesserEq | BiGreaterEq | BiUnEqual | BiAnd | BiOr | BiCons
//:: FunDecl = FunDecl Pos String [String] (Maybe Type) [VarDecl] [Stmt]
//:: FunCall = FunCall String [Expr]
//:: Stmt 
//	= IfStmt Expr [Stmt] [Stmt]
//	| WhileStmt Expr [Stmt]
//	| AssStmt VarDef Expr
//	| FunStmt FunCall
//	| ReturnStmt (Maybe Expr)
//:: Pos = {line :: Int, col :: Int}
//:: AST = AST [VarDecl] [FunDecl]
//:: VarDecl = VarDecl Pos Type String Expr
//:: Type 
//	= TupleType (Type, Type)
//	| ListType Type
//	| IdType String 
//	| IntType 
//	| BoolType
//	| CharType
//	| VarType
//    | VoidType
//    | (->>) infixl 7 Type Type