Typed normalization-by-evaluation using a slowed category action

TypedNBE.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

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_