{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE UnboxedTuples       #-}
{-# OPTIONS_HADDOCK not-home #-}
-- |
-- Copyright: (c) 2021 Xy Ren
-- License: BSD3
-- Maintainer: xy.r@outlook.com
-- Stability: unstable
-- Portability: non-portable (GHC only)
--
-- This module defines the effect stack as an immutable extensible stack type, and provides functions for manipulating
-- it. The effect stack type has the following time complexities:
--
-- * Lookup: Amortized \( O(1) \).
-- * Update: \( O(n) \).
-- * Shrink: \( O(1) \).
-- * Append: \( O(n) \).
--
-- __This is an /internal/ module and its API may change even between minor versions.__ Therefore you should be
-- extra careful if you're to depend on this module.
module Cleff.Internal.Stack
  ( Effect
  , Stack
  , HandlerPtr (HandlerPtr, unHandlerPtr)
    -- * Construction
  , type (++)
  , empty
  , cons
  , concat
    -- * Deconstruction
  , KnownList
  , head
  , take
  , tail
  , drop
    -- * Retrieval and updating
  , (:>)
  , (:>>)
  , Subset
  , index
  , pick
  , update
  ) where

import           Data.Coerce  (coerce)
import           Data.Kind    (Constraint, Type)
import           Data.PrimVec (PrimVec)
import qualified Data.PrimVec as Vec
import           GHC.TypeLits (ErrorMessage (ShowType, Text, (:<>:)), TypeError)
import           Prelude      hiding (concat, drop, head, tail, take)

-- | The type of effects. An effect @e m a@ takes an effect monad type @m :: 'Type' -> 'Type'@ and a result type
-- @a :: 'Type'@.
type Effect = (Type -> Type) -> Type -> Type

-- | The effect stack, storing pointers to handlers. It is essentially an extensible stack type supporting
-- efficient \( O(1) \) reads.
type role Stack nominal
newtype Stack (es :: [Effect]) = Stack (PrimVec Int)

-- | A pointer to an effect handler.
type role HandlerPtr nominal
newtype HandlerPtr (e :: Effect) = HandlerPtr { HandlerPtr e -> Int
unHandlerPtr :: Int }

unreifiable :: String -> String -> String -> a
unreifiable :: String -> String -> String -> a
unreifiable String
clsName String
funName String
comp = String -> a
forall a. HasCallStack => String -> a
error (String -> a) -> String -> a
forall a b. (a -> b) -> a -> b
$
  String
funName String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
": Attempting to access " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
comp String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
" without a reflected value. This is perhaps because you are trying \
  \to define an instance for the '" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
clsName String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
"' typeclass, which you should not be doing whatsoever. If that or \
  \other shenanigans seem unlikely, please report this as a bug."

-- | Create an empty stack. \( O(1) \).
empty :: Stack '[]
empty :: Stack '[]
empty = PrimVec Int -> Stack '[]
coerce (Prim Int => PrimVec Int
forall a. Prim a => PrimVec a
Vec.empty @Int)

-- | Prepend one entry to the stack. \( O(n) \).
cons :: HandlerPtr e -> Stack es -> Stack (e : es)
cons :: HandlerPtr e -> Stack es -> Stack (e : es)
cons = (Int -> PrimVec Int -> PrimVec Int)
-> HandlerPtr e -> Stack es -> Stack (e : es)
coerce (Prim Int => Int -> PrimVec Int -> PrimVec Int
forall a. Prim a => a -> PrimVec a -> PrimVec a
Vec.cons @Int)

-- | Type level list concatenation.
type family xs ++ ys where
  '[] ++ ys = ys
  (x : xs) ++ ys = x : (xs ++ ys)
infixr 5 ++

-- | Concatenate two stacks. \( O(m+n) \).
concat :: Stack es -> Stack es' -> Stack (es ++ es')
concat :: Stack es -> Stack es' -> Stack (es ++ es')
concat = (PrimVec Int -> PrimVec Int -> PrimVec Int)
-> Stack es -> Stack es' -> Stack (es ++ es')
coerce (Prim Int => PrimVec Int -> PrimVec Int -> PrimVec Int
forall a. Prim a => PrimVec a -> PrimVec a -> PrimVec a
Vec.concat @Int)

