implementation module check

import StdEnv

import Control.Monad => qualified join
import Control.Monad.State
import Control.Monad.Trans
import Control.Monad.Writer
import Data.Either
import Data.Func
import Data.List
import Data.Tuple
import Data.Map => qualified put, union, difference, find, updateAt
import Data.Maybe
import Text

import ast, scc

import StdDebug

check :: ![Either TypeDef Function] -> Either [String] (Expression, [([Char], Scheme)])
check tdfs
	# dups = filter (\x->length x > 1) (groupBy (\(Function i _ _) (Function j _ _)->i == j) functions)
	| length dups > 0 = Left ["Duplicate functions: ":[toString n\\[(Function n _ _):_]<-dups]]
	= case partition (\a->a=:(Function ['start'] _ _)) functions of
		([], _) = Left ["No start function defined"]
		([Function _ [] e:_], fs)
			# e = makeExpression fs e
			= tuple e <$> runInfer (infer (fromList (conses ++ builtin)) e)
		([Function _ _ _:_], _) = Left ["Start cannot have arguments"]
where
	functions = rights tdfs
	conses = flatten $ map (\(TypeDef n t cs)->
		let cons = Forall t o foldr (-->) (foldl TApp (TVar n) (map TVar t))
		in map (appSnd cons) cs) $ lefts tdfs
	builtin =
		[(['_if'],  Forall [['a']] $ TBool --> TVar ['a'] --> TVar ['a'] --> TVar ['a'])
		,(['_eq'],  Forall [] $ TInt --> TInt --> TBool)
		,(['_mul'], Forall [] $ TInt --> TInt --> TInt)
		,(['_add'], Forall [] $ TInt --> TInt --> TInt)
		,(['_sub'], Forall [] $ TInt --> TInt --> TInt)
		,(['_div'], Forall [] $ TInt --> TInt --> TInt)
		]

makeExpression :: [Function] Expression -> Expression
makeExpression fs start = foldr mkExpr start $ scc $ map (appSnd vars) nicefuns
where
	mkExpr :: [[Char]] -> (Expression -> Expression)
	mkExpr scc = Let [(l, e)\\(l, e)<-nicefuns, s<-scc | s == l]

	nicefuns :: [([Char], Expression)]
	nicefuns = [(l, foldr ((o) o Lambda) id i e)\\(Function l i e)<-fs]

	vars :: Expression -> [[Char]]
	vars (Var v) = [v]
	vars (App l r) = vars l ++ vars r
	vars (Lambda l e) = flt l e
	vars (Let ns e) = flatten [[v\\v<-vars e | not (isMember v (map fst ns))]:map (uncurry flt) ns]
	vars _ = []

	flt i e = [v\\v<-vars e | v <> i]

instance toString Scheme where
	toString (Forall [] t) = toString t
	toString (Forall as t) = concat ["A.", join " " (map toString as), ": ", toString t]

:: TypeEnv :== Map [Char] Scheme
:: Subst :== Map [Char] Type

:: Infer a :== StateT [Int] (WriterT [([Char], Scheme)] (Either [String])) a

runInfer :: (Infer (Subst, Type)) -> Either [String] [([Char], Scheme)]
runInfer i = case runWriterT (evalStateT i [0..]) of
	Left e = Left e
	Right ((s, t), w) = pure [(['start'], generalize newMap (apply s t)):w]

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

(oo) infixl 9 :: Subst Subst -> Subst
(oo) 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) = fromMaybe t (get v s)
	apply s (t1 --> t2) = apply s t1 --> apply s t2
	apply s (TApp t1 t2) = TApp (apply s t1) (apply s t2)
	apply _ x = x
	
	ftv (TVar v) = [v]
	ftv (t1 --> t2) = on union ftv t1 t2
	ftv (TApp t1 t2) = on union ftv t1 t2
	ftv _ = []

instance Substitutable Scheme where
	apply s (Forall as t) = Forall as $ apply (foldr 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 (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

err :: [String] -> Infer a
err e = liftT (liftT (Left e))

unify :: Type Type -> Infer Subst
unify (l --> r) (l` --> r`)
	=        unify l l`
	>>= \s1->on unify (apply s1) r r`
	>>= \s2->pure (s1 oo s2)
unify (TVar a) (TVar t)
	| a == t = pure newMap
unify (TVar a) t
	| occursCheck a t = err ["Infinite type: ", toString a, " to ", toString t]
	= pure (singleton a t)
unify t (TVar a) = unify (TVar a) t
unify TInt TInt = pure newMap
unify TBool TBool = pure newMap
unify (TApp l r) (TApp l` r`)
	= unify l l`
	>>= \s1->on unify (apply s1) r r`
	>>= \s2->pure (s1 oo s2)
unify t1 t2 = err ["Cannot unify: ", toString t1, " with ", toString t2]

unifyl :: [Type] -> Infer Subst
unifyl [t1,t2:ts] = unify t1 t2 >>= \s->unifyl (map (apply s) [t2:ts])
unifyl _ = pure newMap

instantiate :: Scheme -> Infer Type
instantiate (Forall as t)
	=         sequence [fresh\\_<-as]
	>>= \as`->pure (apply (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 (newMap, TInt)
infer env (Lit (Bool _)) = pure (newMap, TBool)
infer env (Var x) = maybe (err ["Unbound variable: ", toString x])
	(\s->tuple newMap <$> instantiate s) $ get x env
infer env (App e1 e2)
	=              fresh
	>>= \tv->      infer env e1
	>>= \(s1, t1)->infer (apply s1 env) e2
	>>= \(s2, t2)->unify (apply s2 t1) (t2 --> tv)
	>>= \s3->      pure (s3 oo s2 oo s1, apply s3 tv)
infer env (Lambda x b)
	=              fresh
	>>= \tv->      infer ('Data.Map'.put x (Forall [] tv) env) b
	>>= \(s1, t1)->pure (s1, apply s1 tv --> t1)
//Non recursion
//infer env (Let [(x, e1)] e2)
//	=              infer env e1
//	>>= \(s1, t1)->infer ('Data.Map'.put x (generalize (apply s1 env) t1) env) e2
//	>>= \(s2, t2)->liftT (tell [(x, Forall [] t1)])
//	>>|            pure (s1 oo s2, t2)
//Single recursion
//infer env (Let [(x, e1)] e2)
//	=              fresh
//	>>= \tv->      let env` = 'Data.Map'.put x (Forall [] tv) env
//	               in infer env` e1
//	>>= \(s1,t1)-> infer ('Data.Map'.put x (generalize (apply s1 env`) t1) env`) e2
//	>>= \(s2, t2)->pure (s1 oo s2, t2)
//Multiple recursion
infer env (Let xs e2)
	# (ns, bs) = unzip xs
	=              sequence [fresh\\_<-ns]
	>>= \tvs->     let env` = foldr (\(k, v)->'Data.Map'.put k (Forall [] v)) env (zip2 ns tvs)
	               in  unzip <$> sequence (map (infer env`) bs)
	>>= \(ss,ts)-> unifyl ts
	>>= \s->       liftT (tell [(n, generalize (apply s env`) t)\\t<-ts & n<-ns])
	>>|            let env`` = foldr (\(n, s, t) m->'Data.Map'.put n (generalize (apply s env`) t) m) env` (zip3 ns ss ts)
	               in infer env`` e2
	>>= \(s2, t2)->pure (s oo s2, t2)