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 Control.Applicative
import Control.Monad
import Control.Monad.Trans
import qualified Control.Monad.State as MS
import Control.Monad.RWST

import ast

:: 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

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}->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

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)