src/HM.hs

   1 -- |
   2 -- Module      :  HM
   3 -- Copyright   :  Tomáš Musil 2014
   4 -- License     :  BSD-3
   5 --
   6 -- Maintainer  :  tomik.musil@gmail.com
   7 -- Stability   :  experimental
   8 --
   9 -- This is a toy implementation of \-calculus with Hindley-Milner type system.
  10
  11 module HM
  12   ( -- * Types
  13     Type(..)
  14   , TypeScheme(..)
  15   , Term(..)
  16   , TypedTerm(..)
  17     -- * Type inference
  18   , algW
  19   ) where
  20
  21 import Control.Monad.Except
  22 import Control.Monad.State
  23 import qualified Data.Set as Set
  24 import qualified Data.Map as Map
  25
  26 import HM.Term
  27 import HM.Parser
  28
  29 type Substitution = Map.Map TypeVarName Type
  30 type VarS a = State Int a
  31 type TypeEnv = Map.Map VarName TypeScheme
  32 data TIState = TIState {tiSupply :: Int} deriving (Show)
  33 type TI a = ExceptT String (State TIState) a
  34
  35 runTI :: TI a -> (Either String a, TIState)
  36 runTI t = runState (runExceptT t) $ TIState 0
  37
  38 newVar :: TI Type
  39 newVar = do
  40   s <- get
  41   put s {tiSupply = tiSupply s + 1}
  42   return (TypeVar $ "a" ++ show (tiSupply s))
  43
  44 freeVarsT :: Type -> Set.Set TypeVarName
  45 freeVarsT (Primitive _) = Set.empty
  46 freeVarsT (TypeVar t) = Set.singleton t
  47 freeVarsT (TypeFunction a b) = freeVarsT a `Set.union` freeVarsT b
  48
  49 freeVarsS :: TypeScheme -> Set.Set TypeVarName
  50 freeVarsS (TScheme t) = freeVarsT t
  51 freeVarsS (TSForAll v s) = v `Set.delete` freeVarsS s
  52
  53 substituteT :: Substitution -> Type -> Type
  54 substituteT _ t@(Primitive _) = t
  55 substituteT s t@(TypeVar v) = Map.findWithDefault t v s
  56 substituteT s (TypeFunction a b) = TypeFunction (substituteT s a) (substituteT s b)
  57
  58 substituteS :: Substitution -> TypeScheme -> TypeScheme
  59 substituteS s (TScheme t) = TScheme $ substituteT s t
  60 substituteS s (TSForAll v t) = TSForAll v $ substituteS (v `Map.delete` s) t
  61
  62 idSub :: Substitution
  63 idSub = Map.empty
  64
  65 composeSub :: Substitution -> Substitution -> Substitution
  66 composeSub s1 s2 = Map.map (substituteT s1) s2 `Map.union` s1
  67
  68 varBind :: TypeVarName -> Type -> TI Substitution
  69 varBind v t | t == TypeVar v = return idSub
  70             | v `Set.member` freeVarsT t = fail $ "occur check failed: " ++ v ++ " in " ++ show t
  71             | otherwise = return $ Map.singleton v t
  72
  73 instantiate :: TypeScheme -> TI Type
  74 instantiate (TScheme t) = return t
  75 instantiate (TSForAll v t) = do
  76   nv <- newVar
  77   instantiate $ substituteS (Map.singleton v nv) t
  78
  79 generalize :: TypeEnv -> Type -> TypeScheme
  80 generalize e t = foldr TSForAll (TScheme t) vars
  81   where
  82     vars = Set.toList $ freeVarsT t Set.\\ Set.unions (map freeVarsS $ Map.elems e)
  83
  84 unify :: Type -> Type -> TI Substitution
  85 unify (TypeVar a) t = varBind a t
  86 unify t (TypeVar a) = varBind a t
  87 unify (TypeFunction a b) (TypeFunction a' b') = do
  88   s1 <- unify a a'
  89   s2 <- unify b b'
  90   return $ s1 `composeSub` s2
  91 unify (Primitive a) (Primitive b) | a == b = return idSub
  92 unify a b = fail $ "cannot unify " ++ show a ++ " with " ++ show b
  93
  94 ti :: TypeEnv -> TypedTerm -> TI (Substitution, Type)
  95 --ti _ (TTerm (Var v) (TScheme t@(Primitive _))) = return (idSub, t)
  96 ti e (TTerm tr sch) = do
  97   (s, t) <- ti e (NTTerm tr)
  98   sch' <- instantiate sch
  99   s' <- unify t sch'
 100   return (s', substituteT s' sch')
 101 ti e (NTTerm (Var v)) = case Map.lookup v e of
 102   Nothing -> fail $ "unbound variable: " ++ v
 103   Just sigma -> do
 104     t <- instantiate sigma
 105     return (idSub, t)
 106 ti e (NTTerm (Lam x y)) = do
 107   tv <- newVar
 108   let e' = Map.insert x (TScheme tv) e
 109   (s, t) <-  ti e' y
 110   return (s, TypeFunction (substituteT s tv) t)
 111 ti e (NTTerm (App a b)) = do
 112   tv <- newVar
 113   (s1, t1) <- ti e a
 114   (s2, t2) <- ti (Map.map (substituteS s1) e) b
 115   s3 <- unify (substituteT s2 t1) (TypeFunction t2 tv)
 116   return (s3 `composeSub` s2 `composeSub` s1, substituteT s3 tv)
 117 ti e (NTTerm (Let x a b)) = do
 118   (s1, t1) <- ti e a
 119   let t' = generalize (Map.map (substituteS s1) e) t1
 120       e' = Map.insert x t' e
 121   (s2, t2) <- ti (Map.map (substituteS s1) e') b
 122   return (s1 `composeSub` s2, t2)
 123
 124
 125 algW :: TypedTerm -> Either String Type
 126 algW t = fst . runTI $ do
 127   (s, u) <- ti Map.empty t
 128   return $ substituteT s u