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 => qualified join
import Control.Monad.Trans
import Control.Monad.State
import Text

import qualified Data.Map
from Data.Map import :: Map(..), newMap, get, instance Functor (Map k)

import ast

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 (foldr ((o) o Lambda) id a e)) e fs
			= case runInfer (infer newMap e) of
				Left e = Left e
				Right s = Left [toString s]
		([Function _ _ _], _) = Left ["Start cannot have arguments"]

instance toString Scheme where
	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 (TFun a b) = concat ["(", toString a, ") -> ", toString b]

:: TypeEnv :== Map [Char] Scheme
:: 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]]))

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) = fromMaybe t (get 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 'Data.Map'.newMap
unify TBool TBool = pure 'Data.Map'.newMap
unify t1 t2 = liftT (Left ["Cannot unify: ", toString t1, " with ", toString t2])

bind :: [Char] Type -> Infer Subst
bind a (TVar t) | a == t = pure 'Data.Map'.newMap
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 ('Data.Map'.newMap, TInt)
infer env (Lit (Bool _)) = pure ('Data.Map'.newMap, TBool)
infer env (Var x) = case 'Data.Map'.get x env of
	Nothing = liftT (Left ["Unbound variable: ", toString x])
	Just s = tuple 'Data.Map'.newMap <$> 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)