simple polymorphic type checker in haskell

Tyck.hs

{-# LANGUAGE LambdaCase #-}

module Tyck
  ( Term (..),
    Type (..),
    emptyContext,
    emptySolution,
    topInfer,
    infer,
    check,
  )
where

-- typechecking monad
type M = Either String

-- report error in monad
report :: String -> M a
report = Left

type Name = String

type Context = [(Name, Type)]

emptyContext :: Context
emptyContext = []

-- solution is a map from free type variables to its solution: here is a type
type Solution = [(Int, Type)]

emptySolution :: Solution
emptySolution = []

data Type
  = TFree Int
  | Type :-> Type
  | TInt
  | TBool
  deriving (Eq, Show)

data Term
  = EVar Name
  | -- lam x. M
    ELam Name Term
  | -- M @ N
    Term :@ Term
  | EInt Int
  | -- M + N
    -- where M, N : Int
    Term :+ Term
  | EBool Bool
  deriving (Eq, Show)

topInfer :: Term -> M Type
topInfer tm = do
  r <- infer emptyContext tm
  solve r
  where
    solve :: (Type, Solution) -> M Type
    solve (TFree i, sol) = case lookup i sol of
      Nothing -> report $ "cannot solve " ++ show (TFree i)
      Just ty -> return ty
    solve (t1 :-> t2, sol) = do
      t1' <- solve (t1, sol)
      t2' <- solve (t2, sol)
      return $ t1' :-> t2'
    solve (ty, _) = return ty

infer :: Context -> Term -> M (Type, Solution)
infer ctx = \case
  EVar x -> case lookup x ctx of
    Nothing -> report $ "Name not found: " ++ x
    Just ty -> return (ty, emptySolution)
  ELam x e -> do
    let xty = TFree (length ctx)
    (ety, sol) <- infer ((x, xty) : ctx) e
    return (xty :-> ety, sol)
  a :@ b -> do
    (aty, _) <- infer ctx a
    case aty of
      l :-> r -> do
        sol <- check ctx l b
        return (r, sol)
      _ -> report "non-appliable"
  EInt _ -> return (TInt, emptySolution)
  a :+ b -> do
    sol1 <- check ctx TInt a
    sol2 <- check ctx TInt b
    return (TInt, sol1 ++ sol2)
  EBool _ -> return (TBool, emptySolution)

check :: Context -> Type -> Term -> M Solution
check ctx ty tm = do
  (ty', sol) <- infer ctx tm
  unify sol ty ty'

unify :: Solution -> Type -> Type -> M Solution
unify solution t1 t2 = case (t1, t2) of
  (TFree i, ty) -> do
    occurs i ty
    return $ (i, ty) : solution
  (_, TFree _) -> unify solution t2 t1
  (a1 :-> a2, b1 :-> b2) -> do
    sol1 <- unify solution a1 b1
    unify sol1 a2 b2
  _ ->
    if t1 == t2
      then return solution
      else report $ "fail to unify: " ++ show t1 ++ " and " ++ show t2

occurs :: Int -> Type -> M ()
occurs idx ty = case ty of
  TFree _ -> report $ show ty ++ " occurs in " ++ show (TFree idx)
  t1 :-> t2 -> occurs idx t1 *> occurs idx t2
  _ -> return ()