Bidirectional typechecking using GADTs and DataKinds

Bidirectional.hs

{-# LANGUAGE DataKinds, KindSignatures, GADTs, StandaloneDeriving, TypeFamilies #-}
module Bidirectional where

import Control.Applicative
import Control.Monad (when)

----------------------------------------
-- 'Flow d a' is like 'Maybe a' except the type index 'd' says whether
-- we're given a value or whether we want one.

data Direction = Wanted | Given

type family Reverse (d :: Direction) :: Direction
type instance Reverse Wanted = Given
type instance Reverse Given = Wanted

data Flow (d :: Direction) a where
  Want :: Flow Wanted a
  Give :: a -> Flow Given a

instance Functor (Flow d) where
  fmap f Want = Want
  fmap f (Give x) = Give (f x)

-- instance Contravariant (Flow Want) where
--  contramap _ Want = Want

deriving instance Show a => Show (Flow d a)

type Input c d = Flow d c

type Output c d = Flow (Reverse d) c


----------------------------------------
-- DeBruijn encoding of STLC with one base type

data Type =
  BaseT | ArrT Type Type
  deriving (Eq, Show)

data Expr =
  V Int | Lam Expr | App Expr Expr
  deriving (Show)

----------------------------------------
-- type checking monad: ReaderT [Type] (Error String)

newtype TC a = TC { runTC :: [Type] -> Either String a }

instance Functor TC where
  fmap f (TC c) = TC $ \env -> fmap f (c env)

instance Applicative TC where
  pure = return
  mf <*> mx = do
    f <- mf
    x <- mx
    return (f x)

instance Monad TC where
  return x = TC $ const (Right x)
  (TC c) >>= f = TC $ \env -> do
    x <- c env
    runTC (f x) env

typeError :: String -> TC a
typeError = TC . const . Left

extendEnv :: Type -> TC a -> TC a
extendEnv t kont = TC $ \env -> runTC kont (t:env)
  
askEnv :: TC [Type]
askEnv = TC return

----------------------------------------
-- Bidirectional typechecking algorithm for STLC

typeCheck :: Expr -> Input Type d -> TC (Output Type d)
typeCheck e i =
  case e of
   V v -> tcVar v i
   Lam body ->
     case i of
      Give (ArrT t1 t2) -> do
        extendEnv t1 $ typeCheck body (Give t2)
      _ -> typeError "lambda not checking or not arrow type"
   App e1 e2 -> do
     (Give tf) <- typeCheck e1 Want
     tres <- case tf of
      ArrT targ tres -> do
        Want <- typeCheck e2 (Give targ)
        return tres
      _ -> typeError "application of non-function"
     putatively tres i

putatively :: Eq t => t -> Input t d -> TC (Output t d)
putatively tactual Want = return (Give tactual)
putatively tactual (Give texpected) = do
  when (texpected /= tactual) $
    typeError "expected and actual types differ"
  return $ Want

tcVar :: Int -> Input Type d -> TC (Output Type d)
tcVar n i = do
  t <- fmap (!! n) askEnv
  putatively t i

----------------------------------------
-- Example
--
-- >>> example
-- Right Want
example =
  let
    e = Lam $ Lam $ App (V 1) (Lam $ V 1)
    ty = ArrT (ArrT (ArrT BaseT BaseT) BaseT)
              (ArrT BaseT BaseT)
    env = []
  in runTC (typeCheck e (Give ty)) env

But this chokes on applications of lambda expressions, e.g.

runTC (typeCheck e (Give ty)) env
  where
    env = []
    e = App (Lam $ V 0) (Lam $ V 0)
    ty = ArrT BaseT BaseT