-- Maintainer : tomik.musil@gmail.com
-- Stability : experimental
--
--- This is a toy implementation of λ-calculus with Hindley-Milner type system.
+-- This is a toy implementation of \-calculus with Hindley-Milner type system.
module HM
( -- * Types
, TypedTerm(..)
-- * Type inference
, algW
+ , runTI
) where
+import Control.Monad.Except
import Control.Monad.State
+import qualified Data.Set as Set
+import qualified Data.Map as Map
import HM.Term
import HM.Parser
-type Substitution = TypeScheme -> TypeScheme
+type Substitution = Map.Map TypeVarName Type
type VarS a = State Int a
+type TypeEnv = Map.Map VarName TypeScheme
+data TIState = TIState {tiSupply :: Int} deriving (Show)
+type TI a = ExceptT String (State TIState) a
-fresh :: TypeVarName
-fresh = undefined
---fresh = do
--- n <- get
--- "a" ++ show n
-
-
-substitute :: TypeScheme -> TypeVarName -> TypeScheme -> TypeScheme
-substitute = undefined
-
-unify :: TypeScheme -> TypeScheme -> Either String Substitution
-unify (TScheme (Primitive a)) (TScheme (Primitive b)) | a == b = Right id
-unify (TScheme (TypeVar a)) (TScheme (TypeVar b)) | a == b = Right id
-unify a b = Left $ "cannot unify " ++ show a ++ " with " ++ show b
-
-algW :: TypedTerm -> Either String TypeScheme
-algW (TTerm (Var _) t) = Right t
---algW (TTerm (Lam x t) (TScheme p)) = do
--- let v = TScheme (TypeVar fresh)
--- np = substitute v x t
--- unify p np
---algW (TTerm (App u v) t) = do
--- tu <- algW u
--- tv <- algW v
--- case tu of
--- (TScheme (TypeFunction a b)) -> do
--- unify a tv
--- return b
--- _ -> Left $ "cannot apply " ++ show tu ++ " to " ++ show tv
+runTI :: TI a -> (Either String a, TIState)
+runTI t = runState (runExceptT t) $ TIState 0
+
+newVar :: TI Type
+newVar = do
+ s <- get
+ put s {tiSupply = tiSupply s + 1}
+ return (TypeVar $ "a" ++ show (tiSupply s))
+
+freeVarsT :: Type -> Set.Set TypeVarName
+freeVarsT (Primitive _) = Set.empty
+freeVarsT (TypeVar t) = Set.singleton t
+freeVarsT (TypeFunction a b) = freeVarsT a `Set.union` freeVarsT b
+
+freeVarsS :: TypeScheme -> Set.Set TypeVarName
+freeVarsS (TScheme t) = freeVarsT t
+freeVarsS (TSForAll v s) = v `Set.delete` freeVarsS s
+
+substituteT :: Substitution -> Type -> Type
+substituteT _ t@(Primitive _) = t
+substituteT s t@(TypeVar v) = Map.findWithDefault t v s
+substituteT s (TypeFunction a b) = TypeFunction (substituteT s a) (substituteT s b)
+
+substituteS :: Substitution -> TypeScheme -> TypeScheme
+substituteS s (TScheme t) = TScheme $ substituteT s t
+substituteS s (TSForAll v t) = TSForAll v $ substituteS (v `Map.delete` s) t
+
+idSub :: Substitution
+idSub = Map.empty
+
+composeSub :: Substitution -> Substitution -> Substitution
+composeSub s1 s2 = Map.map (substituteT s1) s2 `Map.union` s1
+
+varBind :: TypeVarName -> Type -> TI Substitution
+varBind v t | t == TypeVar v = return idSub
+ | v `Set.member` freeVarsT t = fail $ "occur check failed: " ++ v ++ " in " ++ show t
+ | otherwise = return $ Map.singleton v t
+
+instantiate :: TypeScheme -> TI Type
+instantiate (TScheme t) = return t
+instantiate (TSForAll v t) = do
+ nv <- newVar
+ instantiate $ substituteS (Map.singleton v nv) t
+
+unify :: Type -> Type -> TI Substitution
+unify (TypeVar a) t = varBind a t
+unify t (TypeVar a) = varBind a t
+unify (TypeFunction a b) (TypeFunction a' b') = do
+ s1 <- unify a a'
+ s2 <- unify b b'
+ return $ s1 `composeSub` s2
+unify (Primitive a) (Primitive b) | a == b = return idSub
+unify a b = fail $ "cannot unify " ++ show a ++ " with " ++ show b
+
+ti :: TypeEnv -> TypedTerm -> TI (Substitution, Type)
+--ti _ (TTerm (Var v) (TScheme t@(Primitive _))) = return (idSub, t)
+ti e (TTerm tr sch) = do
+ (s, t) <- ti e (NTTerm tr)
+ sch' <- instantiate sch
+ s' <- unify t sch'
+ return (s', substituteT s' sch')
+ti e (NTTerm (Var v)) = case Map.lookup v e of
+ Nothing -> fail $ "unbound variable: " ++ v
+ Just sigma -> do
+ t <- instantiate sigma
+ return (idSub, t)
+ti e (NTTerm (Lam x y)) = do
+ tv <- newVar
+ let e' = Map.insert x (TScheme tv) e
+ (s, t) <- ti e' y
+ return (s, TypeFunction (substituteT s tv) t)
+
+algW :: TypedTerm -> TI Type
+algW t = do
+ (s, u) <- ti Map.empty t
+ return $ substituteT s u