Skip to content

Instantly share code, notes, and snippets.

@ekmett
Last active September 7, 2020 19:37
Show Gist options
  • Save ekmett/586f06db7ccb0d15b4814babc8155b3b to your computer and use it in GitHub Desktop.
Save ekmett/586f06db7ccb0d15b4814babc8155b3b to your computer and use it in GitHub Desktop.
Typed normalization-by-evaluation using a slowed category action
{-# 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
module TypedNBE
(
-- * Names
Name
( NZ
, NS
#ifndef SAFE
, UnsafeName
#endif
)
, indexName
-- * Weakening
, Wk
( WZ
, WS
#ifndef SAFE
, UnsafeWk
#endif
)
, wk
, indexWk
-- * Environments
, Env
( ..
, ECons
)
, evalVar
, wkEnv
-- * Domain
, Val(..)
, Spine(.., SApp)
, wkVal
, wkSp
-- * Evaluation
, eval
, quote
, nf
-- * example
, k, id_, kid_
, main
) where
import Control.Category (Category(..))
import GHC.Types (Type)
import Prelude hiding (id, (.))
#ifndef NOSKEW
import Data.Bits (unsafeShiftR)
#endif
#ifndef SAFE
import Unsafe.Coerce
#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 :: [i] -> i -> Type
type role Name nominal nominal
data Name j a where
NZ :: Name (a ': j) a
NS :: Name j a -> Name (b ': j) a
indexName :: Name j a -> Int
indexName = go 0 where
go :: Int -> Name j a -> Int
go !acc NZ = acc
go acc (NS n) = go (acc + 1) n
instance Show (Name j a) where
showsPrec d = showsPrec d . indexName
#else
type Name' :: [i] -> i -> Type
type role Name' nominal nominal
data Name' j a where
NZ' :: Name' (a ': j) a
NS' :: Name j a -> Name' (b ': j) a
type Name :: [i] -> i -> Type
type role Name nominal nominal
newtype Name j a = UnsafeName { indexName :: Int }
instance Show (Name j a) where
showsPrec d = showsPrec d . indexName
upName :: Name i ~> Name' i
upName (UnsafeName 0) = unsafeCoerce NZ'
upName (UnsafeName n) = unsafeCoerce (NS' $ UnsafeName (n-1))
pattern NZ :: () => (b ~ a) => Name (b ': as) a
pattern NZ <- (upName -> NZ') where
NZ = UnsafeName 0
pattern NS :: () => (a ~ a', as ~ (b ': as')) => Name as' a' -> Name as a
pattern NS n <- (upName -> NS' n) where
NS n = UnsafeName (indexName n + 1)
{-# complete NZ, NS #-}
#endif
-------------------------------------------------------------------------------
-- * Expressions
-------------------------------------------------------------------------------
type Expr :: [Type] -> Type -> Type
type role Expr nominal nominal
data Expr (e :: [Type]) (a :: Type) where
Var :: {-# unpack #-} !(Name i a) -> Expr i a
App :: !(Expr i (a -> b)) -> !(Expr i a) -> Expr i b
Lam :: !(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 (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 :: [i] -> [i] -> Type
type role Wk nominal nominal
data Wk i j where
WZ :: Wk i i
WS :: Wk i j -> Wk i (a ': 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 :: [i] -> [i] -> 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' :: [i] -> [i] -> Type
type role Wk' nominal nominal
data Wk' i j where
WZ' :: Wk' i i
WS' :: Wk i j -> Wk' i (a ': 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 ~ (a ': 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 (a ': i)
wk = UnsafeWk 1
#endif
-------------------------------------------------------------------------------
-- * Semantic Domain
-------------------------------------------------------------------------------
type Val :: [Type] -> Type -> Type
type role Val nominal nominal
data Val i a where
-- VLam :: !(Env i j) -> !(Expr (a ': j) b) -> Val i (a -> b)
VLam :: {-# unpack #-} !(Wk i j) -> (forall k. Wk i k -> Val k a -> Val k b) -> Val j (a -> b)
VVar :: {-# unpack #-} !(Name 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 :: {-# unpack #-} !(Wk i j) -> !(Spine i x (a -> b)) -> !(Val i a) -> Spine j x b
pattern SApp :: Spine i x (a -> b) -> Val i a -> Spine i x b
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 a b -> Spine j a b
wkSp _ SNil = SNil
wkSp δ (SAPP η s v) = SAPP (δ . η) s v
-------------------------------------------------------------------------------
-- * Environments
-------------------------------------------------------------------------------
#ifdef NOSKEW
type role Env nominal nominal
data Env k i where
ENil :: Env '[] i
ECONS :: {-# unpack #-} !(Wk i j) -> !(Val i a) -> !(Env k i) -> Env (a ': k) j
pattern ECons :: () => (l ~ (a ': k)) => Val i a -> 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 :: [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 :: {-# unpack #-} !(Wk i o) -> !(Val i a) -> !(Tree i j k) -> !(Tree i k l) -> Tree o (a ': j) l
pattern TBin :: () => (h ~ (a ': j)) => Val i a -> 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 :: [Type] -> [Type] -> 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 '[] 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' :: [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) 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 ~ (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 -> Name i ~> Val j
#ifdef SAFE
evalVar (ECons a _) NZ = a
evalVar (ECons _ as) (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 (a -> b) -> Val i a -> Val i b
evalApp (VLam w f) v = f w v
evalApp (VVar h s) v = VVar h (SApp s v)
quote :: Val i ~> Expr i
quote (VLam w b) = Lam $ quote $ b (wk . w) (VVar NZ SNil)
quote (VVar h s) = quoteSp s (Var h)
quoteSp :: Spine i a b -> Expr i a -> Expr i b
quoteSp SNil e = e
quoteSp (SApp s x) e = App (quoteSp s e) (quote x)
nf :: Env i j -> Expr i ~> Expr j
nf e t = quote (eval e t)
-------------------------------------------------------------------------------
-- * Tests
-------------------------------------------------------------------------------
id_ :: Val i (a -> a)
id_ = eval ENil $ Lam (Var NZ)
k :: Val i (a -> b -> a)
k = eval ENil $ Lam (Lam (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 kid_
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment