Skip to content

Instantly share code, notes, and snippets.

@silky
Forked from ekmett/SizedNBE.hs
Created April 13, 2021 11:55
Show Gist options
  • Save silky/2ed477a4cf0fdde7539de9a88d3a4ca2 to your computer and use it in GitHub Desktop.
Save silky/2ed477a4cf0fdde7539de9a88d3a4ca2 to your computer and use it in GitHub Desktop.
Sized NBE using the same schema as the typed one, but with natural number sizes on environments
{-# 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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment