implementation module sem

import qualified Data.Map as Map

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

import Control.Applicative
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
:: Typing a :== StateT (Gamma, [TVar]) (Either SemError) a
:: 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..]

defaultGamma :: Gamma //includes all default functions
defaultGamma = extend "print" (Forall ["a"] ((IdType "a") ->> VoidType))
                $ extend "isEmpty" (Forall ["a"] ((ListType (IdType "a")) ->> BoolType))
                $ extend "read" (Forall [] (FuncType CharType))
                $ extend "1printchar" (Forall [] (CharType ->> VoidType))
                $ extend "1printint" (Forall [] (IntType ->> VoidType))
                $ extend "1printbool" (Forall [] (BoolType ->> VoidType))
                zero

sem :: AST -> Either [SemError] (AST, Gamma)
sem (AST fd) = case foldM (const $ hasNoDups fd) () fd 
            		>>| foldM (const isNiceMain) () fd
            		>>| hasMain fd
                    >>| runStateT (unfoldLambda fd >>= type) (defaultGamma, variableStream) of
	Left e = Left [e]
    Right ((_,fds),(gam,_)) = Right (AST fds, gam)
where
		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 ()


//------------------
// LAMBDA UNFOLDING
//------------------
unfoldLambda :: [FunDecl] -> Typing [FunDecl]
unfoldLambda [] = pure []
unfoldLambda [fd:fds] = unfoldL_ fd >>= \(gen1, fs_)-> 
                        unfoldLambda fds >>= \gen2->
                        pure $ gen1 ++ [fs_] ++ gen2

flattenT :: [([a],b)] -> ([a],[b])
flattenT ts = (flatten $ map fst ts, map snd ts)

class unfoldL_ a :: a -> Typing ([FunDecl], a)

instance unfoldL_ FunDecl where
    unfoldL_ (FunDecl p f args mt vds stmts) = 
        flattenT <$> mapM unfoldL_ vds >>= \(fds1,vds_) ->
        flattenT <$> mapM unfoldL_ stmts >>= \(fds2,stmts_)->
        pure (fds1 ++ fds2, FunDecl p f args mt vds_ stmts_)

instance unfoldL_ VarDecl where
    unfoldL_ (VarDecl p mt v e) = unfoldL_ e >>= \(fds, e_)->pure (fds, VarDecl p mt v e_)

instance unfoldL_ Stmt where
    unfoldL_ (IfStmt e th el) = unfoldL_ e >>= \(fds, e_)->pure (fds, IfStmt e_ th el)
    unfoldL_ (WhileStmt e c) = unfoldL_ e >>= \(fds, e_)->pure (fds, WhileStmt e_ c)
    unfoldL_ (AssStmt vd e) = unfoldL_ e >>= \(fds, e_)->pure (fds, AssStmt vd e_)
    unfoldL_ (FunStmt f es fs) = flattenT <$> mapM unfoldL_ es >>= \(fds, es_)->
        pure (fds, FunStmt f es_ fs)
    unfoldL_ (ReturnStmt (Just e)) = unfoldL_ e >>= \(fds, e_) -> 
        pure (fds, ReturnStmt (Just e_))
    unfoldL_ (ReturnStmt Nothing) = pure ([], ReturnStmt Nothing)

instance unfoldL_ Expr where
    unfoldL_ (LambdaExpr p args e) = 
        fresh >>= \(IdType n) ->
        let f = ("2lambda_"+++n) in
        let fd = FunDecl p f args Nothing [] [ReturnStmt $ Just e] in 
        let fe = VarExpr p (VarDef f []) in
        pure ([fd], fe)
    unfoldL_ (FunExpr p f es fs) = flattenT <$> mapM unfoldL_ es >>= \(fds, es_)->
        pure (fds, FunExpr p f es_ fs)
    unfoldL_ (Op2Expr p e1 op e2) = unfoldL_ e1 >>= \(fds1, e1_)->
        unfoldL_ e2 >>= \(fds2, e2_)->
        pure (fds1++fds2, Op2Expr p e1_ op e2_)
    unfoldL_ (Op1Expr p op e1) = unfoldL_ e1 >>= \(fds, e1_)->pure (fds, Op1Expr p op e1_)
    unfoldL_ (TupleExpr p (e1, e2)) = unfoldL_ e1 >>= \(fds1, e1_)->
        unfoldL_ e2 >>= \(fds2, e2_)->
        pure (fds1++fds2, TupleExpr p (e1_, e2_))
    unfoldL_ e = pure ([], e)

//------------
//------------
//      TYPING 
//------------
//------------

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 (FuncType t)            = ftv t
    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 (FuncType t)        = FuncType (subst s t)
    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 t2=:(IdType tv)    | t1 == (IdType tv) = Right zero
                            | occurs tv t1 = Left $ InfiniteTypeError zero t1
                            | otherwise = Right $ 'Map'.singleton tv t1
unify t1=:(IdType tv) t2 = unify t2 t1
unify (ta1->>ta2) (tb1->>tb2) = unify ta1 tb1 >>= \s1->
                                unify (subst s1 ta2) (subst s1 tb2) >>= \s2->
                                Right $ compose s2 s1
