{-# LANGUAGE CPP #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedRecordDot #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE RoleAnnotations #-}
{-# LANGUAGE ViewPatterns #-}
{-# OPTIONS_GHC -Wno-duplicate-exports #-}

module Data.DeBruijn.Environment.Fast (
  -- * Environments
  Env (Nil, (:>)),
  (!),

  -- * Fast
  EnvRep,
  Env (UnsafeEnv, envRep),
) where

import Data.DeBruijn.Index.Fast (Ix (ixRep), ixRepToInt)
import Data.Kind (Type)
import Data.Type.Nat (Nat (..), Pos, Pred)
import Unsafe.Coerce (unsafeCoerce)

--------------------------------------------------------------------------------
-- Conditional Imports
--------------------------------------------------------------------------------

#if defined(ENV_AS_SEQ)
import Data.Sequence (Seq)
import Data.Sequence qualified as Seq
import Data.Maybe (fromJust)
#elif defined(ENV_AS_SKEW_LIST)
import Data.SkewList.Strict (SkewList)
import Data.SkewList.Strict qualified as SkewList
#endif

--------------------------------------------------------------------------------
-- Environment Representation
--------------------------------------------------------------------------------

--------------------------------------------------------------------------------
-- Environment Representation: Finger Tree
#if defined(ENV_AS_SEQ)

type EnvRep a = Seq a

mkNilRep :: EnvRep a
mkNilRep :: forall a. EnvRep a
mkNilRep = Seq a
forall a. EnvRep a
Seq.empty
{-# INLINE mkNilRep #-}

mkSnocRep :: EnvRep a -> a -> EnvRep a
mkSnocRep :: forall a. EnvRep a -> a -> EnvRep a
mkSnocRep EnvRep a
xs a
x = EnvRep a
xs EnvRep a -> a -> EnvRep a
forall a. EnvRep a -> a -> EnvRep a
Seq.:|> a
x
{-# INLINE mkSnocRep #-}

elEnvRep :: b -> (EnvRep a -> a -> b) -> EnvRep a -> b
elEnvRep :: forall b a. b -> (EnvRep a -> a -> b) -> EnvRep a -> b
elEnvRep b
ifNil EnvRep a -> a -> b
ifSnoc = \case
  EnvRep a
Seq.Empty -> b
ifNil
  EnvRep a
xs Seq.:|> a
x -> EnvRep a -> a -> b
ifSnoc EnvRep a
xs a
x
{-# INLINE elEnvRep #-}

lookupRep :: Int -> EnvRep a -> a
lookupRep :: forall a. Int -> EnvRep a -> a
lookupRep = (Maybe a -> a
forall a. HasCallStack => Maybe a -> a
fromJust .) ((EnvRep a -> Maybe a) -> EnvRep a -> a)
-> (Int -> EnvRep a -> Maybe a) -> Int -> EnvRep a -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> EnvRep a -> Maybe a
forall a. Int -> Seq a -> Maybe a
Seq.lookup
{-# INLINE lookupRep #-}
{-# ANN lookupRep ("HLint: ignore Avoid partial function" :: String) #-}

--------------------------------------------------------------------------------
-- Environment Representation: Skew List
#elif defined(ENV_AS_SKEW_LIST)

type EnvRep a = SkewList a

mkNilRep :: EnvRep a
mkNilRep = SkewList.Nil
{-# INLINE mkNilRep #-}

mkSnocRep :: EnvRep a -> a -> EnvRep a
mkSnocRep xs x = SkewList.Cons x xs
{-# INLINE mkSnocRep #-}

elEnvRep :: b -> (EnvRep a -> a -> b) -> EnvRep a -> b
elEnvRep ifNil ifSnoc = \case
  SkewList.Nil -> ifNil
  SkewList.Cons x xs -> ifSnoc xs x
{-# INLINE elEnvRep #-}

lookupRep :: Int -> EnvRep a -> a
lookupRep = flip (SkewList.!)
{-# INLINE lookupRep #-}
{-# ANN lookupRep ("HLint: ignore Avoid partial function" :: String) #-}

--------------------------------------------------------------------------------
-- Environment Representation: None
#elif !defined(__HLINT__)
#error "cpp: define one of [ENV_AS_SEQ, ENV_AS_SKEW_LIST]"
#endif

--------------------------------------------------------------------------------
-- Environments
--------------------------------------------------------------------------------

type Env :: Nat -> Type -> Type
newtype Env n a = UnsafeEnv {forall (n :: Nat) a. Env n a -> EnvRep a
envRep :: EnvRep a}

type role Env nominal representational

deriving stock instance Functor (Env n)

deriving stock instance Foldable (Env n)

deriving stock instance Traversable (Env n)

mkNil :: Env Z a
mkNil :: forall a. Env Z a
mkNil = EnvRep a -> Env Z a
forall (n :: Nat) a. EnvRep a -> Env n a
UnsafeEnv EnvRep a
forall a. EnvRep a
mkNilRep
{-# INLINE mkNil #-}

mkSnoc :: Env n a -> a -> Env (S n) a
mkSnoc :: forall (n :: Nat) a. Env n a -> a -> Env (S n) a
mkSnoc = (EnvRep a -> Env (S n) a
forall (n :: Nat) a. EnvRep a -> Env n a
UnsafeEnv .) ((a -> EnvRep a) -> a -> Env (S n) a)
-> (Env n a -> a -> EnvRep a) -> Env n a -> a -> Env (S n) a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. EnvRep a -> a -> EnvRep a
forall a. EnvRep a -> a -> EnvRep a
mkSnocRep (EnvRep a -> a -> EnvRep a)
-> (Env n a -> EnvRep a) -> Env n a -> a -> EnvRep a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (.envRep)
{-# INLINE mkSnoc #-}

elEnv :: b -> (Env (Pred n) a -> a -> b) -> Env n a -> b
elEnv :: forall b (n :: Nat) a.
b -> (Env (Pred n) a -> a -> b) -> Env n a -> b
elEnv b
ifNil Env (Pred n) a -> a -> b
ifSnoc = b -> (EnvRep a -> a -> b) -> EnvRep a -> b
forall b a. b -> (EnvRep a -> a -> b) -> EnvRep a -> b
elEnvRep b
ifNil (Env (Pred n) a -> a -> b
ifSnoc (Env (Pred n) a -> a -> b)
-> (EnvRep a -> Env (Pred n) a) -> EnvRep a -> a -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. EnvRep a -> Env (Pred n) a
forall (n :: Nat) a. EnvRep a -> Env n a
UnsafeEnv) (EnvRep a -> b) -> (Env n a -> EnvRep a) -> Env n a -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (.envRep)
{-# INLINE elEnv #-}

type EnvF :: (Nat -> Type -> Type) -> Nat -> Type -> Type
data EnvF env n a where
  NilF :: EnvF env Z a
  SnocF :: env n a -> a -> EnvF env (S n) a

projectEnv :: Env n a -> EnvF Env n a
projectEnv :: forall (n :: Nat) a. Env n a -> EnvF Env n a
projectEnv = EnvF Env n a
-> (Env (Pred n) a -> a -> EnvF Env n a) -> Env n a -> EnvF Env n a
forall b (n :: Nat) a.
b -> (Env (Pred n) a -> a -> b) -> Env n a -> b
elEnv (EnvF Any Z Any -> EnvF Env n a
forall a b. a -> b
unsafeCoerce EnvF Any Z Any
forall (env :: Nat -> * -> *) a. EnvF env Z a
NilF) ((a -> EnvF Env (S (Pred n)) a) -> a -> EnvF Env n a
forall a b. a -> b
unsafeCoerce ((a -> EnvF Env (S (Pred n)) a) -> a -> EnvF Env n a)
-> (Env (Pred n) a -> a -> EnvF Env (S (Pred n)) a)
-> Env (Pred n) a
-> a
-> EnvF Env n a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Env (Pred n) a -> a -> EnvF Env (S (Pred n)) a
forall (env :: Nat -> * -> *) (n :: Nat) a.
env n a -> a -> EnvF env (S n) a
SnocF)
{-# INLINE projectEnv #-}

embedEnv :: EnvF Env n a -> Env n a
embedEnv :: forall (n :: Nat) a. EnvF Env n a -> Env n a
embedEnv = \case
  EnvF Env n a
NilF -> Env n a
Env Z a
forall a. Env Z a
mkNil
  SnocF Env n a
xs a
x -> Env n a -> a -> Env (S n) a
forall (n :: Nat) a. Env n a -> a -> Env (S n) a
mkSnoc Env n a
xs a
x
{-# INLINE embedEnv #-}

pattern Nil :: () => (n ~ Z) => Env n a
pattern $mNil :: forall {r} {n :: Nat} {a}.
Env n a -> ((n ~ Z) => r) -> ((# #) -> r) -> r
$bNil :: forall (n :: Nat) a. (n ~ Z) => Env n a
Nil <- (projectEnv -> NilF) where Nil = EnvF Env n a -> Env n a
forall (n :: Nat) a. EnvF Env n a -> Env n a
embedEnv EnvF Env n a
EnvF Env Z a
forall (env :: Nat -> * -> *) a. EnvF env Z a
NilF
{-# INLINE Nil #-}

pattern (:>) :: () => (Pos n) => Env (Pred n) a -> a -> Env n a
pattern $m:> :: forall {r} {n :: Nat} {a}.
Env n a -> (Pos n => Env (Pred n) a -> a -> r) -> ((# #) -> r) -> r
$b:> :: forall (n :: Nat) a. Pos n => Env (Pred n) a -> a -> Env n a
(:>) xs x <- (projectEnv -> SnocF xs x) where (:>) Env (Pred n) a
xs a
x = EnvF Env n a -> Env n a
forall (n :: Nat) a. EnvF Env n a -> Env n a
embedEnv (Env (Pred n) a -> a -> EnvF Env (S (Pred n)) a
forall (env :: Nat -> * -> *) (n :: Nat) a.
env n a -> a -> EnvF env (S n) a
SnocF Env (Pred n) a
xs a
x)
{-# INLINE (:>) #-}

{-# COMPLETE Nil, (:>) #-}

(!) :: Env n a -> Ix n -> a
Env n a
xs ! :: forall (n :: Nat) a. Env n a -> Ix n -> a
! Ix n
i = Int -> EnvRep a -> a
forall a. Int -> EnvRep a -> a
lookupRep (Int -> Int
ixRepToInt Ix n
i.ixRep) Env n a
xs.envRep