Strongly Typed Typechecker for Simply Typed Lambda Calculus in Haskell

TypedTypechecking.hs

{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}

import Data.List (intercalate)
import Data.Type.Equality ((:~:)(Refl))
import GHC.Exts (IsString(..))

main :: IO ()
main = maybe (putStrLn "error") print
  $ uncurry typecheck example

example :: (TEnv, UExp)
example =

  -- In the context:
  --   f : int -> int,
  --   y : int
  ( TEnv
    [ ("f", UTInt :-> UTInt)
    , ("y", UTInt)
    ]

  -- This term:
  --   (\x:int. f (x + y)) 1
  , (UELam "x" UTInt ("f" :@ ("x" + "y"))) :@ 1
  )

  -- Should have the type 'int':
  --
  -- f : (int -> int), y : int
  --   |- (\x : int. f (x + y)) 1
  --     :: int

type Name = String

newtype TEnv = TEnv [(Name, UType)]

instance Show TEnv where
  show (TEnv bs) = intercalate ", "
    $ fmap
      (\ (x, t) -> concat
        [ x
        , " : "
        , showsPrec (arrPrec + 1) t ""
        ])
      bs

lookupTEnv :: Name -> TEnv -> Maybe UType
lookupTEnv x (TEnv bs) = lookup x bs

insertTEnv :: Name -> UType -> TEnv -> TEnv
insertTEnv x t (TEnv bs) = TEnv ((x, t) : bs)

-- | Untyped expression
data UExp

  -- | Integer literal
  = UEInt Int

  -- | Lambda abstraction (with explicitly typed binding)
  | UELam Name UType UExp

  -- | Function application
  | UExp :@ UExp

  -- | Variable
  | UEVar Name

  -- | Arithmetic
  | UExp :+ UExp
  | UExp :* UExp
  | UExp :- UExp

infixl 6 :+
infixl 6 :-
infixl 7 :*
infixl 9 :@

-- | Convenience instance for variable names
instance IsString UExp where
  fromString = UEVar

-- | Convenience instance for arithmetic operations
instance Num UExp where
  fromInteger = UEInt . fromInteger
  (+) = (:+)
  (*) = (:*)
  (-) = (:-)
  abs = ("abs" :@)
  signum = ("signum" :@)
  negate = ("negate" :@)

instance Show UExp where
  showsPrec p = \ case
    UEInt n -> shows n
    UELam x t e -> showParen (p > lamPrec)
      $ showString "\\"
      . showString x
      . showString " : "
      . showsPrec (arrPrec + 1) t
      . showString ". "
      . showsPrec appPrec e
    e1 :@ e2 -> showParen (p > appPrec)
      $ showsPrec appPrec e1
      . showString " "
      . showsPrec (appPrec + 1) e2
    UEVar x -> showString x
    e1 :+ e2 -> showParen (p > addPrec)
      $ showsPrec addPrec e1
      . showString " + "
      . showsPrec (addPrec + 1) e2
    e1 :* e2 -> showParen (p > mulPrec)
      $ showsPrec mulPrec e1
      . showString " * "
      . showsPrec (mulPrec + 1) e2
    e1 :- e2 -> showParen (p > subPrec)
      $ showsPrec subPrec e1
      . showString " - "
      . showsPrec (subPrec + 1) e2

