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 Control.Monad.State
import Data.Either
import Data.Maybe
import Data.Monoid
import Data.List
import Data.Functor
import Data.Tuple

import StdString
import StdTuple
import StdList
import StdMisc
import StdEnum
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)]
:: SemError
	= ParseError Pos String 
	| UnifyError Pos Type Type 
    | InfiniteTypeError Pos 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 :: [TVar]
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]
    _ = Right []
	//_ = case execRWST (constraints fd) zero variableStream of
	//	Left e = Left [e]
	//	Right (a, b) = Right b
where
		constraints :: [FunDecl] -> Typing ()
        constraints _ = pure ()
        //TODO: fix
		//constraints fds = mapM_ funconstraint fds >>| pure ()

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

		vardeclconstraint :: VarDecl -> Typing ()
        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 ()

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

extend :: String Scheme Gamma -> Gamma
extend k t g = 'Map'.put k t g

//// ------------------------
//// algorithm U, Unification 
//// ------------------------
instance zero Substitution where zero = 'Map'.newMap

compose :: Substitution Substitution -> Substitution
compose s1 s2 = 'Map'.union (Mapmap (subst s1) s2) s1
//Note: just like function compositon compose does snd first

occurs :: TVar a -> Bool | Typeable a
occurs tvar a = elem tvar (ftv a)

unify :: Type Type -> Either SemError Substitution
unify t1=:(IdType tv) t2 = unify t2 t1
unify t1 t2=:(IdType tv)    | t1 == (IdType tv) = Right zero
                            | occurs tv t1 = Left $ InfiniteTypeError zero t1
                            | otherwise = Right $ 'Map'.singleton tv t1
unify (ta1->>ta2) (tb1->>tb2) = unify ta1 tb1 >>= \s1->
                                unify tb1 tb2 >>= \s2->
                                Right $ compose s1 s2
unify (TupleType (ta1,ta2)) (TupleType (tb1,tb2)) = unify ta1 tb1 >>= \s1->
                                                    unify ta2 tb2 >>= \s2->
                                                    Right $ compose s1 s2
unify (ListType t1) (ListType t2) = unify t1 t2
unify t1 t2 | t1 == t2  = Right zero
            | otherwise = Left $ UnifyError zero t1 t2

//// ------------------------
//// Algorithm M, Inference and Solving
//// ------------------------
//The typing monad
:: Typing a :== StateT (Gamma, [TVar]) (Either SemError) a
gamma :: Typing Gamma
gamma = gets fst
putGamma :: Gamma -> Typing ()
putGamma g = modify (appFst $ const g) >>| pure ()
changeGamma :: (Gamma -> Gamma) -> Typing Gamma
changeGamma f = modify (appFst f) >>| gamma
withGamma :: (Gamma -> a) -> Typing a
withGamma f = f <$> gamma
fresh :: Typing Type
fresh = gets snd >>= \vars-> 
        modify (appSnd $ const $ tail vars) >>| 
        pure (IdType (head vars))

lift :: (Either SemError a) -> Typing a
lift (Left e) = liftT $ Left e 
lift (Right v) = pure v

//instantiate maps a schemes type variables to variables with fresh names
//and drops the quantification: i.e. forall a,b.a->[b] becomes c->[d]
instantiate :: Scheme -> Typing Type
instantiate (Forall bound t) = 
    mapM (const fresh) bound >>= \newVars->
    let s = 'Map'.fromList (zip (bound,newVars)) in
    pure (subst s t)

//generalize quentifies all free type variables in a type which are not
//in the gamma
generalize :: Type -> Typing Scheme
generalize t = gamma >>= \g-> pure $ Forall (difference (ftv t) (ftv g)) t

lookup :: String -> Typing Type 
lookup k = gamma >>= \g-> case 'Map'.member k g of
    False = liftT (Left $ UndeclaredVariableError zero k)
    True = instantiate $ 'Map'.find k g

//The inference class
//When tying it all together we will treat the program is a big 
//let x=e1 in let y=e2 in .... 
class infer a :: a -> Typing (Substitution, Type)

////---- Inference for Expressions ----

instance infer Expr where 
 infer e = case e of
    VarExpr _ (VarDef k fs) = (\t->(zero,t)) <$> lookup k
    //instantiate is key for the let polymorphism!
    //TODO: field selectors

    Op2Expr _ e1 op e2 =
        infer e1 >>= \(s1, t1) -> 
        infer e2 >>= \(s2, t2) ->
        fresh >>= \tv ->
        let given = t1 ->> t2 ->> tv in 
        op2Type op >>= \expected ->
        lift (unify expected given) >>= \s3 ->
        pure ((compose s3 $ compose s2 s1), subst s3 tv)

    Op1Expr _ op e1 =
        infer e1 >>= \(s1, t1) -> 
        fresh >>= \tv ->
        let given = t1 ->> tv in 
        op1Type op >>= \expected ->
        lift (unify expected given) >>= \s2 ->
        pure (compose s2 s1, subst s2 tv)

    EmptyListExpr _ = (\tv->(zero,tv)) <$> fresh

    TupleExpr _ (e1, e2) = 
        infer e1 >>= \(s1, t1) -> 
        infer e2 >>= \(s2, t2) ->
        pure (compose s2 s1, TupleType (t1,t2))

    FunExpr _ f args fs = //todo: fieldselectors
        lookup f >>= \expected ->
        let accST = (\(s,ts) e->infer e >>= \(s_,et)->pure (compose s_ s,ts++[et])) in
        foldM accST (zero,[]) args >>= \(s1, argTs)->
        fresh >>= \tv->
        let given = foldr (->>) tv argTs in
        lift (unify expected given) >>= \s2->
        pure (compose s2 s1, subst s2 tv)

    IntExpr _ _ = pure $ (zero, IntType)
    BoolExpr _ _ = pure $ (zero, BoolType)
    CharExpr _ _ = pure $ (zero, CharType)


op2Type :: Op2 -> Typing 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 -> Typing Type
op1Type UnNegation = pure $ (BoolType ->> BoolType)
op1Type UnMinus = pure $ (IntType ->> IntType)

////----- Inference for Statements -----
applySubst :: Substitution -> Typing Gamma
applySubst s = changeGamma (subst s)

instance infer Stmt where
 infer s = case s of
    IfStmt e th el = 
        infer e >>= \(s1, et)->
        lift (unify et BoolType) >>= \s2 ->
        applySubst (compose s2 s1) >>|        
        infer th >>= \(s3, tht)->
        applySubst s3 >>|
        infer el >>= \(s4, elt)->
        applySubst s4 >>|
        lift (unify tht elt) >>= \s5-> 
        pure (compose s5 $ compose s4 $ compose s3 $ compose s1 s2, subst s5 tht)

    WhileStmt e wh = 
        infer e >>= \(s1, et)->
        lift (unify et BoolType) >>= \s2 ->
        applySubst (compose s2 s1) >>|
        infer wh >>= \(s3, wht)->
        pure (compose s3 $ compose s1 s2, subst s3 wht)

    AssStmt (VarDef k fs) e =
        infer e >>= \(s1, et)->
        applySubst s1 >>|
        changeGamma (extend k (Forall [] et)) >>| //todo: fieldselectors
        pure (s1, VoidType)

    FunStmt f es = undef //what is this? 

    ReturnStmt Nothing = pure (zero, VoidType)
    ReturnStmt (Just e) = infer e


instance infer [a] | infer a where
    infer _ = undef

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)

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: "

instance MonadTrans (StateT (Gamma, [TVar])) where
    liftT m = StateT \s-> m >>= \a-> return (a, s)

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