unify (TupleType (ta1,ta2)) (TupleType (tb1,tb2)) = unify ta1 tb1 >>= \s1->
                                                    unify (subst s1 ta2) (subst s1 tb2) >>= \s2->
                                                    Right $ compose s2 s1
unify (ListType t1) (ListType t2) = unify t1 t2
unify (FuncType t1) (FuncType t2) = unify t1 t2
unify t1 t2 | t1 == t2  = Right zero
            | otherwise = Left $ UnifyError zero t1 t2

//// ------------------------
//// Algorithm M, Inference and Solving
//// ------------------------
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, a)

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

instance infer Expr where 
 infer e = case e of
    VarExpr _ (VarDef k fs) = lookup k >>= \t ->
        foldM foldFieldSelectors t fs >>= \finalT ->
        pure (zero, finalT, e)

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

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

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

    TupleExpr p (e1, e2) = 
        infer e1 >>= \(s1, t1, e1_) -> 
        applySubst s1 >>|
        infer e2 >>= \(s2, t2, e2_) ->
        applySubst s2 >>|
        pure (compose s2 s1, TupleType (t1,t2), TupleExpr p (e1_,e2_))

    LambdaExpr _ _ _ = liftT $ Left $ Error "PANIC: lambdas should be Unfolded"

    FunExpr p f args fs =
        lookup f >>= \expected ->
        let accST = (\(s,ts,es) e->infer e >>= \(s_,et,e_)-> pure (compose s_ s,ts++[et],es++[e_])) in
        foldM accST (zero,[],[]) args >>= \(s1, argTs, args_)->
        applySubst s1 >>|
        (case f of
            "print" = case head argTs of
                IntType = pure "1printint"
                CharType = pure "1printchar"
                BoolType = pure "1printbool"
                ListType (CharType) = pure "1printstr"
                t = liftT $ Left $ SanityError p ("can not print " +++ toString t)
            _ = pure f
        ) >>= \newF->
        fresh >>= \tv->case expected of
            FuncType t = foldM foldFieldSelectors t fs >>= \returnType ->
                pure (s1, returnType, (FunExpr p newF args fs))
            _ = (let given = foldr (->>) tv argTs in
                lift (unify expected given) >>= \s2->
                applySubst s2 >>|
                let fReturnType = subst s2 tv in
                foldM foldFieldSelectors fReturnType fs >>= \returnType ->
                pure (compose s2 s1, returnType, FunExpr p newF args_ fs))

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

foldFieldSelectors :: Type FieldSelector -> Typing Type 
foldFieldSelectors (ListType t) (FieldHd) = pure t
foldFieldSelectors t=:(ListType _) (FieldTl) = pure t 
foldFieldSelectors (TupleType (t1, _)) (FieldFst) = pure t1
foldFieldSelectors (TupleType (_, t2)) (FieldSnd) = pure t2
foldFieldSelectors t fs = liftT $ Left $ FieldSelectorError zero t fs

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, e_)->
        lift (unify et BoolType) >>= \s2 ->
        applySubst (compose s2 s1) >>|        
        infer th >>= \(s3, tht, th_)->
        applySubst s3 >>|
        infer el >>= \(s4, elt, el_)->
        applySubst s4 >>|
        lift (unify tht elt) >>= \s5-> 
        let sub = compose s5 $ compose s4 $ compose s3 $ compose s2 s1 in
        pure (sub, subst s5 tht, IfStmt e_ th_ el_)

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

    AssStmt vd=:(VarDef k fs) e =
        lookup k >>= \expected ->
        infer e >>= \(s1, given, e_)->
        foldM reverseFs given (reverse fs) >>= \varType->
        lift (unify expected varType) >>= \s2->
        let s = compose s2 s1 in
        applySubst s >>|
        changeGamma (extend k (Forall [] (subst s varType))) >>| 
        pure (s, VoidType, AssStmt vd e_)

    FunStmt f args fs = 
        lookup f >>= \expected ->
        let accST = (\(s,ts,es) e->infer e >>= \(s_,et,e_)-> pure (compose s_ s,ts++[et],es++[e_])) in
        foldM accST (zero,[],[]) args >>= \(s1, argTs, args_)->
        fresh >>= \tv->
        let given = foldr (->>) tv argTs in
        lift (unify expected given) >>= \s2->
        let fReturnType = subst s2 tv in
        foldM foldFieldSelectors fReturnType fs >>= \returnType ->
        (case f of
            "print" = case head argTs of
                IntType = pure "1printint"
                CharType = pure "1printchar"
                BoolType = pure "1printbool"
                ListType (CharType) = pure "1printstr"
                t = liftT $ Left $ SanityError zero ("can not print " +++ toString t)
            _ = pure f) >>= \newF->
        pure (compose s2 s1, VoidType, FunStmt newF args_ fs)

    ReturnStmt Nothing = pure (zero, VoidType, s)
    //hier ook sub applyen
    ReturnStmt (Just e) = infer e >>= \(sub, t, e_)-> pure (sub, t, ReturnStmt (Just e_))

