implementation module check

import StdEnv

import qualified Data.Map as DM
from Data.Map import instance Functor (Map k)
import qualified Data.Set as DS
import Data.Functor
import Data.Func
import Data.Either
import Data.List
import Data.Maybe
import Control.Applicative
import Control.Monad
import Control.Monad.Trans
import qualified Control.Monad.State as MS
import Control.Monad.State => qualified gets, put, modify
import Control.Monad.RWST => qualified put

import ast

check :: AST -> Either [String] (AST, [([Char], Scheme)])
check (AST fs) = pure (AST fs, [])/*case inferAST preamble fs of
	Left e = Left e
	Right s = Right (AST fs, 'DM'.toList s)
where
	preamble = 'DM'.fromList
		[(['if'], Forall [['a']] $ TFun TBool $ TFun (TVar ['a']) $ TFun (TVar ['a']) $ TVar ['a'])
		,(['eq'], Forall [] $ TFun TInt $ TFun TInt TBool)
		,(['mul'], Forall [] $ TFun TInt $ TFun TInt TInt)
		,(['div'], Forall [] $ TFun TInt $ TFun TInt TInt)
		,(['add'], Forall [] $ TFun TInt $ TFun TInt TInt)
		,(['sub'], Forall [] $ TFun TInt $ TFun TInt TInt)
		]
*/

:: TypeEnv    :== 'DM'.Map [Char] Scheme
:: Constraint :== (Type, Type)
:: Subst      :== 'DM'.Map [Char] Type
:: Unifier    :== (Subst, [Constraint])
:: Infer a    :== RWST TypeEnv [Constraint] InferState (Either [String]) a
:: InferState =   { count :: Int }
:: Scheme     =   Forall [[Char]] Type
:: Solve a    :== StateT Unifier (Either [String]) a

nullSubst :: Subst
nullSubst = 'DM'.newMap

uni :: Type Type -> Infer ()
uni t1 t2 = tell [(t1, t2)]

inEnv :: ([Char], Scheme) (Infer a) -> Infer a
inEnv (x, sc) m = local (\e->'DM'.put x sc $ 'DM'.del x e) m

letters :: [[Char]]
letters = [1..] >>= flip replicateM ['a'..'z']

fresh :: Infer Type
fresh = get >>= \s=:{count}->'Control.Monad.RWST'.put {s & count=count + 1} >>| pure (TVar $ letters !! count)

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

instance Substitutable Type
where
	apply s t=:(TVar a) = maybe t id $ 'DM'.get a s
	apply s (TFun t1 t2)  = TFun (apply s t1) (apply s t2)
	apply _ t = t

	ftv (TVar a)     = [a]
	ftv (TFun t1 t2) = union (ftv t1) (ftv t2)
	ftv t            = []

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

instance Substitutable [a] | Substitutable a
where
	apply s ls = map (apply s) ls
	ftv t = foldr (union o ftv) [] t

instance Substitutable TypeEnv where
	apply s env = fmap (apply s) env
	ftv env = ftv $ 'DM'.elems env

instance Substitutable Constraint where
	apply s (t1, t2) = (apply s t1, apply s t2)
	ftv (t1, t2) = union (ftv t1) (ftv t2)

instantiate ::  Scheme -> Infer Type
instantiate (Forall as t) = mapM (const fresh) as
	>>= \as`->let s = 'DM'.fromList $ zip2 as as` in pure $ apply s t

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

//:: Expression
//	= Lit Value
//	| Var [Char]
//	| App Expression Expression
//	| Lambda [Char] Expression
//	| Builtin [Char] [Expression]
inferExpr :: TypeEnv Expression -> Either [String] Scheme
inferExpr env ex = case runRWST (infer ex) env {count=0} of
	Left e = Left e
	Right (ty, st, cs) = case runStateT solver ('DM'.newMap, cs) of
		Left e = Left e
		Right (s, _) = Right (closeOver (apply s ty))

closeOver :: Type -> Scheme
closeOver t = normalize (generalize 'DM'.newMap t)

normalize :: Scheme -> Scheme
normalize t = t 

inferAST :: TypeEnv [Function] -> Either [String] TypeEnv
inferAST env [] = Right env
inferAST env [Function name args body:rest] = case inferExpr env (foldr Lambda body args) of
	Left e = Left e
	Right ty = inferAST ('DM'.put name ty env) rest

inferFunc :: [Function] -> Infer ()
inferFunc [] = pure () 
inferFunc [Function name args body:rest]
	=      ask
	>>= \en->infer (foldr Lambda body args)
	>>= \t1->inEnv (name, generalize en t1) (inferFunc rest)
	>>= \_->pure ()

infer :: Expression -> Infer Type
infer (Lit v) = case v of
	Int  _ = pure TInt
	Bool _ = pure TBool
infer (Var s) = asks ('DM'.get s)
	>>= maybe (liftT $ Left ["Unbound variable " +++ toString s]) instantiate
infer (App e1 e2)
	=        infer e1
	>>= \t1->infer e2
	>>= \t2->fresh
	>>= \tv->uni t1 (TFun t2 tv)
	>>|      pure tv
infer (Lambda s e)
	=        fresh
	>>= \tv->inEnv (s, Forall [] tv) (infer e)
	>>= \t-> pure (TFun tv t)
//infer (Let x e1 e2)
//	=        ask
//	>>= \en->infer e1
//	>>= \t1->inEnv (x, generalize en t1) (infer e2)

unifies :: Type Type -> Solve Unifier
unifies TInt TInt = pure ('DM'.newMap, [])
unifies TBool TBool = pure ('DM'.newMap, [])
unifies (TVar a) (TVar b) 
	| a == b = pure ('DM'.newMap, [])
unifies (TVar v) t = tbind v t
unifies t (TVar v) = tbind v t
unifies (TFun t1 t2) (TFun t3 t4) = unifyMany [t1, t2] [t3, t4]
unifies t1 t2 = liftT $ Left ["Cannot unify " +++ toString t1 +++ " with " +++ toString t2]

unifyMany :: [Type] [Type] -> Solve Unifier
unifyMany [] [] = pure ('DM'.newMap, [])
unifyMany [t1:ts1] [t2:ts2] = unifies t1 t2
	>>= \(su1, cs1)->unifyMany (apply su1 ts1) (apply su1 ts2)
	>>= \(su2, cs2)->pure (su2 `compose` su1, cs1 ++ cs2)
unifyMany t1 t2 = liftT $ Left ["Length difference in unifyMany"]

(`compose`) infix 1 :: Subst Subst -> Subst
(`compose`) s1 s2 = 'DM'.union (apply s1 <$> s2) s1

tbind ::  [Char] Type -> Solve Unifier
tbind a (TVar b)
	| a == b = pure ('DM'.newMap, [])
tbind a t
	| occursCheck a t = liftT $ Left ["Infinite type " +++ toString a +++ toString t]
	= pure $ ('DM'.singleton a t, [])

occursCheck ::  [Char] a -> Bool | Substitutable a
occursCheck a t = isMember a $ ftv t

solver :: Solve Subst
solver = getState >>= \(su, cs)->case cs of
	[] = pure su
	[(t1, t2):cs0] = unifies t1 t2
		>>= \(su1, cs1)->'MS'.put (su1 `compose` su, cs1 ++ (apply su1 cs0))
		>>| solver