a pragmatic mix of de bruijn indices and levels

IndicesAndLevels.hs

{-# language PolyKinds #-}
{-# language BlockArguments #-}
{-# language AllowAmbiguousTypes #-}
{-# language StrictData #-}
{-# language DerivingStrategies #-}
{-# language GeneralizedNewtypeDeriving #-}
{-# language TypeApplications #-}
{-# language BangPatterns #-}
{-# language NPlusKPatterns #-}
{-# language TypeFamilies #-}
{-# language RoleAnnotations #-}
{-# language ViewPatterns #-}
{-# language StandaloneKindSignatures #-}
{-# language OverloadedStrings #-}
{-# language ScopedTypeVariables #-}
{-# language RankNTypes #-}
{-# language DataKinds #-}
{-# language GADTs #-}
{-# language TypeOperators #-}
{-# language PatternSynonyms #-}
{-# language TypeApplications #-}
{-# language TemplateHaskell #-}
{-# language LambdaCase #-}
{-# language ImportQualifiedPost #-}

module Simple where

import Control.Category
import Control.Exception
import Control.Monad
import Control.Monad.IO.Class
import Data.Bits
import Data.Coerce
import Data.Functor
import Data.IORef
import Data.Kind
import Data.Maybe
import Data.Proxy
import Text.Show.Deriving
import Unsafe.Coerce
import Numeric.Natural
import System.IO.Unsafe
import Prelude hiding (id, (.))

-- see github.com/ekmett/haskell types, mostly used because I can't be bothered to strip it out, not important
import Data.Type hiding (SNil, SCons)
import Data.Type.Unsafe
import Data.Type.Equality hiding (apply)


data Icit = I | E
  deriving (Eq,Ord,Show,Read,Enum,Bounded)

makeSing ''Icit

--------------------------------------------------------------------------------
-- * De Bruijn Indices
--------------------------------------------------------------------------------

type Ix :: Int -> Type        
newtype Ix i = Ix Int     
  deriving newtype Show

type Ix' :: Int -> Type
data Ix' i where
  IZ' :: Ix' (S n)
  IS' :: Ix n -> Ix' (S n)

ix' :: Ix i -> Ix' i
ix' (Ix 0) = unsafeCoerce IZ'
ix' (Ix n) = unsafeCoerce (IS' (Ix (n-1)))

pattern IZ :: () => n ~ (S n') => Ix n
pattern IZ <- (ix' -> IZ') where
  IZ = Ix 0

pattern IS :: () => n ~ S n' => Ix n' -> Ix n
pattern IS n <- (ix' -> IS' n) where
  IS (Ix n) = Ix (n+1)

--------------------------------------------------------------------------------
-- * De Bruijn Levels
--------------------------------------------------------------------------------

type Lvl :: Type
newtype Lvl = Lvl Int
  deriving newtype (Eq,Ord,Show,Read,Num,Enum,Real,Integral)

lvlIx :: forall i. Sing i => Lvl -> Ix i
lvlIx (Lvl l) = Ix (reflect @_ @i - l - 1)

ixLvl :: forall i. Sing i => Ix i -> Lvl
ixLvl (Ix l) = Lvl (reflect @_ @i - l - 1)

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

type Tree :: Int -> Int -> Type -> Type
type role Tree nominal nominal representational
data Tree i j a where
  TTip :: ~a -> Tree (S j) j a
  TBin :: ~a -> Tree j k a -> Tree k l a -> Tree (S j) l a

instance Functor (Tree i j) where
  fmap f (TTip a) = TTip (f a)
  fmap f (TBin a l r) = TBin (f a) (fmap f l) (fmap f r)

instance Foldable (Tree i j) where
  foldMap f (TTip a) = f a
  foldMap f (TBin a l r) = f a <> foldMap f l <> foldMap f r

instance Traversable (Tree i j) where
  traverse f (TTip a) = TTip <$> f a
  traverse f (TBin a l r) = TBin <$> f a <*> traverse f l <*> traverse f r

-- TODO: functor, foldable, traversable, applicative, monad

type Vec :: Int -> Type -> Type
type role Vec nominal representational
data Vec n a where
  TCons :: Int -> Int -> Tree n m a -> Vec m a -> Vec n a
  Nil   :: Vec Z a

instance Functor (Vec n) where
  fmap _ Nil = Nil
  fmap f (TCons s n t xs) = TCons s n (fmap f t) (fmap f xs)

instance Foldable (Vec n) where
  foldMap _ Nil = mempty
  foldMap f (TCons s n t xs) = foldMap f t <> foldMap f xs
  null Nil = True
  null _ = False
  length Nil = 0
  length (TCons s _ _ _) = s

-- TODO: functor, foldable, traversable, applicative, monad

type Vec' :: Int -> Type -> Type
type role Vec' nominal representational
data Vec' n a where
  Nil'  :: Vec' Z i
  Cons' :: ~a -> Vec n a -> Vec' (S n) a

vcons :: a -> Vec j a -> Vec (S j) a
vcons a (TCons s n l (TCons _ m r xs))
  | n == m = TCons (s+1) (2*n+1) (TBin a l r) xs
vcons a xs = TCons (1 + length xs) 1 (TTip a) xs

upVec :: Vec i a -> Vec' i a
upVec Nil = Nil'
upVec (TCons _ _ (TTip a) xs)     = Cons' a xs
upVec (TCons s n (TBin a l r) xs) = Cons' a $ TCons (s-1) n' l $ TCons (s-1-n') n' r xs
  where n' = unsafeShiftR n 1

pattern (:*) :: () => n ~ S m => a -> Vec m a -> Vec n a
pattern v :* e <- (upVec -> Cons' v e) where
  v :* e = vcons v e

infixr 5 :*

{-# complete Nil, (:*) #-}

index :: Vec i a -> Ix i -> a
index = go where
  go :: Vec i a -> Ix i -> a
  go (TCons _ n t xs) !m@(Ix im)
    | n <= im = go xs (Ix (im-n))
    | otherwise = goTree n t m
  go Nil _ = panic

  goTree :: Int -> Tree j k a -> Ix j -> a
  goTree _ (TTip a) IZ = a
  goTree _ (TTip _) _ = panic
  goTree _ (TBin a _ _) IZ = a
  goTree n (TBin _ l r) m@(Ix im)
    | im <= n' = goTree n' l $ Ix (im - 1)
    | otherwise = goTree n' r $ Ix (im - n' - 1)
    where n' = unsafeShiftR n 1

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

type Name = String

type Ty = Tm
data Tm i where
  Var  :: Ix i -> Tm i
  App  :: Icit -> Tm i -> Tm i -> Tm i
  Lam  :: Name -> Icit -> Ty i -> Tm (S i) -> Tm i
  Let  :: Name -> Ty i -> Tm i -> Tm (S i) -> Tm i
  Pi   :: Name -> Icit -> Ty i -> Ty (S i) -> Tm i
  U    :: Tm i

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

type Val :: Type
data Val where
  VLam  :: Name -> Icit -> ~VTy -> EVal -> Val -- type is lazy
  VPi   :: Name -> Icit -> ~VTy -> EVTy -> Val -- type is lazy
  VVar  :: Lvl -> Sp -> Val
  VU    :: Val

type VTy = Val
type EVal = Val -> IO Val
type EVTy = EVal

lvlVar :: Sing i => Lvl -> Tm i
lvlVar h = Var (lvlIx h)

topLvl :: forall (i::Int). SingI i => Lvl
topLvl = Lvl (reflect @Int @i)

topVal :: forall (i::Int). SingI i => Val
topVal = VVar (topLvl @i) SNil

data Sp where
  SNil :: Sp
  SApp :: Icit -> Sp -> ~Val -> Sp -- arg is lazy

{-# complete SNil, SApp #-}

panic :: a
panic = error "panic"

-------------------------------------------------------------------------------
-- * Utilities
-------------------------------------------------------------------------------

lazily :: IO a -> IO a
lazily = unsafeInterleaveIO

type Vals :: Int -> Type
type Vals i = Vec i Val

vskip :: Vals i -> Vals (S i)
vskip vs = VVar (Lvl $ length vs) SNil :* vs

-------------------------------------------------------------------------------
-- * Evaluation
-------------------------------------------------------------------------------

lazilyEval :: Vals i -> Tm i -> IO Val
lazilyEval e t = lazily (eval e t)

eval :: Vals i -> Tm i -> IO Val
eval e (Var n) = pure $ index e n
eval e (App i f x) = do
  fv <- eval e f
  xv <- lazilyEval e x
  apply i fv xv
eval e (Lam n i t b) = do
  tv <- lazilyEval e t
  pure $ VLam n i tv \v -> eval (v :* e) b
eval e (Let _ _ d b) = do
  v <- lazilyEval e d
  eval (v :* e) b
eval e (Pi n i t b) = do
  tv <- lazilyEval e t
  pure $ VPi n i tv \v -> eval (v :* e) b
eval _ U = pure VU

apply :: Icit -> Val -> Val -> IO Val
apply _ (VLam _ _ _ f) v = f v
apply i (VVar h s)     v = pure $ VVar h (SApp i s v)
apply _ _              _ = panic

applySp :: Val -> Sp -> IO Val
applySp h SNil = pure h
applySp h (SApp i sp u) = do
  sp' <- applySp h sp
  apply i sp' u

quote :: forall n. Sing n => Val -> IO (Tm n)
quote (VLam n i t b) = do
  t' <- quote t
  Lam n i t' <$> do
    v' <- b (topVal @n)
    quote v'
quote (VPi n i t b) = do
  t' <- quote t
  Pi n i t' <$> do
    v' <- b (topVal @n)
    quote v'
quote (VVar h s) = quoteSp s (lvlVar h)
quote VU = pure U

quoteSp :: Sing i => Sp -> Tm i -> IO (Tm i)
quoteSp SNil e = pure e
quoteSp (SApp i s x) e = App i <$> quoteSp s e <*> quote x

nf :: Sing j => Vals i -> Tm i -> IO (Tm j)
nf e t = eval e t >>= quote

deriveShow ''Tm

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

ki :: IO Val
ki = do
  i <- lazilyEval Nil $ Lam "x" E U $ Var IZ
  k <- lazilyEval Nil $ Lam "x" E U $ Lam "y" E U $ Var (IS IZ)
  eval (i :* k :* Nil) $ App E (Var (IS IZ)) (Var IZ)

main :: IO ()
main = do
  test <- ki
  test' <- quote @Z test
  print test'