{-# LANGUAGE Safe #-}
module DeBruijn.Env (
    Env (EmptyEnv, (:>)),
    lookupEnv,
    sizeEnv,
    tabulateEnv,
) where

import Data.Kind (Type)

import DeBruijn.Ctx
import DeBruijn.Idx
import DeBruijn.Size

-- $setup
-- >>> import DeBruijn
-- >>> import Data.Foldable (toList)

-------------------------------------------------------------------------------
-- Environment
-------------------------------------------------------------------------------

-- | Environment
--
-- >>> EmptyEnv :> 'a' :> 'b'
-- EmptyEnv :> 'a' :> 'b'
--
type Env :: Ctx -> Type -> Type
data Env ctx a where
    EmptyEnv :: Env EmptyCtx a
    (:>)     :: Env ctx a -> a -> Env (S ctx) a

infixl 5 :>

-------------------------------------------------------------------------------
-- Instances
-------------------------------------------------------------------------------

deriving instance Functor (Env ctx)

-- |
--
-- >>> toList (tabulateEnv S3 id)
-- [2,1,0]
--
deriving instance Foldable (Env ctx)

-- |
--
-- >>> traverse print (tabulateEnv S3 id)
-- 2
-- 1
-- 0
-- EmptyEnv :> () :> () :> ()
--
deriving instance Traversable (Env ctx)

instance Show a => Show (Env ctx a) where
    showsPrec :: Int -> Env ctx a -> ShowS
showsPrec Int
_ Env ctx a
EmptyEnv = String -> ShowS
showString String
"EmptyEnv"
    showsPrec Int
d (Env ctx a
xs :> a
x) = Bool -> ShowS -> ShowS
showParen (Int
d Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
5)
        (ShowS -> ShowS) -> ShowS -> ShowS
forall a b. (a -> b) -> a -> b
$ Int -> Env ctx a -> ShowS
forall a. Show a => Int -> a -> ShowS
showsPrec Int
5 Env ctx a
xs
        ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> ShowS
showString String
" :> "
        ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> a -> ShowS
forall a. Show a => Int -> a -> ShowS
showsPrec Int
6 a
x

-------------------------------------------------------------------------------
-- Combinators
-------------------------------------------------------------------------------

-- | Lookup in the context.
--
-- >>> lookupEnv IZ (EmptyEnv :> 'a' :> 'b')
-- 'b'
--
lookupEnv :: Idx ctx -> Env ctx a -> a
lookupEnv :: forall (ctx :: Ctx) a. Idx ctx -> Env ctx a -> a
lookupEnv Idx ctx
IZ     (Env ctx a
_  :> a
x)  = a
x
lookupEnv (IS Idx n1
n) (Env ctx a
xs :> a
_) = Idx n1 -> Env n1 a -> a
forall (ctx :: Ctx) a. Idx ctx -> Env ctx a -> a
lookupEnv Idx n1
n Env n1 a
Env ctx a
xs

-- | Size of the environment.
--
-- >>> sizeEnv (EmptyEnv :> 'a' :> 'b')
-- 2
--
sizeEnv :: Env n a -> Size n
sizeEnv :: forall (n :: Ctx) a. Env n a -> Size n
sizeEnv Env n a
EmptyEnv  = Size n
Size 'Z
SZ
sizeEnv (Env ctx a
xs :> a
_) = Size ctx -> Size ('S ctx)
forall (ctx1 :: Ctx). Size ctx1 -> Size ('S ctx1)
SS (Env ctx a -> Size ctx
forall (n :: Ctx) a. Env n a -> Size n
sizeEnv Env ctx a
xs)

tabulateEnv :: Size ctx -> (Idx ctx -> a) -> Env ctx a
tabulateEnv :: forall (ctx :: Ctx) a. Size ctx -> (Idx ctx -> a) -> Env ctx a
tabulateEnv Size ctx
SZ     Idx ctx -> a
_ = Env ctx a
Env 'Z a
forall a. Env 'Z a
EmptyEnv
tabulateEnv (SS Size ctx1
s) Idx ctx -> a
f = Size ctx1 -> (Idx ctx1 -> a) -> Env ctx1 a
forall (ctx :: Ctx) a. Size ctx -> (Idx ctx -> a) -> Env ctx a
tabulateEnv Size ctx1
s (Idx ctx -> a
f (Idx ctx -> a) -> (Idx ctx1 -> Idx ctx) -> Idx ctx1 -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Idx ctx1 -> Idx ctx
Idx ctx1 -> Idx ('S ctx1)
forall (n1 :: Ctx). Idx n1 -> Idx ('S n1)
IS) Env ctx1 a -> a -> Env ('S ctx1) a
forall (ctx :: Ctx) a. Env ctx a -> a -> Env (S ctx) a
:> Idx ctx -> a
f Idx ctx
Idx ('S ctx1)
forall (n1 :: Ctx). Idx ('S n1)
IZ