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 ()

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

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 
//// ------------------------
instance zero Substitution where zero = 'Map'.newMap

compose :: Substitution Substitution -> Substitution
compose s1 s2 = 'Map'.union (Mapmap (subst s1) s2) s1
//Note: unlike function composition, compose is left biased!

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 ()
withGamma :: (Gamma -> Gamma) -> Typing ()
withGamma f = modify (appFst f) >>| pure ()
fresh :: Typing Type
fresh = gets snd >>= \vars-> 
        modify (appSnd $ const $ tail vars) >>| 
        pure (IdType (head vars))

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





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