implementation module check

import StdEnv

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

import ast

check :: [Function] -> Either [String] (Expression, Scheme)
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)
//			= (\x->(e, x)) <$> runInfer (infer preamble (makeExpression fs e))
			= pure (makeExpression fs e, undef)
		([Function _ _ _], _) = Left ["Start cannot have arguments"]


:: Node a :== (a, [a])
:: SCCState a =
	{ index :: Int
	, stack :: [a]
	, map   :: Map a (Int, Int, Bool)
	, sccs  :: [[a]]
	}

import StdDebug
import Text.GenPrint
scc :: [Node a] -> [[a]] | Eq, Ord a
scc nodes = (foldr scc` {index=0,stack=[],map=newMap,sccs=[]} nodes).sccs
where
//	scc` :: (Node a) (SCCState a) -> SCCState a | Eq, Ord a
	scc` (v, suc) s = maybe (strongconnect s (v, suc)) (\_->s) $ get v s.map

//	strongconnect :: (SCCState a) (Node a)-> SCCState a | Eq, Ord a
	strongconnect s (v, suc)
		# s = flip (foldr processSucc) suc
			{ s
			& map   = 'Data.Map'.put v (s.index, s.index, True) s.map
			, stack = [v:s.stack]
			, index = s.index + 1
			}
		# (Just (iv, lv, lo)) = get v s.map
		| iv == lv
			# (scc, [sl:stack]) = span ((<>) v) s.stack
			# scc = scc ++ [sl]
			= { s
			  & sccs = [scc:s.sccs]
			  , stack= stack
			  , map  = foldr (\w m->'Data.Map'.put w (appThd3 (\_->False) $ fromJust (get w m)) m) s.map scc
			  }
		= s
	where
//		processSucc :: a (SCCState a) -> SCCState a | Eq, Ord a
		processSucc w s = case get w s.map of
			Nothing
				# s = strongconnect s $ hd [l\\l=:(n, _)<-nodes | n == w]
				# (Just (iw, lw, ow)) = get w s.map
				# (Just (iv, lv, ov)) = get v s.map
				= {s & map='Data.Map'.put v (iv, min lv lw, ov) s.map}
			Just (iw, lw, True)
				# (Just (iv, lv, ov)) = get v s.map
				= {s & map='Data.Map'.put v (iv, min iw lv, ov) s.map}
			Just _ = s

makeExpression :: [Function] Expression -> Expression
makeExpression fs start
	= mkExpr $ scc [(l, vars e [])\\(l, e)<-nicefuns]
where
	mkExpr :: [[[Char]]] -> Expression
	mkExpr t = trace_n (printToString t) start
	nicefuns = [(l, foldr ((o) o Lambda) id i e)\\(Function l i e)<-fs]

	vars :: Expression [[Char]] -> [[Char]]
	vars (Var v=:[m:_]) c
		| m <> '_' = [v:c]
	vars (App l r) c = vars l $ vars r c
	vars (Lambda l e) c = [v\\v<-vars e c | v <> l]
	vars (Let ns e) c = vars e c // TODO
	vars _ c = c

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

instance toString Type where
	toString (TVar a) = toString a
	toString TInt = "Int"
	toString TBool = "Bool"
	toString (a --> b) = concat ["(", toString a, ") -> ", toString b]

:: TypeEnv :== Map [Char] Scheme
preamble :: TypeEnv
preamble = fromList
	[(['_if'],  Forall [['_ift']]
		$ TBool --> TVar ['_ift'] --> TVar ['_ift'] --> TVar ['_ift'])
	,(['_eq'],  Forall [['_eq']]  $ TInt --> TInt --> TBool)
	,(['_mul'], Forall [['_mul']] $ TInt --> TInt --> TInt)
	,(['_add'], Forall [['_add']] $ TInt --> TInt --> TInt)
	,(['_sub'], Forall [['_sub']] $ TInt --> TInt --> TInt)
	]
:: Subst :== Map [Char] Type

:: Infer a :== StateT [Int] (Either [String]) a
runInfer :: (Infer (Subst, Type)) -> Either [String] Scheme
runInfer i = uncurry ((o) (generalize newMap) o apply)
	<$> evalStateT i [0..]

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 _ x = x
	
	ftv (TVar v) = [v]
	ftv (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

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 = liftT (Left ["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 t1 t2 = liftT (Left ["Cannot unify: ", toString t1, " with ", toString t2])

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) = case get x env of
	Nothing = liftT (Left ["Unbound variable: ", toString x])
	Just s = (\x->(newMap, x)) <$> 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) (t2 --> tv)
	>>= \s3->      pure (s1 oo s2 oo s3, 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)
//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)->pure (s1 oo s2, t2)
infer env (Let [(x, e1)] e2)
	=              fresh
	>>= \tv->      let env` = 'Data.Map'.put x (Forall [] tv) env
	               in infer env` e1
	>>= \(s1,t1)-> unify t1 tv
	>>= \t->infer ('Data.Map'.put x (generalize (apply s1 env`) t1) env`) e2
	>>= \(s2, t2)->pure (s1 oo s2, t2)
//infer env (Let xs e2)
//	# (ns, bs) = unzip xs
//	=              sequence [fresh\\_<-ns]
//	>>= \tvs->     let env` = foldr (uncurry putenv) env (zip2 ns tvs)
//	               in  unzip <$> sequence (map infer env`) bs
//	>>= \(ss,ts)-> let s = foldr (oo) newMap ss
//	               in  //unify t1 tv
//	>>= \t->infer ('Data.Map'.put x (generalize (apply s1 env`) t1) env`) e2
//	>>= \(s2, t2)->pure (s1 oo s2, t2)
where
	putenv :: [Char] -> (Type TypeEnv -> TypeEnv)
	putenv k = 'Data.Map'.put k o Forall []

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