-- | Typed expression
data TExp (t :: UType) where
  TEInt :: Int -> TExp 'UTInt
  TELam :: Name -> TType a -> TExp b -> TExp (a ':-> b)
  (::@) :: TExp (a ':-> b) -> TExp a -> TExp b
  TEVar :: Name -> TExp a
  (::+) :: TExp 'UTInt -> TExp 'UTInt -> TExp 'UTInt
  (::*) :: TExp 'UTInt -> TExp 'UTInt -> TExp 'UTInt
  (::-) :: TExp 'UTInt -> TExp 'UTInt -> TExp 'UTInt

infixr 6 ::+
infixl 6 ::-
infixl 7 ::*
infixl 9 ::@

-- | Untyped type
data UType
  = UTInt
  | UType :-> UType

infixr 9 :->

instance Show UType where
  showsPrec p = \ case
    UTInt -> showString "int"
    a :-> b -> showParen (p > arrPrec)
      $ showsPrec (arrPrec + 1) a
      . showString " -> "
      . showsPrec arrPrec b

lamPrec, addPrec, subPrec, mulPrec, arrPrec, appPrec :: Int
lamPrec = 0
addPrec = 6
subPrec = 6
mulPrec = 7
arrPrec = 9
appPrec = 10

-- | Typed type (promoted/singleton version of ‘UType’)
data TType (t :: UType) where
  TTInt :: TType 'UTInt
  (::->) :: TType a -> TType b -> TType (a ':-> b)

infixr 9 ::->

-- | Convert a typed expression back to an untyped one
demoteExp :: TExp t -> UExp
demoteExp = \ case
  TEInt n -> UEInt n
  TELam x t e -> UELam x (demoteType t) (demoteExp e)
  e1 ::@ e2 -> demoteExp e1 :@ demoteExp e2
  TEVar x -> UEVar x
  e1 ::+ e2 -> demoteExp e1 :+ demoteExp e2
  e1 ::* e2 -> demoteExp e1 :* demoteExp e2
  e1 ::- e2 -> demoteExp e1 :- demoteExp e2

-- | Convert a typed type back to an untyped one
demoteType :: TType t -> UType
demoteType TTInt = UTInt
demoteType (a ::-> b) = demoteType a :-> demoteType b

-- | For pretty output in 'main'
data Judgement = TEnv :|- Annotated

infixl 0 :|-

instance Show Judgement where
  show (tenv :|- e) = concat
    [ show tenv
    , " |- "
    , show e
    ]

-- | A dependent pair of a typed expression, with a
-- representation of its type, allowing us to discover the
-- type of a typed term at runtime by pattern-matching.
data Annotated where
  (:::) :: forall t. TExp t -> TType t -> Annotated

infixl 1 :::

instance Show Annotated where
  show (e ::: t) = concat
    [ show (demoteExp e)
    , " :: "
    , show (demoteType t)
    ]

-- | Abstract/hidden type
data Type where
  Type :: forall t. TType t -> Type

-- | Given an environment of the types of free variables,
-- and an untyped expression, try to produce a judgement
-- that the expression is well typed in the environment,
-- annotated with a type to prove it
typecheck :: TEnv -> UExp -> Maybe Judgement
typecheck tenv expr = (tenv :|-) <$> check expr
  where
    check = \ case

      -- For simple terms we just annotate the type
      UEInt n -> pure (TEInt n ::: TTInt)

      -- For more complex terms we check subterms, then
      -- match to produce type-level equalities, then return
      -- an annotated term
      e1 :@ e2 -> do
        e1' ::: a ::-> b <- check e1
        e2' ::: a' <- check e2
        Refl <- equateTypes a a'
        pure (e1' ::@ e2' ::: b)

      -- We need the explicit type of the binder to check
      -- the body here; without it, we'd need to infer it
      UELam x t e -> case liftType t of
        Type a -> do
          let tenv' = insertTEnv x t tenv
          _tenv' :|- e' ::: b <- typecheck tenv' e
          Refl <- equateTypes a b
          pure (TELam x a e' ::: a ::-> b)

      -- We just look up variables in the environment, and
      -- lift their types to the type level
      UEVar x -> do
        a <- lookupTEnv x tenv
        case liftType a of
          Type a' -> pure (TEVar x ::: a')

      -- Arithmetic expressions are monomorphic and
      -- uncurried; that is, we have this typing rule:
      --
      -- Γ ⊢ e₁ : int    Γ ⊢ e₂ : int
      -- ────────────────────────────
      -- Γ ⊢ e₁ + e₂ : int
      --
      -- Not this one:
      --
      -- ───────────────────────────
      -- Γ ⊢ (+) : int -> int -> int
      --
      --
      -- You can construct the curried form from the
      -- primitive one, though:
      --
      -- UELam "x" UTInt (UELam "y" UTInt ("x" :+ "y"))
      --
      e1 :+ e2 -> do
        e1' ::: TTInt <- check e1
        e2' ::: TTInt <- check e2
        pure (e1' ::+ e2' ::: TTInt)

      e1 :* e2 -> do
        e1' ::: TTInt <- check e1
        e2' ::: TTInt <- check e2
        pure (e1' ::* e2' ::: TTInt)

      e1 :- e2 -> do
        e1' ::: TTInt <- check e1
        e2' ::: TTInt <- check e2
        pure (e1' ::- e2' ::: TTInt)

-- | Lift a type to the type level
liftType :: UType -> Type
liftType UTInt = Type TTInt
liftType (a :-> b) = case liftType a of
  Type a' -> case liftType b of
    Type b' -> Type (a' ::-> b')

-- | Check two types for equality, returning 'Just' a proof
-- that they're equal, or else 'Nothing'
equateTypes :: TType a -> TType b -> Maybe (a :~: b)
equateTypes TTInt TTInt = pure Refl
equateTypes (a ::-> b) (a' ::-> b') = do
  Refl <- equateTypes a a'
  Refl <- equateTypes b b'
  pure Refl
equateTypes _ _ = Nothing