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

import Data.Coerce          (coerce)
import Data.Functor.Reverse (Reverse (..))
import Data.Kind            (Type)
import Data.SkewList.Lazy   (SkewList)
import Unsafe.Coerce        (unsafeCoerce)

import qualified Data.SkewList.Lazy as SL

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

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

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

-- | Environment
--
-- >>> EmptyEnv :> 'a' :> 'b'
-- EmptyEnv :> 'a' :> 'b'
--
type Env :: Ctx -> Type -> Type
type role Env nominal representational
newtype Env ctx a = UnsafeEnv { forall (ctx :: Ctx) a. Env ctx a -> SkewList a
unEnv :: SkewList a }

-------------------------------------------------------------------------------
-- Pattern synonyms
-------------------------------------------------------------------------------

-- We need a GADT to implement pattern synonyms.
type ViewEnv :: Ctx -> Type -> Type
type role ViewEnv nominal representational
data ViewEnv ctx a where
    EmptyViewEnv :: ViewEnv EmptyCtx a
    (:>>)     :: Env ctx a -> a -> ViewEnv (S ctx) a

viewEnv :: Env ctx a -> ViewEnv ctx a
viewEnv :: forall (ctx :: Ctx) a. Env ctx a -> ViewEnv ctx a
viewEnv (UnsafeEnv SkewList a
env) = case SkewList a -> Maybe (a, SkewList a)
forall a. SkewList a -> Maybe (a, SkewList a)
SL.uncons SkewList a
env of
    Maybe (a, SkewList a)
Nothing      -> ViewEnv EmptyCtx Any -> ViewEnv ctx a
forall a b. a -> b
unsafeCoerce ViewEnv EmptyCtx Any
forall a. ViewEnv EmptyCtx a
EmptyViewEnv
    Just (a
x, SkewList a
xs) -> ViewEnv (S Any) a -> ViewEnv ctx a
forall a b. a -> b
unsafeCoerce (SkewList a -> Env Any a
forall (ctx :: Ctx) a. SkewList a -> Env ctx a
UnsafeEnv SkewList a
xs Env Any a -> a -> ViewEnv (S Any) a
forall (ctx :: Ctx) a. Env ctx a -> a -> ViewEnv (S ctx) a
:>> a
x)

