Type-check STLC into a GADT

STLC.hs

{-# LANGUAGE BlockArguments, DataKinds, LambdaCase, MonoLocalBinds #-}
module STLC where

-- https://gist.github.com/pedrominicz/656cbe1a8309af8ef3226b40093bf409
-- https://github.com/goldfirere/glambda/blob/master/src/Language/Glambda/Check.hs
-- https://github.com/goldfirere/glambda/blob/master/src/Language/Glambda/Type.hs
-- https://stackoverflow.com/questions/11104536/how-to-parse-strings-to-syntax-tree-using-gadts

import Data.Type.Equality

data Type a where
  Int :: Type Int
  Fun :: Type a -> Type b -> Type (a -> b)

data Ctx ctx where
  Empty  :: Ctx '[]
  Extend :: Type a -> Ctx ctx -> Ctx (a:ctx)

data Var ctx a where
  Zero :: Var (a:ctx) a
  Succ :: Var ctx a -> Var (b:ctx) a

data Expr ctx a where
  Var :: Var ctx a -> Expr ctx a
  Lam :: Expr (a:ctx) b -> Expr ctx (a -> b)
  App :: Expr ctx (a -> b) -> Expr ctx a -> Expr ctx b
  Num :: Int -> Expr ctx Int

data Type' where
  Int' :: Type'
  Fun' :: Type' -> Type' -> Type'

data Expr' where
  Var' :: Int -> Expr'
  Lam' :: Type' -> Expr' -> Expr'
  App' :: Expr' -> Expr' -> Expr'
  Num' :: Int -> Expr'

type' :: Type' -> (forall a. Type a -> b) -> b
type' Int' k = k Int
type' (Fun' a b) k =
  type' a \a ->
    type' b \b ->
      k (Fun a b)

eq :: Type a -> Type b -> Maybe (a :~: b)
eq Int Int = Just Refl
eq (Fun a1 b1) (Fun a2 b2) =
  case (eq a1 a2, eq b1 b2) of
    (Just Refl, Just Refl) -> Just Refl
    _ -> Nothing
eq _ _ = Nothing

typeOf :: Ctx ctx -> Int -> (forall a. Var ctx a -> Type a -> Maybe b) -> Maybe b
typeOf (Extend x _) 0 k = k Zero x
typeOf (Extend _ ctx) x k = typeOf ctx (x - 1) \x a -> k (Succ x) a
typeOf Empty _ _ = Nothing

check :: Expr' -> (forall a. Expr '[] a -> b) -> Maybe b
check e k = go Empty e \e _ -> Just (k e)
  where
  go :: Ctx ctx -> Expr' -> (forall a. Expr ctx a -> Type a -> Maybe b) -> Maybe b
  go ctx (Var' x) k = typeOf ctx x \x a -> k (Var x) a
  go ctx (Lam' a f) k =
    type' a \a ->
      go (Extend a ctx) f \f b ->
        k (Lam f) (Fun a b)
  go ctx (App' f x) k =
    go ctx f \f -> \case
      Fun a b ->
        go ctx x \x a' ->
          case eq a a' of
            Just Refl -> k (App f x) b
            _ -> Nothing
      _ -> Nothing
  go _ (Num' x) k = k (Num x) Int