reverseFs :: Type FieldSelector -> Typing Type 
reverseFs t FieldHd = pure $ ListType t
reverseFs t FieldTl = pure t 
reverseFs t FieldFst = fresh >>= \tv -> pure $ TupleType (t, tv)
reverseFs t FieldSnd = fresh >>= \tv -> pure $ TupleType (tv, t)

//The type of a list of statements is either an encountered
//return, or VoidType
instance infer [a] | infer a where
    infer []        = pure (zero, VoidType, [])
    infer [stmt:ss] = 
        infer stmt >>= \(s1, t1, s_) ->
        applySubst s1 >>|
        infer ss >>= \(s2, t2, ss_) ->
        applySubst s2 >>|
        case t1 of
            VoidType = pure (compose s2 s1, t2, [s_:ss_])
            _ = case t2 of
                VoidType = pure (compose s2 s1, t1, [s_:ss_])
                _ = lift (unify t1 t2) >>= \s3 -> 
                    pure (compose s3 $ compose s2 s1, t1, [s_:ss_])

//the type class inferes the type of an AST element (VarDecl or FunDecl)
//and adds it to the AST element
class type a :: a -> Typing (Substitution, a)

instance type VarDecl where
    type (VarDecl p expected k e) = 
        infer e >>= \(s1, given, e_) ->
        applySubst s1 >>| 
        case expected of
            Nothing = pure zero
            Just expected_ = lift (unify expected_ given)
        >>= \s2->
        applySubst s2 >>|
        let vtype = subst (compose s2 s1) given in
        generalize vtype >>= \t ->
        changeGamma (extend k t) >>| 
        pure (compose s2 s1, VarDecl p (Just vtype) k e_)

instance type FunDecl where
    type fd=:(FunDecl p f args expected vds stmts) = 
        gamma >>= \outerScope-> //functions are infered in their own scopde
        introduce f >>|
        mapM introduce args >>= \argTs->
        fresh >>= \tempTv ->
        let temp = foldr (->>) tempTv argTs in 
        (case expected of
            Just expected_ = lift (unify expected_ temp)
            _   = pure zero
        ) >>= \s0->
        applySubst s0 >>|
        type vds >>= \(s1, tVds)->
        applySubst s1 >>|
        infer stmts >>= \(s2, result, stmts_)->
        applySubst s1 >>|
        let argTs_ = map (subst $ compose s2 $ compose s1 s0) argTs in 
        let given = foldr (->>) result argTs_ in 
        (case expected of
            Nothing = pure zero
            Just (FuncType expected_) = lift (unify expected_ given)
            Just expected_ = lift (unify expected_ given)
        ) >>= \s3 ->
        let ftype = subst (compose s3 $ compose s2 $ compose s1 s0) given in
        (case ftype of
            _ ->> _ = pure ftype
            _       = pure $ FuncType ftype
        ) >>= \ftype_->
        generalize ftype_ >>= \t->
        putGamma outerScope >>|
        changeGamma (extend f t) >>| 
        pure (compose s3 $ compose s2 $ compose s1 s0, 
                FunDecl p f args (Just ftype_) tVds stmts_)

instance type [a] | type a where
    type []     = pure (zero, [])
    type [v:vs] = 
        type v >>= \(s1, v_)->
        applySubst s1 >>|
        type vs >>= \(s2, vs_)->
        applySubst (compose s2 s1) >>|
        pure (compose s2 s1, [v_:vs_])

introduce :: String -> Typing Type
introduce k = 
    fresh >>= \tv ->
    changeGamma (extend k (Forall [] tv)) >>|
    pure tv

instance toString Scheme where
    toString (Forall x t) = 
        concat ["Forall ": intersperse "," x] +++ concat [". ", toString t];

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

instance toString Substitution where
    toString subs = 
        concat [concat [k, ": ", toString t, "\n"]\\(k, t)<-'Map'.toList subs]

instance toString SemError where
    toString (SanityError p e) = concat [toString p, 
        "SemError: SanityError: ", e]
    toString (ParseError p s) = concat [toString p, 
        "ParseError: ", s]
    toString (UnifyError p t1 t2) = concat [toString p, 
        "Can not unify types, expected|given:\n", toString t1,
        "\n", toString t2]
    toString (InfiniteTypeError p t) = concat [toString p, 
        "Infinite type: ", toString t]
    toString (FieldSelectorError p t fs) = concat [toString p,
        "Can not run fieldselector '", toString fs, "' on type: ",
        toString t]
    toString (OperatorError p op t) = concat [toString p,
        "Operator error, operator '", toString op, "' can not be",
        "used on type: ", toString t]
    toString (UndeclaredVariableError p k) = concat [toString p,
        "Undeclared identifier: ", k]
    toString (ArgumentMisMatchError p str) = concat [toString p,
        "Argument mismatch: ", str]
    toString (Error e) = concat ["Unknown error during semantical",
        "analysis: ", e]

instance toString (Maybe a) | toString a where
    toString Nothing = "Nothing"
    toString (Just e) = concat ["Just ", toString e]

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

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)