Simply Typed Lambda Calculus with GADTs and Singletons (doodle made while reading code that uses singletons).

Single.hs

{-# LANGUAGE DataKinds #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TypeOperators #-}

-- https://stackoverflow.com/questions/27831223/simply-typed-lambda-calculus-with-failure-in-haskell
-- https://www.youtube.com/watch?v=6snteFntvjM
-- https://github.com/goldfirere/glambda/blob/master/src/Language/Glambda/Type.hs
-- https://github.com/goldfirere/glambda/blob/master/src/Language/Glambda/Check.hs

module Single where

import Data.Type.Equality

data Elem as a where
  EZ :: Elem (x:xs) x
  ES :: Elem xs x -> Elem (y:xs) x

deriving instance Show (Elem as a)

data Expr as a where
  Ref  :: Elem as a -> Expr as a
  Lam  :: Expr (a:as) b -> Expr as (a -> b)
  App  :: Expr as (a -> b) -> Expr as a -> Expr as b
  Num  :: Integer -> Expr as Integer
  Bool :: Bool -> Expr as Bool

deriving instance Show (Expr as a)

data UExpr
  = URef Int
  | ULam Type UExpr
  | UApp UExpr UExpr
  | UNum Integer
  | UBool Bool
  deriving Show

data Type
  = LamT Type Type
  | NumT
  | BoolT
  deriving (Eq, Show)

data SType a where
  SLamT  :: SType a -> SType b -> SType (a -> b)
  SNumT  :: SType Integer
  SBoolT :: SType Bool

deriving instance Show (SType a)

refine :: Type -> (forall a. SType a -> b) -> b
refine (LamT x y) f =
  refine x $ \x' ->
    refine y $ \y' -> f (SLamT x' y')
refine NumT f       = f SNumT
refine BoolT f      = f SBoolT

instance TestEquality SType where
  testEquality (SLamT x y) (SLamT x' y') = do
    Refl <- testEquality x x'
    Refl <- testEquality y y'
    pure Refl
  testEquality SNumT SNumT   = Just Refl
  testEquality SBoolT SBoolT = Just Refl
  testEquality _ _           = Nothing

data TEnv as where
  TZ :: TEnv '[]
  TS :: SType a -> TEnv as -> TEnv (a:as)

deriving instance Show (TEnv as)

data SomeExpr as where
  SomeExpr :: SType a -> Expr as a -> SomeExpr as

deriving instance Show (SomeExpr as)

data SomeRef as where
  SomeRef :: SType a -> Elem as a -> SomeRef as

deriving instance Show (SomeRef as)

check :: UExpr -> SType a -> Maybe (Expr '[] a)
check expr te = do
    SomeExpr te' expr' <- check' TZ expr
    Refl <- testEquality te te'
    pure expr'
  where
    check' :: TEnv as -> UExpr -> Maybe (SomeExpr as)
    check' env (URef x) = do
      SomeRef t' x' <- lookup' x env
      pure $ SomeExpr t' (Ref x')
    check' env (ULam t x) =
      refine t $ \t' -> do
        SomeExpr t'' x' <- check' (TS t' env) x
        pure $ SomeExpr (SLamT t' t'') (Lam x')
    check' env (UApp x y) = do
      SomeExpr tx x' <- check' env x
      SomeExpr ty y' <- check' env y
      case tx of
        (SLamT ty' t) -> do
          Refl <- testEquality ty ty'
          pure $ SomeExpr t (App x' y')
        _ -> Nothing
    check' _ (UNum x)  = pure $ SomeExpr SNumT (Num x)
    check' _ (UBool x) = pure $ SomeExpr SBoolT (Bool x)

    lookup' :: Int -> TEnv as -> Maybe (SomeRef as)
    lookup' 0 (TS t _)  = pure $ SomeRef t EZ
    lookup' x (TS _ ts) = do
      SomeRef x' ts' <- lookup' (x - 1) ts
      pure $ SomeRef x' (ES ts')
    lookup' _ _         = Nothing