implementation module check

import StdEnv

import Data.Either
import Data.List
import Data.Functor
import Data.Func
import Data.Maybe
import Data.Tuple
import Control.Monad
import Control.Monad.Trans
import Control.Monad.State

import qualified Data.Map
from Data.Map import instance Functor (Map k)

import ast

import StdDebug
check :: [Function] -> Either [String] Expression
check fs
	# dups = filter (\x->length x > 1) (groupBy (\(Function i _ _) (Function j _ _)->i == j) fs)
	| length dups > 0 = Left ["Duplicate functions: ":[toString n\\[(Function n _ _):_]<-dups]]
	= case partition (\a->a=:(Function ['start'] _ _)) fs of
		([], _) = Left ["No start function defined"]
		([Function _ [] e], fs)
			# e = foldr (\(Function i a e)->Let i (mkLambda a e)) e fs
			= case runInfer (infer 'Data.Map'.newMap e) of
				Left e = Left e
				Right s
					= Left [printToString s]
		([Function _ _ _], _) = Left ["Start cannot have arguments"]

mkLambda :: [[Char]] Expression -> Expression
mkLambda [] e = e
mkLambda [a:as] e = Lambda a (mkLambda as e)

import Text.GenPrint
derive gPrint Scheme, Type

//Polytypes
:: Scheme = Forall [[Char]] Type
:: TypeEnv :== 'Data.Map'.Map [Char] Scheme
:: Subst   :== 'Data.Map'.Map [Char] Type
nullSubst = 'Data.Map'.newMap

:: Infer a :== StateT [Int] (Either [String]) a
runInfer :: (Infer (Subst, Type)) -> Either [String] Scheme
runInfer i = uncurry closeOver <$> evalStateT i [0..]
where
	closeOver :: Subst Type -> Scheme
	closeOver sub ty = normalize (generalize 'Data.Map'.newMap (apply sub ty))

	normalize :: Scheme -> Scheme
	normalize i = i
//	normalize (Forall ts body) = Forall (snd <$> ord) (normtype body)
//	where
//		ord = zip2 (removeDup $ fv body) (fmap letters)
//
//	fv (TVar a) = [a]
//	fv (TFun a b) = fv a ++ fv b
//	fv _ = []
//
//	normtype (TFun a b) = TFun (normtype a) (normtype b)
//	normtype (TCon a)   = TCon a
//	normtype (TVar a)   =
//		  case lookup a ord of
//			Just x = TVar x
//			Nothing = Left ["type variable not in signature"]

fresh :: Infer Type
fresh = getState >>= \[s:ss]->put ss >>| pure (TVar (['v':[c\\c<-:toString s]]))

compose :: Subst Subst -> Subst
compose s1 s2 = 'Data.Map'.union (apply s1 <$> s2) s1

class Substitutable a where
	apply :: Subst a -> a
	ftv :: a -> [[Char]]

instance Substitutable Type where
	apply s t=:(TVar v) = 'Data.Map'.findWithDefault t v s
	apply s (TFun t1 t2) = on TFun (apply s) t1 t2
	apply _ x = x
	
	ftv (TVar v) = [v]
	ftv (TFun t1 t2) = on union ftv t1 t2
	ftv _ = []

instance Substitutable Scheme where
	apply s (Forall as t) = Forall as $ apply (foldr 'Data.Map'.del s as) t
	ftv (Forall as t) = difference (ftv t) (removeDup as)

instance Substitutable TypeEnv where
	apply s env = apply s <$> env
	ftv env = ftv ('Data.Map'.elems env)

instance Substitutable [a] | Substitutable a where
	apply s l = apply s <$>  l
	ftv t = foldr (union o ftv) [] t

occursCheck :: [Char] -> (a -> Bool) | Substitutable a
occursCheck a = isMember a o ftv

unify :: Type Type -> Infer Subst
unify (TFun l r) (TFun l` r`)
	=        unify l l`
	>>= \s1->on unify (apply s1) r r`
	>>= \s2->pure (compose s1 s2)
unify (TVar a) t = bind a t
unify t (TVar a) = bind a t
unify TInt TInt = pure nullSubst
unify TBool TBool = pure nullSubst
unify t1 t2 = liftT (Left ["Cannot unify: ", toString t1, " with ", toString t2])

bind :: [Char] Type -> Infer Subst
bind a (TVar t) | a == t = pure nullSubst
bind a t
	| occursCheck a t = liftT (Left ["Infinite type: ", toString a, " and ", toString t])
	= pure ('Data.Map'.singleton a t)

instantiate :: Scheme -> Infer Type
instantiate (Forall as t)
	=         sequence [fresh\\_<-as]
	>>= \as`->pure (apply ('Data.Map'.fromList $ zip2 as as`) t)

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

infer :: TypeEnv Expression -> Infer (Subst, Type)
infer env (Lit (Int _)) = pure (nullSubst, TInt)
infer env (Lit (Bool _)) = pure (nullSubst, TBool)
infer env (Var x) = case 'Data.Map'.get x env of
	Nothing = liftT (Left ["Unbound variable: ", toString x])
	Just s = tuple nullSubst <$> instantiate s
infer env (App e1 e2)
	=              fresh
	>>= \tv->      infer env e1
	>>= \(s1, t1)->infer (apply s1 env) e2
	>>= \(s2, t2)->unify (apply s2 t1) (TFun t2 tv)
	>>= \s3->      pure (compose (compose s3 s2) s1, apply s3 tv)
infer env (Lambda x b)
	=              fresh
	>>= \tv->      infer ('Data.Map'.put x (Forall [] tv) env) b
	>>= \(s1, t1)->pure (s1, TFun (apply s1 tv) t1)
infer env (Builtin c a) = undef
infer env (Let x e1 e2)
	=              infer env e1
	>>= \(s1, t1)->let env` = apply s1 env in infer ('Data.Map'.put x (generalize env` t1) env) e2
	>>= \(s2, t2)->pure (compose s1 s2, t2)