src/Core/HindleyMilner.hs

@@ -1,3 +1,7 @@
+{-|
+Module      : Core.HindleyMilner
+Description : Hindley-Milner inference
+-}
 {-# LANGUAGE LambdaCase #-}
 module Core.HindleyMilner
     ( infer
@@ -15,28 +19,65 @@ import Control.Monad.State
 import Core.Syntax
 ----------------------------------------------------------------------------------
 
+-- | Annotated typing context -- I have a feeling we're going to want this in the
+-- future.
 type Context b = [(b, Type)]
 
+-- | Unannotated typing context, AKA our beloved Γ.
 type Context' = Context Name
 
 -- TODO: Errorful monad?
-data TypeError = TyErrCouldNotUnify Type Type
-               | TyErrRecursiveType Name Type
-               deriving Show
+
+-- | Type error enum.
+data TypeError
+    -- | Two types could not be unified
+    = TyErrCouldNotUnify Type Type
+    -- | @x@ could not be unified with @t@ because @x@ occurs in @t@
+    | TyErrRecursiveType Name Type
+    -- | Untyped, potentially undefined variable
+    | TyErrUntypedVariable Name
+    deriving (Show, Eq)
+
+-- | Synonym for @Either TypeError@
+type HMError = Either TypeError
+
+-- | Infer the type of an expression under some context.
+--
+-- >>> let g1 = [("id", TyVar "a" :-> TyVar "a")]
+-- >>> let g2 = [("id", (TyVar "a" :-> TyVar "a") :-> TyVar "a" :-> TyVar "a")]
+-- >>> infer g1 [coreExpr|id 3|]
+-- Right TyInt
+-- >>> infer g2 [coreExpr|id 3|]
+-- Left (TyErrCouldNotUnify (TyVar "a" :-> TyVar "a") TyInt)
 
 infer :: Context' -> Expr' -> Either TypeError Type
-infer g e = foldr (uncurry subst) t <$> unify cs where
-    (t,cs) = gather g e
+infer g e = do
+    (t,cs) <- gather g e
+    foldr (uncurry subst) t <$> unify cs
 
+-- | A @Constraint@ between two types describes the requirement that the pair
+-- must unify
 type Constraint = (Type, Type)
 
-gather :: Context' -> Expr' -> (Type, [Constraint])
-gather = \g e -> let (t,(cs,_)) = runState (go g e) ([],0) in (t,cs) where
-    go :: Context' -> Expr' -> State ([Constraint], Int) Type
+-- | Type of an expression under some context, and gather the constraints
+-- necessary to unify. Note that this is not the same as @infer@, as the
+-- expression will likely be given a fresh type variable along with a
+-- constraint, rather than the solved type.
+--
+-- For example, if the context says "@id@ has type a -> a," in an application of
+-- @id 3@, the whole application is assigned type @$a0@ and the constraint that
+-- @id@ must unify with type @Int -> $a0@ is generated.
+--
+-- >>> gather [("id", TyVar "a" :-> TyVar "a")] [coreExpr|id 3|]
+-- (TyVar "$a0",[(TyVar "a" :-> TyVar "a",TyInt :-> TyVar "$a0")])
+
+gather :: Context' -> Expr' -> HMError (Type, [Constraint])
+gather = \g e -> runStateT (go g e) ([],0) <&> \ (t,(cs,_)) -> (t,cs) where
+    go :: Context' -> Expr' -> StateT ([Constraint], Int) HMError Type
     go g = \case
         LitE (IntL _) -> pure TyInt
-        Var k         -> maybe e pure $ lookup k g
-            where e = error $ "variable `" <> k <> "' untyped in Γ"
+        Var k         -> lift $ maybe e Right $ lookup k g
+            where e = Left (TyErrUntypedVariable k)
         App f x       -> do
             tf <- go g f
             tx <- go g x
@@ -44,18 +85,23 @@ gather = \g e -> let (t,(cs,_)) = runState (go g e) ([],0) in (t,cs) where
             addConstraint tf (tx :-> tfx)
             pure tfx
 
-uniqueVar :: State ([Constraint], Int) Type
+uniqueVar :: StateT ([Constraint], Int) HMError Type
 uniqueVar = do
     n <- use _2
     _2 %= succ
     pure (TyVar $ '$' : 'a' : show n)
 
-addConstraint :: Type -> Type -> State ([Constraint], Int) ()
+addConstraint :: Type -> Type -> StateT ([Constraint], Int) HMError ()
 addConstraint t u = _1 %= ((t, u):)
 
-unify :: [Constraint] -> Either TypeError Context'
+-- | Unify a list of constraints, meaning that pairs between types are turned
+-- into pairs of type variables and types. A useful thought model is to think of
+-- it like solving an equation such that the unknown variable is the left-hand
+-- side.
+
+unify :: [Constraint] -> HMError Context'
 unify = go mempty where
-    go :: Context' -> [Constraint] -> Either TypeError Context'
+    go :: Context' -> [Constraint] -> HMError Context'
 
     -- nothing left! return accumulated context
     go g [] = Right g
@@ -90,7 +136,10 @@ unify = go mempty where
                 | x == y     = True
             occurs _         = False
 
-subst :: String -> Type -> Type -> Type
+-- | The expression @subst x t e@ substitutes all occurences of @x@ in @e@ with
+-- @t@
+
+subst :: Name -> Type -> Type -> Type
 subst x t (TyVar y) | x == y  = t
 subst x t (a :-> b)           = subst x t a :-> subst x t b
 subst _ _ e                   = e

src/Core/Syntax.hs

@@ -44,7 +44,7 @@ import Lens.Micro
 ----------------------------------------------------------------------------------
 
 data Expr b = Var Name
-            | Con Tag Int -- Con Tag Arity
+            | Con Tag Int -- ^ Con Tag Arity
             | Case (Expr b) [Alter b]
             | Lam [b] (Expr b)
             | Let Rec [Binding b] (Expr b)

tst/Core/HindleyMilnerSpec.hs

@@ -15,5 +15,5 @@ spec :: Spec
 spec = do
     it "should infer `id 3` :: Int" $
         let g = [ ("id", TyVar "a" :-> TyVar "a") ]
-        in infer g [coreExpr|id 3|] `shouldBe` Just TyInt
+        in infer g [coreExpr|id 3|] `shouldBe` Right TyInt