Sized NBE using the same schema as the typed one, but with natural number sizes on environments

SizedNBE.hs

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

-- #define SAFE 1
-- #define NOSKEW 1
-- #define PEANO 1

module SizedNBE where

import Control.Category (Category(..))
import GHC.Types (Type)
import Prelude hiding (id, (.))

#ifndef PEANO
import GHC.TypeLits 
#endif

#ifndef NOSKEW
import Data.Bits (unsafeShiftR)
#endif

#ifndef SAFE
import Unsafe.Coerce
#endif

#ifdef PEANO
data Nat = Z | S !Nat
type Z = 'Z
type S = 'S
#else
type Z = 0
type S n = 1 + n
#endif

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

-------------------------------------------------------------------------------
-- * Names (unsafe version)
-------------------------------------------------------------------------------

#ifdef SAFE

type Name :: Nat -> Type
type role Name nominal
data Name j where
  NZ :: Name (S j)
  NS :: Name j -> Name (S j) a

indexName :: Name j -> Int
indexName = go 0 where
  go :: Int -> Name j -> Int
  go !acc NZ = acc
  go acc (NS n) = go (acc + 1) n

instance Show (Name j) where
  showsPrec d = showsPrec d . indexName

#else

type Name' :: Nat -> Type
type role Name' nominal
data Name' j where
  NZ' :: Name' (S j)
  NS' :: Name j -> Name' (S j)

type Name :: Nat -> Type
type role Name nominal
newtype Name j = UnsafeName { indexName :: Int }

instance Show (Name j) where
  showsPrec d = showsPrec d . indexName

