implementation module infer

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 Data.GenEq
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

derive gEq Type
:: TypeEnv :== 'DM'.Map String Scheme
:: Constraint :== (Type, Type)
:: Subst :== 'DM'.Map String Type
:: Unifier :== (Subst, [Constraint])
:: Infer a :== RWST TypeEnv [Constraint] InferState (Either TypeError) a
:: InferState = { count :: Int }
:: Solve a :== 'MS'.StateT Unifier (Either TypeError) a
:: TypeError
	= UnboundVariable String
	| UnificationFail Type Type
	| UnificationMismatch [Type] [Type]
	| InfiniteType String Type

nullSubst :: Subst
nullSubst = 'DM'.newMap

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

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

letters :: [String]
letters = map toString $ [1..] >>= flip replicateM ['a'..'z']

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

lookupEnv :: String -> Infer Type
lookupEnv x = asks ('DM'.get x)
	>>= maybe (liftT $ Left $ UnboundVariable x) instantiate

class Substitutable a
where
	apply :: Subst a -> a
	ftv   :: a -> 'DS'.Set String

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

	ftv (VarType a)   = 'DS'.singleton a
	ftv (Arrow t1 t2) = 'DS'.union (ftv t1) (ftv t2)
	ftv t             = 'DS'.newSet

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

instance Substitutable [a] | Substitutable a
where
	apply s ls = map (apply s) ls
	ftv t = foldr ('DS'.union o ftv) 'DS'.newSet 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) = 'DS'.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 ('DS'.toList $ 'DS'.difference (ftv t) (ftv env)) t

infer :: Expression -> Infer Type
infer (Literal v) = case v of
	Int  _ = pure IntType
	Bool _ = pure BoolType
	Char _ = pure CharType
infer (Variable s) = lookupEnv s
infer (Apply e1 e2)
	=        infer e1
	>>= \t1->infer e2
	>>= \t2->fresh
	>>= \tv->uni t1 (Arrow t2 tv)
	>>|      pure tv
infer (Lambda s e)
	=        fresh
	>>= \tv->inEnv (s, Forall [] tv) (infer e)
	>>= \t-> pure (Arrow tv t)
infer (Let x e1 e2)
	=        ask
	>>= \en->infer e1
	>>= \t1->inEnv (x, generalize en t1) (infer e2)
infer (Code c) = fresh >>= \v->pure case c of
	"add" = Arrow v (Arrow v v)
	"sub" = Arrow v (Arrow v v)
	"mul" = Arrow v (Arrow v v)
	"div" = Arrow v (Arrow v v)
	"and" = Arrow v (Arrow v BoolType)
	"or"  = Arrow v (Arrow v BoolType)

unifies :: Type Type -> Solve Unifier
unifies t1 t2
	| t1 === t2 = pure ('DM'.newMap, [])
unifies (VarType v) t = tbind v t
unifies t (VarType v) = tbind v t
unifies (Arrow t1 t2) (Arrow t3 t4) = unifyMany [t1, t2] [t3, t4]
unifies t1 t2 = liftT (Left (UnificationFail t1 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 (UnificationMismatch t1 t2))

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

tbind ::  String Type -> Solve Unifier
tbind a t
	| t === VarType a = pure ('DM'.newMap, [])
	| occursCheck a t = liftT $ Left $ InfiniteType a t
	= pure $ ('DM'.singleton a t, [])

occursCheck ::  String a -> Bool | Substitutable a
occursCheck a t = 'DS'.member 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