Bidirectional type elaboration for the simply-typed lambda calculus with unit values & types.

Bidi.hs

{-# LANGUAGE DeriveFunctor #-}
module Bidi where

-- For 'guard'.
import Control.Monad
-- We use Cofree to represent type-annotated terms.
import Control.Comonad.Cofree
import Data.Functor.Classes
-- We use Fix to represent unannotated terms.
import Data.Functor.Foldable

-- | The universe of simple types is defined by a simple tree structure.
data Type
  = Function Type Type -- ^ Function types, i.e. types from a -> b where a and b are both 'Type's.
  | UnitT -- ^ The unit type, i.e. the type of ().
  deriving (Eq, Show)

-- | The universe of simply-typed is defined by a slightly more complicated tree structure.
-- 
-- This is a non-recursive functor so that we can use the same type for both our regular term structure and the elaborated term structure using the Term and ATerm type synonyms (below).
data TermF a
  = App a a -- ^ Function application.
  | Ann a Type -- ^ Annotation of a term by its Type.
  | Lambda Int a -- ^ A De Bruijn–indexed lambda. The variable named by the Int is in scope in the body of the lambda (the second field, of type 'a').
  | Var Int -- ^ A reference to a variable. This must be in scope (see docs for Lambda) or it’s ill-formed (and therefore also ill-typed).
  | Unit -- ^ The unit value.
  deriving (Eq, Functor, Show)

-- | The Term type synonym represents an unannotated simply-typed lambda term. This is the structure you’d write your programs in, probably using smart constructors to abstract away 'Fix'/'unfix' boilerplate.
type Term = Fix TermF

-- | The ATerm type synonym represents an annotated simply-typed lambda term, for example a term elaborated with the types of every subterm.
type ATerm = Cofree TermF

-- | To check a term against a given type, we elaborate it with 'Just' that type.
check :: Type -> Term -> Context -> Maybe (ATerm Type)
check = elaborate . Just

-- | To infer the type of a term, we elaborate it without a type to check against.
infer :: Term -> Context -> Maybe (ATerm Type)
infer = elaborate Nothing

-- | Typing contexts here are simplistic. We do O(n) lookups in a linked list of Int, Type pairs! Do something smarter in real type systems; this is sort of the minimum viable typing context.
type Context = [(Int, Type)]

-- | Elaborate a term against a type in a context. Failures are represented as Nothing, so there are no error messages! A smarter typechecker would use Either to give the poor user some hint of what went wrong.
--
-- The type is given in Maybe so that we perform both inference and type checking in one place.
elaborate :: Maybe Type -> Term -> Context -> Maybe (ATerm Type)
elaborate ty term context = case (ty, unfix term) of
  -- For an application to be well-typed…
  (Nothing, App f a) -> do
    a'@(aT :< _) <- infer a context
    -- …its operator must be some function a -> b…
    f'@(Function from to :< _) <- infer f context
    -- …and its argument must be of type a.
    guard (from == aT)
    pure (to :< App f' a')
  -- Inference of annotations reduces to typechecking.
  (Nothing, Ann a t) -> check t a context
  -- For a lambda to be well-typed against some function type a -> b…
  (Just (Function from to), Lambda i body) -> do
    -- …its body must be well-typed in the context extended with the type a for the variable bound by the lambda.
    elaborated@(to' :< _) <- check to body ((i, from) : context)
    pure (Function from to' :< Lambda i elaborated)
  -- Unit values can be inferred to have unit type.
  (Nothing, Unit) -> Just (UnitT :< Unit)
  -- Variables just get looked up in the typing context. If a variable is missing (i.e. a term is ill-formed), we’ll just return Nothing, i.e. ill-formed implies ill-typed.
  (Nothing, Var i) -> (:< Var i) <$> i `lookup` context
  -- We can’t infer the type of lambdas. We could if we did something fancier, like unification, but we don’t.
  (Nothing, Lambda _ _) -> Nothing
  -- Finally, all other checks reduce to inference & equality of the inferred type against the expected one.
  (Just tExpected, _) -> do
    elaborated@(tActual :< _) <- infer term context
    guard (tActual == tExpected)
    pure elaborated

-- | Our old friend the identity function!
--
-- But wait: we don’t know how to infer the type of lambdas. And so while this term is well-formed, it’s (on its own) ill-typed: we can’t assign it a type.
-- 
-- Inference fails:
-- 
--   λ: infer idTerm []
--   Nothing
-- 
-- But since we’re using a bidirectional elaborator we can always just check it against some expected type. This is like annotating a top-level function with its type in some Haskell code (like this file!):
-- 
--   λ: check (Function UnitT UnitT) idTerm []
--   Just (Function UnitT UnitT :< Lambda 0 (UnitT :< Var 0))
-- 
-- Another thing we can do is to annotate the function with its type inline. That gives us enough information to assign it a type:
-- 
--   λ: infer (Fix (Ann idTerm (Function UnitT UnitT))) []
--   Just (Function UnitT UnitT :< Lambda 0 (UnitT :< Var 0))
idTerm :: Term
idTerm = Fix (Lambda 0 (Fix (Var 0)))


-- Implementation details follow

-- | We need this to show 'Fix TermF' values. 
instance Show1 TermF where
  liftShowsPrec sp _ d term = case term of
    App a b -> showsBinaryWith sp sp "App" d a b
    Ann a t -> showsBinaryWith sp showsPrec "Ann" d a t
    Lambda i a -> showsBinaryWith showsPrec sp "Lambda" d i a
    Var i -> showsUnaryWith showsPrec "Var" d i
    Unit -> showString "Unit"