upName :: Name ~> Name'
upName (UnsafeName 0) = unsafeCoerce NZ'
upName (UnsafeName n) = unsafeCoerce (NS' $ UnsafeName (n-1))

pattern NZ :: () => () => Name (S as)
pattern NZ <- (upName -> NZ') where
  NZ = UnsafeName 0

pattern NS :: () => (j ~ S i) => Name i -> Name j
pattern NS n <- (upName -> NS' n) where
  NS n = UnsafeName (indexName n + 1)

{-# complete NZ, NS #-}

#endif

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

type Expr :: Nat ->  Type
type role Expr nominal
data Expr i where
  Var :: {-# unpack #-} !(Name i) -> Expr i
  App :: !(Expr i) -> !(Expr i) -> Expr i
  Lam :: !(Expr (S i)) -> Expr i

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

-------------------------------------------------------------------------------
-- * Weakening the head of the environment
-------------------------------------------------------------------------------

#ifdef SAFE

type Wk :: Nat -> Nat -> Type
type role Wk nominal nominal
data Wk i j where
  WZ :: Wk i i
  WS :: Wk i j -> Wk i (S j)

instance Category Wk where
  id = WZ
  WZ . f = f
  f . WZ = f
  WS x . y = WS (x . y)

instance Eq (Wk i j) where
  _ == _ = True

wkName :: Wk i j -> Name i ~> Name j
wkName WZ n = n
wkName (WS n) o = NS (wkName n o)

wk :: Wk i (a ': i)
wk = WS WZ

indexWk :: Wk i j -> Int
indexWk = go 0 where
  go :: Int -> Wk i j -> Int
  go !acc WZ = acc
  go acc (WS n) = go (acc + 1) n

#else

type Wk :: Nat -> Nat -> Type
type role Wk nominal nominal
newtype Wk i j = UnsafeWk { indexWk :: Int }

instance Show (Wk i j) where
  showsPrec d = showsPrec d . indexWk

instance Eq (Wk i j) where
  _ == _ = True

instance Category Wk where
  id = UnsafeWk 0
  UnsafeWk m . UnsafeWk n = UnsafeWk (m + n)

wkName :: Wk i j -> Name i -> Name j
wkName (UnsafeWk i) (UnsafeName j) = UnsafeName (i + j)

type Wk' :: Nat -> Nat -> Type
type role Wk' nominal nominal
data Wk' i j where
  WZ' :: Wk' i i
  WS' :: Wk i j -> Wk' i (S j)

upWk :: Wk i j -> Wk' i j
upWk (UnsafeWk 0) = unsafeCoerce WZ'
upWk (UnsafeWk n) = unsafeCoerce (WS' $ UnsafeWk (n-1))

pattern WZ :: () => (i ~ j) => Wk i j
pattern WZ <- (upWk -> WZ') where
  WZ = UnsafeWk 0

pattern WS :: () => (k ~ S j) => Wk i j -> Wk i k
pattern WS n <- (upWk -> WS' n) where
  WS n = UnsafeWk (indexWk n + 1)

{-# complete WZ, WS #-}

wk :: Wk i (S i)
wk = UnsafeWk 1

#endif

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

type Val :: Nat -> Type
type role Val nominal
data Val i where
  VLam :: {-# unpack #-} !(Wk i j) -> (forall k. Wk i k -> Val k -> Val k) -> Val j
  VVar :: {-# unpack #-} !(Name i) -> !(Spine i) -> Val i

type Spine :: Nat -> Type
type role Spine nominal
data Spine i where
  SNil :: Spine i
  SAPP :: {-# unpack #-} !(Wk i j) -> !(Spine i) -> !(Val i) -> Spine j

pattern SApp :: Spine i -> Val i -> Spine i
pattern SApp s v <- SAPP η (wkSp η -> s) (wkVal η -> v) where
  SApp s v = SAPP id s v

infixl 9 `SApp`

{-# complete SNil, SApp #-}

wkVal :: Wk i j -> Val i -> Val j
wkVal δ (VLam η b) = VLam (δ . η) b
wkVal δ (VVar n s) = VVar (wkName δ n) (wkSp δ s)

wkSp :: Wk i j -> Spine i -> Spine j
wkSp _ SNil = SNil
wkSp δ (SAPP η s v) = SAPP (δ . η) s v

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

#ifdef NOSKEW

type role Env nominal nominal
type Env :: Nat -> Nat -> Type
data Env k i where
  ENil :: Env '[] i
  ECONS :: {-# unpack #-} !(Wk i j) -> !(Val i) -> !(Env k i) -> Env (S k) j

pattern ECons :: () => (l ~ S k) => Val i -> Env k i -> Env l i
pattern ECons v e <- ECONS η (wkVal η -> v) (wkEnv η -> e) where
  ECons v e = ECONS id v e

infixr 5 `ECons`

{-# complete ENil, ECons :: Env #-}

evalVar :: Env i j -> Name i -> Val j
evalVar = go id where
  go :: Wk j k -> Env i j -> Name i ~> Val k
  go δ (ECONS η a _) NZ = wkVal (δ . η) a
  go δ (ECONS η _ as) (NS n) = go (δ . η) as n
  go _ ENil _ = error "impossible name"

wkEnv :: Wk j k -> Env i j -> Env i k
wkEnv WZ x = x
wkEnv _ ENil = ENil
wkEnv η (ECONS δ h t) = ECONS (η . δ) h t

#else

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 :: {-# unpack #-} !(Wk i o) -> !(Val i) -> !(Tree i j k) -> !(Tree i k l) -> Tree o (S j) l

pattern TBin :: () => (h ~ S j) => Val i -> Tree i j k -> Tree i k l -> Tree i h l
pattern TBin a l r <- TBIN η (wkVal η -> a) (wkTree η -> l) (wkTree η -> r) where
  TBin a l r = TBIN id a l r

{-# complete TTip, TBin #-}

wkTree :: Wk i o -> Tree i j k -> Tree o j k
wkTree δ (TTip a) = TTip (wkVal δ a)
wkTree δ (TBIN η a l r) = TBIN (δ . η) a l r

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

wkEnv :: Wk j k -> Env i j -> Env i k
wkEnv WZ xs = xs
wkEnv δ (ESkCons n η t ts) = ESkCons n (δ . η) t ts
wkEnv _ ENil = ENil

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) id (TBin a (wkTree η l) (wkTree γ r)) $ wkEnv γ xs
econs a xs = ESkCons 1 id (TTip a) xs

upEnv :: Env i j -> Env' i j
upEnv ENil = ENil'
upEnv (ESkCons _ η (TTip a) xs)     = ECons' (wkVal η a) $ wkEnv η xs
upEnv (ESkCons n η (TBin a l r) xs) = ECons' (wkVal η a) $ ESkCons n' η l $ ESkCons n' id 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`

{-# complete ENil, ECons #-}

evalVar :: Env i j -> Name i -> Val j
#ifdef SAFE
evalVar (ECons a _) = \case
  NZ -> a
  NS n -> evalVar as n
evalVar ENil = error "impossible name"
#else
evalVar = go id where
  go :: Wk j k -> Env i j -> Name i -> Val k
  go !δ (ESkCons n η t xs) !m
    | n <= indexName m = go (δ . η) xs (UnsafeName (indexName m-n))
    | otherwise = goTree (δ . η) n t m
  go _ ENil _ = error "impossible name"

  goTree :: Wk i l -> Int -> Tree i j k -> Name j -> Val l
  goTree δ _ (TTip a) NZ = wkVal δ a
  goTree _ _ (TTip _) _ = error "impossible name"
  goTree δ _ (TBIN η a _ _) NZ = wkVal (δ . η)  a
  goTree δ n (TBIN η _ l r) m
    | indexName m <= n' = goTree (δ . η) n' l $ UnsafeName (indexName m - 1)
    | otherwise = goTree (δ . η) n' r $ UnsafeName (indexName m - n' - 1)
    where n' = unsafeShiftR n 1
#endif

#endif

-------------------------------------------------------------------------------
-- * 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 b) = VLam id \w v -> eval (v `ECons` wkEnv w 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)

quote :: Val i -> Expr i
quote (VLam η b) = Lam $ quote $ b (wk . η) (VVar NZ SNil)
quote (VVar h s) = quoteSp s (Var h)

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

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

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

ki :: Val i
ki = eval (i `ECons` k `ECons` ENil) (App (Var (NS NZ)) (Var NZ)) where
  i = eval ENil $ Lam (Var NZ)
  k = eval ENil $ Lam (Lam (Var (NS NZ)))

main :: IO ()
main = do
  print $ quote ki