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 [TVar] Type
:: Gamma :== 'Map'.Map String Scheme //map from Variables! to types
:: Substitution :== 'Map'.Map TVar Type
:: 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

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 _ = pure ()
        //TODO: fix
		//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 _ = pure ()
        //TODO: fix!
		//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: "

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

class Typeable a where
    ftv :: a -> [TVar]
    subst :: Substitution a -> a

instance Typeable Scheme where
    ftv (Forall bound t) = difference (ftv t) bound
    subst s (Forall bound t) = Forall bound $ subst s_ t
        where s_ = 'Map'.filterWithKey (\k _ -> not (elem k bound)) s

instance Typeable [a] | Typeable a where
    ftv types   = foldr (\t ts-> ftv t ++ ts) [] types
    subst s ts  = map (\t->subst s t) ts

instance Typeable Type where
    ftv (TupleType (t1, t2))    = ftv t1 ++ ftv t2
    ftv (ListType t)            = ftv t
    ftv (IdType tvar)           = [tvar]
    ftv (t1 ->> t2)             = ftv t1 ++ ftv t2
    ftv _                       = []
    subst s (TupleType (t1, t2))= TupleType (subst s t1, subst s t2)
    subst s (ListType t1)       = ListType (subst s t1)
    subst s (t1 ->> t2)         = (subst s t1) ->> (subst s t2)
    subst s t1=:(IdType tvar)   = 'Map'.findWithDefault t1 tvar s
    subst s t                   = t

instance Typeable Gamma where
    ftv gamma       = concatMap id $ map ftv ('Map'.elems gamma)
    subst s gamma   = Mapmap (subst s) gamma

//// ------------------------
//// algorithm U, Unification (outside of Monad)
//// ------------------------
instance zero Substitution where zero = 'Map'.newMap
compose :: Substitution Substitution -> Substitution
compose s1 s2 = undef
//unify :: 




Mapmap :: (a->b) ('Map'.Map k a) -> ('Map'.Map k b)
Mapmap _ 'Map'.Tip = 'Map'.Tip
Mapmap f ('Map'.Bin sz k v ml mr) = 'Map'.Bin sz k (f v) 
                                        (Mapmap f ml) 
                                        (Mapmap f mr)

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

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

//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