-- | Slice off one entry from the top of the stack. \( O(1) \).
tail :: Stack (e : es) -> Stack es
tail :: Stack (e : es) -> Stack es
tail = (PrimVec Int -> PrimVec Int) -> Stack (e : es) -> Stack es
coerce (PrimVec Int -> PrimVec Int
forall a. PrimVec a -> PrimVec a
Vec.tail @Int)

-- | @'KnownList' es@ means the list @es@ is concrete, /i.e./ is of the form @'[a1, a2, ..., an]@ instead of a type
-- variable.
class KnownList (es :: [Effect]) where
  -- | Get the length of the list.
  reifyLen :: Int
  reifyLen = String -> String -> String -> Int
forall a. String -> String -> String -> a
unreifiable String
"KnownList" String
"Cleff.Internal.Stack.reifyLen" String
"the length of a type-level list"

instance KnownList '[] where
  reifyLen :: Int
reifyLen = Int
0

instance KnownList es => KnownList (e : es) where
  reifyLen :: Int
reifyLen = Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ KnownList es => Int
forall (es :: [Effect]). KnownList es => Int
reifyLen @es

-- | Slice off several entries from the top of the stack. \( O(1) \).
drop :: ∀ es es'. KnownList es => Stack (es ++ es') -> Stack es'
drop :: Stack (es ++ es') -> Stack es'
drop = (Int -> PrimVec Int -> PrimVec Int)
-> Int -> Stack (es ++ es') -> Stack es'
coerce (Int -> PrimVec Int -> PrimVec Int
forall a. Int -> PrimVec a -> PrimVec a
Vec.drop @Int) (KnownList es => Int
forall (es :: [Effect]). KnownList es => Int
reifyLen @es)

-- | Get the head of the stack. \( O(1) \).
head :: Stack (e : es) -> HandlerPtr e
head :: Stack (e : es) -> HandlerPtr e
head = (PrimVec Int -> Int) -> Stack (e : es) -> HandlerPtr e
coerce (Prim Int => PrimVec Int -> Int
forall a. Prim a => PrimVec a -> a
Vec.head @Int)

-- | Take elements from the top of the stack. \( O(m) \).
take :: ∀ es es'. KnownList es => Stack (es ++ es') -> Stack es
take :: Stack (es ++ es') -> Stack es
take = (Int -> PrimVec Int -> PrimVec Int)
-> Int -> Stack (es ++ es') -> Stack es
coerce (Prim Int => Int -> PrimVec Int -> PrimVec Int
forall a. Prim a => Int -> PrimVec a -> PrimVec a
Vec.take @Int) (KnownList es => Int
forall (es :: [Effect]). KnownList es => Int
reifyLen @es)

-- | @e ':>' es@ means the effect @e@ is present in the effect stack @es@, and therefore can be 'Cleff.send'ed in an
-- @'Cleff.Eff' es@ computation.
class (e :: Effect) :> (es :: [Effect]) where
  -- | Get the index of the element.
  reifyIndex :: Int
  reifyIndex = String -> String -> String -> Int
forall a. String -> String -> String -> a
unreifiable String
"Elem" String
"Cleff.Internal.Stack.reifyIndex" String
"the index of an effect in the effect stack"
infix 0 :>

-- | The element closer to the head takes priority.
instance {-# OVERLAPPING #-} e :> e : es where
  reifyIndex :: Int
reifyIndex = Int
0

instance e :> es => e :> e' : es where
  reifyIndex :: Int
reifyIndex = Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ (e :> es) => Int
forall (e :: Effect) (es :: [Effect]). (e :> es) => Int
reifyIndex @e @es

type ElemNotFound e = Text "The element '" :<>: ShowType e :<>: Text "' is not present in the constraint"

instance TypeError (ElemNotFound e) => e :> '[] where
  reifyIndex :: Int
reifyIndex = String -> Int
forall a. HasCallStack => String -> a
error
    String
"Cleff.Internal.Stack.reifyIndex: Attempting to refer to a nonexistent member. Please report this as a bug."

-- | @xs ':>>' es@ means the list of effects @xs@ are all present in the effect stack @es@. This is a convenient type
-- alias for @(e1 ':>' es, ..., en ':>' es)@.
type family xs :>> es :: Constraint where
  '[] :>> _ = ()
  (x : xs) :>> es = (x :> es, xs :>> es)
infix 0 :>>

-- | Get an element in the stack. Amortized \( O(1) \).
index :: ∀ e es. e :> es => Stack es -> HandlerPtr e
index :: Stack es -> HandlerPtr e
index = (Int -> PrimVec Int -> Int) -> Int -> Stack es -> HandlerPtr e
coerce (Prim Int => Int -> PrimVec Int -> Int
forall a. Prim a => Int -> PrimVec a -> a
Vec.index @Int) ((e :> es) => Int
forall (e :: Effect) (es :: [Effect]). (e :> es) => Int
reifyIndex @e @es)

-- | @es@ is a subset of @es'@, /i.e./ all elements of @es@ are in @es'@.
class KnownList es => Subset (es :: [Effect]) (es' :: [Effect]) where
  -- | Get a list of indices of the elements.
  reifyIndices :: [Int]
  reifyIndices = String -> String -> String -> [Int]
forall a. String -> String -> String -> a
unreifiable
    String
"Subset" String
"Cleff.Internal.Stack.reifyIndices" String
"the indices of a subset of the effect stack"

instance Subset '[] es where
  reifyIndices :: [Int]
reifyIndices = []

instance (Subset es es', e :> es') => Subset (e : es) es' where
  reifyIndices :: [Int]
reifyIndices = (e :> es') => Int
forall (e :: Effect) (es :: [Effect]). (e :> es) => Int
reifyIndex @e @es' Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: Subset es es' => [Int]
forall (es :: [Effect]) (es' :: [Effect]). Subset es es' => [Int]
reifyIndices @es @es'

-- | Get a subset of the stack. Amortized \( O(m) \).
pick :: ∀ es es'. Subset es es' => Stack es' -> Stack es
pick :: Stack es' -> Stack es
pick = (Int -> [Int] -> PrimVec Int -> PrimVec Int)
-> Int -> [Int] -> Stack es' -> Stack es
coerce (Prim Int => Int -> [Int] -> PrimVec Int -> PrimVec Int
forall a. Prim a => Int -> [Int] -> PrimVec a -> PrimVec a
Vec.pick @Int) (KnownList es => Int
forall (es :: [Effect]). KnownList es => Int
reifyLen @es) (Subset es es' => [Int]
forall (es :: [Effect]) (es' :: [Effect]). Subset es es' => [Int]
reifyIndices @es @es')

-- | Update an entry in the stack. \( O(n) \).
update :: ∀ e es. e :> es => HandlerPtr e -> Stack es -> Stack es
update :: HandlerPtr e -> Stack es -> Stack es
update = (Int -> Int -> PrimVec Int -> PrimVec Int)
-> Int -> HandlerPtr e -> Stack es -> Stack es
coerce (Prim Int => Int -> Int -> PrimVec Int -> PrimVec Int
forall a. Prim a => Int -> a -> PrimVec a -> PrimVec a
Vec.update @Int) ((e :> es) => Int
forall (e :: Effect) (es :: [Effect]). (e :> es) => Int
reifyIndex @e @es)