There are De Bruijn levels of survival that we'd be willing to accept

TypedNBEWithLevels.hs

{-# Language CPP #-}
{-# Language BlockArguments #-}
{-# Language GADTs #-}
{-# Language RankNTypes #-}
{-# Language MagicHash #-}
{-# Language ViewPatterns #-}
{-# Language TypeApplications #-}
{-# Language BangPatterns #-}
{-# Language TypeOperators #-}
{-# Language TypeFamilyDependencies #-}
{-# Language DataKinds #-}
{-# Language ConstraintKinds #-}
{-# Language FlexibleInstances #-}
{-# Language AllowAmbiguousTypes #-}
{-# Language PolyKinds #-}
{-# Language TypeFamilies #-}
{-# Language FlexibleContexts #-}
{-# Language UndecidableInstances #-}
{-# Language PatternSynonyms #-}
{-# Language StandaloneKindSignatures #-}
{-# Language ScopedTypeVariables #-}
{-# Language DerivingStrategies #-}
{-# Language RoleAnnotations #-}
{-# Language MultiParamTypeClasses #-}
{-# options_ghc -Wall #-} 

module TypedNBEWithLevels where

import Control.Category (Category(..))
import Data.Bits (unsafeShiftR)
import Data.Coerce
import GHC.Types (Type)
import Prelude hiding (id, (.))
import Unsafe.Coerce

type (~>) :: (i -> Type) -> (i -> Type) -> Type
type f ~> g = forall x. f x -> g x
infixr 0 ~>

-------------------------------------------------------------------------------
-- * Typed De Bruijn indices
-------------------------------------------------------------------------------

type Ix' :: [i] -> i -> Type
type role Ix' nominal nominal
data Ix' j a where
  NZ' :: Ix' (a ': j) a
  NS' :: Ix j a -> Ix' (b ': j) a

type Ix :: [i] -> i -> Type
type role Ix nominal nominal
newtype Ix j a = UnsafeIx { indexIx :: Int }

instance Show (Ix j a) where
  showsPrec d = showsPrec d . indexIx

upIx :: Ix i ~> Ix' i
upIx (UnsafeIx 0) = unsafeCoerce NZ'
upIx (UnsafeIx n) = unsafeCoerce (NS' $ UnsafeIx (n-1))

pattern NZ :: () => (b ~ a) => Ix (b ': as) a
pattern NZ <- (upIx -> NZ') where
  NZ = UnsafeIx 0

pattern NS :: () => (a ~ a', as ~ (b ': as')) => Ix as' a' -> Ix as a
pattern NS n <- (upIx -> NS' n) where
  NS n = UnsafeIx (indexIx n + 1)

{-# complete NZ, NS #-}

-------------------------------------------------------------------------------
-- * Levels
-------------------------------------------------------------------------------

type Len :: [i] -> Type
type role Len nominal
newtype Len i = UnsafeLen Int

lz :: Len '[]
lz = UnsafeLen 0

ls :: Len i -> Len (a ': i)
ls (UnsafeLen i) = UnsafeLen (i + 1)

type Lvl :: [i] -> i -> Type
type role Lvl nominal nominal
newtype Lvl i a = UnsafeLvl Int

lvlIx :: Len i -> Lvl i a -> Ix i a
lvlIx (UnsafeLen i) (UnsafeLvl v) = UnsafeIx (i - v - 1)

ixLvl :: Len i -> Ix i a -> Lvl i a
ixLvl (UnsafeLen i) (UnsafeIx v) = UnsafeLvl (i - v - 1)

lvl0 :: Len i -> Lvl (a ': i) a
lvl0 = coerce

-------------------------------------------------------------------------------
-- * Expressions
-------------------------------------------------------------------------------

type Name = String

type Expr :: [Type] -> Type -> Type
type role Expr nominal nominal
data Expr e a where
  Var :: {-# unpack #-} !(Ix i a) -> Expr i a
  App :: !(Expr i (a -> b)) -> !(Expr i a) -> Expr i b
  Lam :: !Name -> !(Expr (a ': i) b) -> Expr i (a -> b)

instance Show (Expr i a) where
  showsPrec d (Var n) = showParen (d > 10) $
    showString "Var " . showsPrec 11 (indexIx n)
  showsPrec d (App f x) = showParen (d > 10) $
    showString "App " . showsPrec 11 f . showChar ' ' . showsPrec 11 x
  showsPrec d (Lam n b) = showParen (d > 10) $
    showString "Lam " . showsPrec 11 n . showChar ' ' . showsPrec 11 b

-------------------------------------------------------------------------------
-- * List suffixes
-------------------------------------------------------------------------------

-- | @Wk i j@ is a claim that i is a suffix of j
type Wk :: [i] -> [i] -> Type
type role Wk nominal nominal
data Wk i j = UnsafeWk

wk :: Wk i (a ': i)
wk = UnsafeWk

instance Category Wk where
  id = UnsafeWk
  UnsafeWk . UnsafeWk = UnsafeWk

-------------------------------------------------------------------------------
-- * Semantic Domain
-------------------------------------------------------------------------------

type Val :: [Type] -> Type -> Type
type role Val nominal nominal
data Val i a where
  VLam :: !Name -> (forall j. Wk i j -> Val j a -> Val j b) -> Val i (a -> b)
  VVar :: {-# unpack #-} !(Lvl i a) -> !(Spine i a b) -> Val i b

type Spine :: [Type] -> Type -> Type -> Type
type role Spine nominal nominal nominal
data Spine i a b where
  SNil :: Spine i a a
  SApp :: !(Spine i x (a -> b)) -> !(Val i a) -> Spine i x b

infixl 9 `SApp`

{-# complete SNil, SApp #-}

-------------------------------------------------------------------------------
-- * Environments
-------------------------------------------------------------------------------

type Tree :: [Type] -> [Type] -> [Type] -> Type
type role Tree nominal nominal nominal
data Tree i j k where
  TTip :: !(Val i a) -> Tree i (a ': j) j
  TBin :: !(Val i a) -> !(Tree i j k) -> !(Tree i k l) -> Tree i (a ': j) l

type Env :: [Type] -> [Type] -> Type
type role Env nominal nominal
data Env j i where
  ESkCons :: {-# unpack #-} !Int -> !(Tree i j k) -> !(Env k i) -> Env j i
  ENil  :: Env '[] i

wkEnv :: Wk j k -> Env i j -> Env i k
wkEnv _ = unsafeCoerce

type Env' :: [Type] -> [Type] -> Type
type role Env' nominal nominal
data Env' i j where
  ENil' :: Env' '[] i
  ECons' :: Val i a -> Env k i -> Env' (a ': k) i

econs :: Val i a -> Env j i -> Env (a ': j) i
econs a (ESkCons n l (ESkCons m r xs))
  | n == m = ESkCons (2*n+1) (TBin a l r) xs
econs a xs = ESkCons 1 (TTip a) xs

upEnv :: Env i j -> Env' i j
upEnv ENil = ENil'
upEnv (ESkCons _ (TTip a) xs)     = ECons' a xs
upEnv (ESkCons n (TBin a l r) xs) = ECons' a $ ESkCons n' l $ ESkCons n' r xs
  where n' = unsafeShiftR n 1

pattern ECons :: () => (l ~ (a ': k)) => Val i a -> Env k i -> Env l i
pattern ECons v e <- (upEnv -> ECons' v e) where
  ECons v e = econs v e

infixr 5 `ECons`

{-# complete ENil, ECons #-}

evalVar :: Env i j -> Ix i ~> Val j
evalVar = go where
  go :: Env i j -> Ix i ~> Val j
  go (ESkCons n t xs) !m
    | n <= indexIx m = go xs (UnsafeIx (indexIx m-n))
    | otherwise = goTree n t m
  go ENil _ = error "impossible name"

  goTree :: Int -> Tree i j k -> Ix j ~> Val i
  goTree _ (TTip a) NZ = a
  goTree _ (TTip _) _ = error "impossible name"
  goTree _ (TBin a _ _) NZ = a
  goTree n (TBin _ l r) m
    | indexIx m <= n' = goTree n' l $ UnsafeIx (indexIx m - 1)
    | otherwise = goTree n' r $ UnsafeIx (indexIx m - n' - 1)
    where n' = unsafeShiftR n 1

-------------------------------------------------------------------------------
-- * Eval
-------------------------------------------------------------------------------

eval :: Env i j -> Expr i ~> Val j
eval e (Var n) = evalVar e n
eval e (App f x) = evalApp (eval e f) (eval e x)
eval e (Lam n b) = VLam n \w v -> eval (v `ECons` wkEnv w e) b

evalApp :: Val i (a -> b) -> Val i a -> Val i b
evalApp (VLam _ f) v = f id v
evalApp (VVar h s) v = VVar h (SApp s v)

quote :: Len i -> Val i a -> Expr i a
quote d (VLam n b) = Lam n $ quote (ls d) $ b wk $ VVar (lvl0 d) SNil
quote d (VVar h s) = quoteSp d s (Var (lvlIx d h))

quoteSp :: Len i -> Spine i a b -> Expr i a -> Expr i b
quoteSp _ SNil e = e
quoteSp d (SApp s x) e = App (quoteSp d s e) (quote d x)

nf :: Len j -> Env i j -> Expr i ~> Expr j
nf d e t = quote d (eval e t)

-------------------------------------------------------------------------------
-- * Tests
-------------------------------------------------------------------------------

id_ :: Val i (a -> a)
id_ = eval ENil $ Lam "x" (Var NZ)

k :: Val i (a -> b -> a)
k = eval ENil $ Lam "x" (Lam "y" (Var (NS NZ)))

kid_ :: Val i (a -> b -> b)
kid_ = eval (id_ `ECons` k `ECons` ENil) (App (Var (NS NZ)) (Var NZ))

main :: IO ()
main = do
  print $ quote lz kid_

UntypedNBEWithLevels.hs

{-# Language CPP #-}
{-# Language AllowAmbiguousTypes #-}
{-# Language BangPatterns #-}
{-# Language BlockArguments #-}
{-# Language ConstraintKinds #-}
{-# Language DataKinds #-}
{-# Language DerivingStrategies #-}
{-# Language FlexibleContexts #-}
{-# Language FlexibleInstances #-}
{-# Language GADTs #-}
{-# Language GeneralizedNewtypeDeriving #-}
{-# Language MagicHash #-}
{-# Language PatternSynonyms #-}
{-# Language PolyKinds #-}
{-# Language RankNTypes #-}
{-# Language RoleAnnotations #-}
{-# Language ScopedTypeVariables #-}
{-# Language StandaloneKindSignatures #-}
{-# Language TypeApplications #-}
{-# Language TypeFamilies #-}
{-# Language TypeFamilyDependencies #-}
{-# Language TypeOperators #-}
{-# Language UndecidableInstances #-}
{-# Language ViewPatterns #-}
{-# options_ghc -Wall -fobject-code #-}

module SizedNBE where

import Control.Category (Category(..))
import GHC.Types (Type)
import Prelude hiding (id, (.))
import GHC.TypeLits
import GHC.Exts (Proxy#, proxy#, magicDict, Constraint)
import Data.Bits (unsafeShiftR)
import Unsafe.Coerce

class Index t where
  index :: t -> Int

type Z = 0
type S n = 1 + n

type Dict :: Constraint -> Type
data Dict p where
  Dict :: p => Dict p

axiom :: forall a b. Dict (a ~ b)
axiom = unsafeCoerce (Dict :: Dict (a ~ a))

axiomLe :: forall a b. Dict (a <= b)
axiomLe = axiom

wk :: forall i. Dict (i <= S i)
wk = axiomLe @i @(S i)

type Name = String

type (~>) :: (i -> Type) -> (i -> Type) -> Type
type f ~> g = forall x. f x -> g x
infixr 0 ~>

-- singleton type (for small nats)
type SmNat :: Nat -> Type
type role SmNat nominal
newtype SmNat n = UnsafeSmNat Int
instance Index (SmNat n) where index (UnsafeSmNat n) = n

unsafeCoerce1 :: f a -> f b
unsafeCoerce1 = unsafeCoerce
{-# inline unsafeCoerce1 #-}

type SmNat' :: Nat -> Type
type role SmNat' nominal
data SmNat' n where
  SZ' :: SmNat' Z
  SS' :: SmNat n -> SmNat' (S n)

upNat :: SmNat n -> SmNat' n
upNat (UnsafeSmNat 0) = unsafeCoerce1 SZ'
upNat (UnsafeSmNat n) = unsafeCoerce1 (SS' $ UnsafeSmNat (n-1))

pattern SZ :: () => (n ~ Z) => SmNat n
pattern SZ <- (upNat -> SZ') where
  SZ = UnsafeSmNat 0

pattern SS :: () => (n ~ S n') => SmNat n' -> SmNat n
pattern SS n <- (upNat -> SS' n) where
  SS n = UnsafeSmNat (index n + 1)

type N :: Nat -> Constraint
class N n where
  knownN :: SmNat n

instance KnownNat n => N n where
  knownN = UnsafeSmNat $ fromIntegral $ natVal' (proxy# @n)

data Wrap n r = Wrap (N n => Proxy# n -> r)

withN# :: (N n => Proxy# n -> r) -> SmNat n -> Proxy# n -> r
withN# f x y = magicDict (Wrap f) x y

withN :: SmNat n -> (N n => r) -> r
withN n r = withN# (\_ -> r) n proxy#

type (:-) :: Constraint -> Constraint -> Type
newtype p :- q = Sub (p => Dict q)

succN :: forall n. N n :- N (S n)
succN = Sub $ withN (SS (knownN @n)) Dict

-------------------------------------------------------------------------------
-- * Fin n
-------------------------------------------------------------------------------

type Fin' :: Nat -> Type
type role Fin' nominal
data Fin' j where
  FZ' :: Fin' (S j)
  FS' :: Fin j -> Fin' (S j)

type Fin :: Nat -> Type
type role Fin nominal
newtype Fin j = UnsafeFin Int
instance Index (Fin j) where index (UnsafeFin n) = n

instance Show (Fin j) where
  showsPrec d = showsPrec d . index

upFin :: Fin ~> Fin'
upFin (UnsafeFin 0) = unsafeCoerce FZ'
upFin (UnsafeFin n) = unsafeCoerce (FS' $ UnsafeFin (n-1))

pattern FZ :: Fin (S j)
pattern FZ <- (upFin -> FZ') where
  FZ = UnsafeFin 0

pattern FS :: () => (j ~ S i) => Fin i -> Fin j
pattern FS n <- (upFin -> FS' n) where
  FS n = UnsafeFin (index n + 1)

{-# complete FZ, FS #-}

-------------------------------------------------------------------------------
-- * Expressions
-------------------------------------------------------------------------------

newtype Ix i = Ix { runIx :: Fin i } deriving newtype (Show,Index)
pattern UnsafeIx :: Int -> Ix i
pattern UnsafeIx i = Ix (UnsafeFin i)
{-# complete UnsafeIx #-}

newtype Lvl i = Lvl { runLvl :: Fin i } deriving newtype (Show,Index)
pattern UnsafeLvl :: Int -> Lvl i
pattern UnsafeLvl i = Lvl (UnsafeFin i)
{-# complete UnsafeLvl #-}

type Expr :: Nat ->  Type
type role Expr nominal
data Expr i where
  Var :: {-# unpack #-} !(Ix i) -> Expr i -- debruijn indexes
  App :: !(Expr i) -> !(Expr i) -> Expr i
  Lam :: !Name -> !(Expr (S i)) -> Expr i -- nice names for printing

instance Show (Expr i) where
  showsPrec d (Var n)   = showParen (d > 10) $ 
    showString "Var " . showsPrec 11 n
  showsPrec d (App f x) = showParen (d > 10) $ 
    showString "App " . showsPrec 11 f . showChar ' ' . showsPrec 11 x
  showsPrec d (Lam n b) = showParen (d > 10) $ 
    showString "Lam " . showsPrec 11 n . showChar ' ' . showsPrec 11 b

-------------------------------------------------------------------------------
-- * Semantic Domain
-------------------------------------------------------------------------------

type Val :: Nat -> Type
type role Val nominal
data Val i where
  VLam :: !Name -> (forall j. i <= j => Val j -> Val j) -> Val i
  VVar :: !(Lvl i) -> !(Spine i) -> Val i

type Spine :: Nat -> Type
type role Spine nominal
data Spine i where
  SNil :: Spine i
  SApp :: !(Spine i) -> !(Val i) -> Spine i

infixl 9 `SApp`

{-# complete SNil, SApp #-}

-------------------------------------------------------------------------------
-- * Environments
-------------------------------------------------------------------------------

type Tree :: Nat -> Nat -> Nat -> Type
type role Tree nominal nominal nominal
data Tree i j k where
  TTip :: !(Val i) -> Tree i (S j) j
  TBin :: !(Val i) -> !(Tree i j k) -> !(Tree i k l) -> Tree i (S j) l

type Env :: Nat -> Nat -> Type
type role Env nominal nominal
data Env i j where
  ESkCons :: {-# unpack #-} !Int -> !(Tree i j k) -> !(Env k i) -> Env j i
  ENil  :: Env Z i

type Env' :: Nat -> Nat -> Type
type role Env' nominal nominal
data Env' i j where
  ENil' :: Env' Z i
  ECons' :: Val i -> Env k i -> Env' (S k) i

econs :: Val i -> Env j i -> Env (S j) i
econs a (ESkCons n l (ESkCons m r xs))
  | n == m = ESkCons (2*n+1) (TBin a l r) xs
econs a xs = ESkCons 1 (TTip a) xs

upEnv :: Env i j -> Env' i j
upEnv ENil = ENil'
upEnv (ESkCons _ (TTip a) xs)     = ECons' a xs
upEnv (ESkCons n (TBin a l r) xs) = ECons' a $ ESkCons n' l $ ESkCons n' r xs
  where n' = unsafeShiftR n 1

pattern ECons :: () => (l ~ S k) => Val i -> Env k i -> Env l i
pattern ECons v e <- (upEnv -> ECons' v e) where
  ECons v e = econs v e

infixr 5 `ECons`

wkEnv :: forall j k i. j <= k => Env i j -> Env i k
wkEnv = unsafeCoerce1

{-# complete ENil, ECons #-}

evalVar :: Env i j -> Ix i -> Val j
evalVar ENil _ = error "impossible name"
evalVar (ESkCons n0 t xs) !m0
  | n0 <= index m0 = evalVar xs (UnsafeIx (index m0-n0))
  | otherwise = go n0 t m0
  where
    go :: Int -> Tree i j k -> Ix j -> Val i
    go _ (TTip a) _ = a
    go n (TBin a l r) m
      | n == 0 = a
      | index m <= n' = go n' l $ UnsafeIx (index m - 1)
      | otherwise     = go n' r $ UnsafeIx (index m - n' - 1)
      where n' = unsafeShiftR n 1

-------------------------------------------------------------------------------
-- * Eval
-------------------------------------------------------------------------------

eval :: Env i j -> Expr i -> Val j
eval e (Var n)   = evalVar e n
eval e (App f x) = evalApp (eval e f) (eval e x)
eval e (Lam n b) = VLam n \ v -> eval (ECons v (wkEnv e)) b

evalApp :: Val i -> Val i -> Val i
evalApp (VLam _ f) v = f v
evalApp (VVar h s) v = VVar h (SApp s v)

lvlIx :: forall i. N i => Lvl i -> Ix i
lvlIx (UnsafeLvl v) = UnsafeIx (index (knownN @i) - v - 1)

quote :: forall i. N i => Val i -> Expr i
quote (VLam n b) = Lam n $ case succN @i of
  Sub Dict -> quote $ case wk @i of
    Dict -> b $ VVar (UnsafeLvl (index (knownN @i))) SNil
quote (VVar h s) = quoteSp s (Var (lvlIx h))

quoteSp :: N i => Spine i -> Expr i -> Expr i
quoteSp SNil e = e
quoteSp (SApp s x) e = App (quoteSp s e) (quote x)

nf :: forall j i. N j => Env i j -> Expr i -> Expr j
nf e t = quote (eval e t)

-------------------------------------------------------------------------------
-- * Tests
-------------------------------------------------------------------------------

ki :: Val i
ki = eval (i `ECons` k `ECons` ENil) (App (Var (Ix (FS FZ))) (Var (Ix FZ))) where
  i = eval ENil $ Lam "x" (Var (Ix FZ))
  k = eval ENil $ Lam "x" (Lam "y" (Var (Ix (FS FZ))))

main :: IO ()
main = do
  print (quote ki :: Expr Z)

This seems to have a bug in it.

Instead of b (VVar (Lvl FZ) SNil) in quote, you should pass i as level to b.

Yeah, Lvl FZ should be Lvl i basically.

Fixed.

I'll paste the typed version too.