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   , runTI
  20   ) where
  21
  22 import Control.Monad.Except
  23 import Control.Monad.State
  24 import qualified Data.Set as Set
  25 import qualified Data.Map as Map
  26
  27 import HM.Term
  28 import HM.Parser
  29
  30 type Substitution = Map.Map TypeVarName Type
  31 type VarS a = State Int a
  32 type TypeEnv = Map.Map VarName TypeScheme
  33 data TIState = TIState {tiSupply :: Int} deriving (Show)
  34 type TI a = ExceptT String (State TIState) a
  35
  36 runTI :: TI a -> (Either String a, TIState)
  37 runTI t = runState (runExceptT t) $ TIState 0
  38
  39 newVar :: TI Type
  40 newVar = do
  41   s <- get
  42   put s {tiSupply = tiSupply s + 1}
  43   return (TypeVar $ "a" ++ show (tiSupply s))
  44
  45 freeVarsT :: Type -> Set.Set TypeVarName
  46 freeVarsT (Primitive _) = Set.empty
  47 freeVarsT (TypeVar t) = Set.singleton t
  48 freeVarsT (TypeFunction a b) = freeVarsT a `Set.union` freeVarsT b
  49
  50 freeVarsS :: TypeScheme -> Set.Set TypeVarName
  51 freeVarsS (TScheme t) = freeVarsT t
  52 freeVarsS (TSForAll v s) = v `Set.delete` freeVarsS s
  53
  54 substituteT :: Substitution -> Type -> Type
  55 substituteT _ t@(Primitive _) = t
  56 substituteT s t@(TypeVar v) = Map.findWithDefault t v s
  57 substituteT s (TypeFunction a b) = TypeFunction (substituteT s a) (substituteT s b)
  58
  59 substituteS :: Substitution -> TypeScheme -> TypeScheme
  60 substituteS s (TScheme t) = TScheme $ substituteT s t
  61 substituteS s (TSForAll v t) = TSForAll v $ substituteS (v `Map.delete` s) t
  62
  63 idSub :: Substitution
  64 idSub = Map.empty
  65
  66 composeSub :: Substitution -> Substitution -> Substitution
  67 composeSub s1 s2 = Map.map (substituteT s1) s2 `Map.union` s1
  68
  69 varBind :: TypeVarName -> Type -> TI Substitution
  70 varBind v t | t == TypeVar v = return idSub
  71             | v `Set.member` freeVarsT t = fail $ "occur check failed: " ++ v ++ " in " ++ show t
  72             | otherwise = return $ Map.singleton v t
  73
  74 instantiate :: TypeScheme -> TI Type
  75 instantiate (TScheme t) = return t
  76 instantiate (TSForAll v t) = do
  77   nv <- newVar
  78   instantiate $ substituteS (Map.singleton v nv) t
  79
  80 unify :: Type -> Type -> TI Substitution
  81 unify (TypeVar a) t = varBind a t
  82 unify t (TypeVar a) = varBind a t
  83 unify (TypeFunction a b) (TypeFunction a' b') = do
  84   s1 <- unify a a'
  85   s2 <- unify b b'
  86   return $ s1 `composeSub` s2
  87 unify (Primitive a) (Primitive b) | a == b = return idSub
  88 unify a b = fail $ "cannot unify " ++ show a ++ " with " ++ show b
  89
  90 ti :: TypeEnv -> TypedTerm -> TI (Substitution, Type)
  91 --ti _ (TTerm (Var v) (TScheme t@(Primitive _))) = return (idSub, t)
  92 ti e (TTerm tr sch) = do
  93   (s, t) <- ti e (NTTerm tr)
  94   sch' <- instantiate sch
  95   s' <- unify t sch'
  96   return (s', substituteT s' sch')
  97 ti e (NTTerm (Var v)) = case Map.lookup v e of
  98   Nothing -> fail $ "unbound variable: " ++ v
  99   Just sigma -> do
 100     t <- instantiate sigma
 101     return (idSub, t)
 102 ti e (NTTerm (Lam x y)) = do
 103   tv <- newVar
 104   let e' = Map.insert x (TScheme tv) e
 105   (s, t) <-  ti e' y
 106   return (s, TypeFunction (substituteT s tv) t)
 107
 108 algW :: TypedTerm -> TI Type
 109 algW t = do
 110   (s, u) <- ti Map.empty t
 111   return $ substituteT s u