implementation module check

import StdEnv
import Data.Either
import Data.List
import Data.Functor
import Data.Func
import Data.Maybe
import Data.Monoid
import Control.Applicative
import Control.Monad
import Control.Monad.Trans
import Control.Monad.RWST
import qualified Data.Map as DM
from Data.Map import instance Functor (Map k)

import ast

check :: AST -> Either [String] AST
check (AST fs) = case find (\f->f=:(Function ['start'] [] _)) fs of
	Nothing = Left ["No start function defined"]
	Just _ = Right (AST fs)

:: Type
	= TVar [Char]
	| TInt
	| TBool
	| TChar
	| TFun Type Type
instance == Type where
	(==) (TVar a) (TVar b) = a == b
	(==) TInt TInt = True
	(==) TBool TBool = True
	(==) TChar TChar = True
	(==) (TFun a1 a2) (TFun b1 b2) = a1 == b1 && a2 == b2
	(==) _ _ = False
instance toString Type where
	toString (TVar s) = toString s
	toString TInt = "Int"
	toString TBool = "Bool"
	toString TChar = "Char"
	toString (TFun t1 t2) = toString t1 +++ " -> " +++ toString t2
:: Scheme = Scheme [[Char]] Type
class Types a where
	ftv :: a -> [[Char]]
	apply :: Subst a -> a
instance Types Type where
	ftv (TVar n) = [n]
	ftv TInt = []
	ftv TBool = []
	ftv TChar = []
	ftv (TFun t1 t2) = union (ftv t1) (ftv t2)

	apply s (TVar n) = case 'DM'.get n s of
		Nothing = TVar n
		Just t = t
	apply s (TFun t1 t2) = TFun (apply s t1) (apply s t2)
	apply s t = t
instance Types Scheme where
	ftv (Scheme vars t) = difference (ftv t) vars
	apply s (Scheme vars t) = Scheme vars (apply (foldr 'DM'.del s vars) t)
instance Types [a] | Types a where
	ftv l = foldr union [] (map ftv l)
	apply s l = map (apply s) l

:: Subst :== 'DM'.Map [Char] Type
composeSubst s1 s2 = 'DM'.union ((apply s1) <$> s2) s1

:: TypeEnv = TypeEnv ('DM'.Map [Char] Scheme)
remove :: TypeEnv [Char] -> TypeEnv
remove (TypeEnv env) var = TypeEnv ('DM'.del var env)

instance Types TypeEnv where
	ftv (TypeEnv env) = ftv ('DM'.elems env)
	apply s (TypeEnv env) = TypeEnv (apply s <$> env)

generalize :: TypeEnv Type -> Scheme
generalize env t = Scheme (difference (ftv t) (ftv env)) t

:: TI a :== RWST TIEnv () TIState (Either [String]) a
:: TIState = {tiSupply :: Int, tiSubst :: Subst}
:: TIEnv :== [Int]

mgu :: Type Type -> TI Subst
mgu (TFun l r) (TFun l` r`) = composeSubst <$> mgu l l` <*> mgu r r`
mgu (TVar u) t = varBind u t
mgu t (TVar u) = varBind u t
mgu TInt TInt = pure 'DM'.newMap
mgu TBool TBool = pure 'DM'.newMap
mgu TChar TChar = pure 'DM'.newMap
mgu t1 t2 = liftT (Left ["cannot unify: " +++ toString t1 +++ " with " +++ toString t2])

varBind :: [Char] Type -> TI Subst
varBind u t
	| t == TVar u = pure 'DM'.newMap
	| isMember u (ftv t) = liftT (Left ["occur check fails: " +++ toString u +++ " vs. " +++ toString t])
	= pure ('DM'.singleton u t)

newTyVar :: [Char] -> TI Type
newTyVar prefix
	=   get
	>>= \t->put {t & tiSupply=t.tiSupply+1}
	>>| pure (TVar (prefix ++ fromString (toString t.tiSupply)))

instantiate :: Scheme -> TI Type
instantiate (Scheme vars t)
	= mapM (\_->newTyVar ['a']) vars
	>>= \nvars->pure (apply ('DM'.fromList (zip2 vars nvars)) t)

class infer a :: TypeEnv a -> TI (Subst, Type)
instance infer Value where
	infer _ (Int _) = pure ('DM'.newMap, TInt)
	infer _ (Bool _) = pure ('DM'.newMap, TBool)
	infer _ (Char _) = pure ('DM'.newMap, TChar)
instance infer Expression where
	infer e (Lit a) = infer e a
	infer (TypeEnv env) (Var v) = case 'DM'.get v env of
		Nothing = liftT (Left ["unbound variable: " +++ toString v])
		Just s = instantiate s >>= \t->pure ('DM'.newMap, t)
	infer env (App e1 e2)
		=              newTyVar ['a']
		>>= \tv      ->infer env e1
		>>= \(s1, t1)->infer (apply s1 env) e2
		>>= \(s2, t2)->mgu (apply s2 t1) (TFun t2 tv)
		>>= \s3->pure (composeSubst s3 (composeSubst s2 s1), apply s3 tv)
	//infer env (Lambda s e)
	//	=        newTyVar ['l']
	//	>>= \tv->pure undef//inEnv (s, Forall [] tv) (infer e)
	//	>>= \t-> pure (TFun tv t)