HM type inference

[fp.git] / src / HM.hs

diff --git a/src/HM.hs b/src/HM.hs
index 065ccbb..9c4aa76 100644 (file)
--- a/src/HM.hs
+++ b/src/HM.hs
@@ -6,7 +6,7 @@
 -- Maintainer  :  tomik.musil@gmail.com
 -- Stability   :  experimental
 --
 -- 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
 
 module HM
   ( -- * Types


@@ -16,42 +16,114 @@ module HM
   , TypedTerm(..)
     -- * Type inference
   , algW
   , TypedTerm(..)
     -- * Type inference
   , algW

+  , runTI

   ) where
 
   ) where
 

+import Control.Monad.Except

 import Control.Monad.State
 import Control.Monad.State

+import qualified Data.Set as Set
+import qualified Data.Map as Map

 
 import HM.Term
 import HM.Parser
 
 
 import HM.Term
 import HM.Parser
 

-type Substitution = TypeScheme -> TypeScheme
+type Substitution = Map.Map TypeVarName Type

 type VarS a = State Int a
 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
-  
+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

 
 

-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
+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
+
+generalize :: TypeEnv -> Type -> TypeScheme
+generalize e t = foldr TSForAll (TScheme t) vars
+  where
+    vars = Set.toList $ freeVarsT t Set.\\ Set.unions (map freeVarsS $ Map.elems e)
+
+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)
+ti e (NTTerm (App a b)) = do
+  tv <- newVar 
+  (s1, t1) <- ti e a
+  (s2, t2) <- ti (Map.map (substituteS s1) e) b
+  s3 <- unify (substituteT s2 t1) (TypeFunction t2 tv)
+  return (s3 `composeSub` s2 `composeSub` s1, substituteT s3 tv)
+ti e (NTTerm (Let x a b)) = do
+  (s1, t1) <- ti e a
+  let t' = generalize (Map.map (substituteS s1) e) t1
+      e' = Map.insert x t' e
+  (s2, t2) <- ti (Map.map (substituteS s1) e') b
+  return (s1 `composeSub` s2, t2)
+  
+ 
+algW :: TypedTerm -> TI Type
+algW t = do
+  (s, u) <- ti Map.empty t
+  return $ substituteT s u