pattern EmptyEnv :: () => (ctx ~ EmptyCtx) => Env ctx a
pattern $mEmptyEnv :: forall {r} {ctx :: Ctx} {a}.
Env ctx a -> ((ctx ~ EmptyCtx) => r) -> ((# #) -> r) -> r
$bEmptyEnv :: forall (ctx :: Ctx) a. (ctx ~ EmptyCtx) => Env ctx a
EmptyEnv <- (viewEnv -> EmptyViewEnv)
  where EmptyEnv = SkewList a -> Env ctx a
forall (ctx :: Ctx) a. SkewList a -> Env ctx a
UnsafeEnv SkewList a
forall a. SkewList a
SL.Nil

infixl 5 :>
pattern (:>) :: () => (ctx ~ S ctx') => Env ctx' a -> a -> Env ctx a
pattern xs $m:> :: forall {r} {ctx :: Ctx} {a}.
Env ctx a
-> (forall {ctx' :: Ctx}. (ctx ~ S ctx') => Env ctx' a -> a -> r)
-> ((# #) -> r)
-> r
$b:> :: forall (ctx :: Ctx) a (ctx' :: Ctx).
(ctx ~ S ctx') =>
Env ctx' a -> a -> Env ctx a
:> x <- (viewEnv -> xs :>> x)
  where Env ctx' a
xs :> a
x = SkewList a -> Env ctx a
forall (ctx :: Ctx) a. SkewList a -> Env ctx a
UnsafeEnv (a -> SkewList a -> SkewList a
forall a. a -> SkewList a -> SkewList a
SL.cons a
x (Env ctx' a -> SkewList a
forall (ctx :: Ctx) a. Env ctx a -> SkewList a
unEnv Env ctx' a
xs))

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

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

deriving instance Functor (Env ctx)

-- |
--
-- >>> toList (tabulateEnv S3 id)
-- [2,1,0]
--
instance Foldable (Env ctx) where
    foldMap :: forall m a. Monoid m => (a -> m) -> Env ctx a -> m
foldMap a -> m
f (UnsafeEnv SkewList a
xs) = (a -> m) -> Reverse SkewList a -> m
forall m a. Monoid m => (a -> m) -> Reverse SkewList a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap a -> m
f (SkewList a -> Reverse SkewList a
forall {k} (f :: k -> *) (a :: k). f a -> Reverse f a
Reverse SkewList a
xs)

-- |
--
-- >>> traverse print (tabulateEnv S3 id)
-- 2
-- 1
-- 0
-- EmptyEnv :> () :> () :> ()
--
instance Traversable (Env ctx) where
    traverse :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Env ctx a -> f (Env ctx b)
traverse a -> f b
f (UnsafeEnv SkewList a
xs) = (Reverse SkewList b -> Env ctx b)
-> f (Reverse SkewList b) -> f (Env ctx b)
forall a b. (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (SkewList b -> Env ctx b
forall (ctx :: Ctx) a. SkewList a -> Env ctx a
UnsafeEnv (SkewList b -> Env ctx b)
-> (Reverse SkewList b -> SkewList b)
-> Reverse SkewList b
-> Env ctx b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Reverse SkewList b -> SkewList b
forall {k} (f :: k -> *) (a :: k). Reverse f a -> f a
getReverse) ((a -> f b) -> Reverse SkewList a -> f (Reverse SkewList b)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Reverse SkewList a -> f (Reverse SkewList b)
traverse a -> f b
f (SkewList a -> Reverse SkewList a
forall {k} (f :: k -> *) (a :: k). f a -> Reverse f a
Reverse SkewList a
xs))

instance Show a => Show (Env ctx a) where
    showsPrec :: Int -> Env ctx a -> ShowS
showsPrec Int
d0 (UnsafeEnv SkewList a
xs0) = Int -> SkewList a -> ShowS
forall {t} {a}. (Ord t, Num t, Show a) => t -> SkewList a -> ShowS
go Int
d0 SkewList a
xs0 where
        go :: t -> SkewList a -> ShowS
go t
d SkewList a
ys = case SkewList a -> Maybe (a, SkewList a)
forall a. SkewList a -> Maybe (a, SkewList a)
SL.uncons SkewList a
ys of
            Maybe (a, SkewList a)
Nothing      -> String -> ShowS
showString String
"EmptyEnv"
            Just (a
x, SkewList a
xs) -> Bool -> ShowS -> ShowS
showParen (t
d t -> t -> Bool
forall a. Ord a => a -> a -> Bool
> t
5)
                (ShowS -> ShowS) -> ShowS -> ShowS
forall a b. (a -> b) -> a -> b
$ t -> SkewList a -> ShowS
go t
5 SkewList 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 (UnsafeIdx Int
i) (UnsafeEnv SkewList a
xs) = SkewList a
xs SkewList a -> Int -> a
forall a. HasCallStack => SkewList a -> Int -> a
SL.! Int
i

-- | 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 (UnsafeEnv SkewList a
xs) = Int -> Size n
forall (ctx :: Ctx). Int -> Size ctx
UnsafeSize (SkewList a -> Int
forall a. SkewList a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length SkewList a
xs)

-- | Create 'Env' by creating elements given an 'Idx'.
--
-- >>> tabulateEnv S4 id
-- EmptyEnv :> 3 :> 2 :> 1 :> 0
--
tabulateEnv :: Size ctx -> (Idx ctx -> a) -> Env ctx a
tabulateEnv :: forall (ctx :: Ctx) a. Size ctx -> (Idx ctx -> a) -> Env ctx a
tabulateEnv (UnsafeSize Int
s) Idx ctx -> a
f = SkewList a -> Env ctx a
forall (ctx :: Ctx) a. SkewList a -> Env ctx a
UnsafeEnv (SkewList a -> Env ctx a) -> SkewList a -> Env ctx a
forall a b. (a -> b) -> a -> b
$ [a] -> SkewList a
forall a. [a] -> SkewList a
SL.fromList ([a] -> SkewList a) -> [a] -> SkewList a
forall a b. (a -> b) -> a -> b
$ (Int -> a) -> [Int] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map ((Idx ctx -> a) -> Int -> a
forall a b. Coercible a b => a -> b
coerce Idx ctx -> a
f) [Int
0 .. Int
s Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]