-- | Support for accumulation of values in a shared array.
--
-- @since 2.7.0.0
module Effectful.Output.Static.Shared.Array
  ( -- * Effect
    Output

    -- ** Handlers
  , runOutput

    -- ** Operations
  , output

    -- * Re-exports
  , Array
  ) where

import Control.Concurrent.MVar.Strict qualified as S
import Control.Monad.Primitive
import Data.Kind
import Data.Primitive.Array

import Effectful
import Effectful.Dispatch.Static
import Effectful.Dispatch.Static.Primitive
import Effectful.Internal.Utils

-- | Provide access to accumulation of values of type @o@ in a shared array.
data Output (o :: Type) :: Effect

data OutputData o = OutputData !Int !(MutableArray RealWorld o)

type instance DispatchOf (Output o) = Static NoSideEffects
newtype instance StaticRep (Output o) = Output (S.MVar (OutputData o))

-- | Run the 'Output' effect and return the final value along with the
-- accumulated array.
runOutput :: HasCallStack => Eff (Output o : es) a -> Eff es (a, Array o)
runOutput :: forall o (es :: [(Type -> Type) -> Type -> Type]) a.
HasCallStack =>
Eff (Output o : es) a -> Eff es (a, Array o)
runOutput = (OutputData o -> IO (Array o))
-> Eff (Output o : es) a -> Eff es (a, Array o)
forall o acc (es :: [(Type -> Type) -> Type -> Type]) a.
HasCallStack =>
(OutputData o -> IO acc)
-> Eff (Output o : es) a -> Eff es (a, acc)
runOutputImpl ((OutputData o -> IO (Array o))
 -> Eff (Output o : es) a -> Eff es (a, Array o))
-> (OutputData o -> IO (Array o))
-> Eff (Output o : es) a
-> Eff es (a, Array o)
forall a b. (a -> b) -> a -> b
$ \(OutputData Int
size MutableArray RealWorld o
arr) -> do
  MutableArray (PrimState IO) o -> Int -> Int -> IO (Array o)
forall (m :: Type -> Type) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> Int -> m (Array a)
freezeArray MutableArray RealWorld o
MutableArray (PrimState IO) o
arr Int
0 Int
size

-- | Append the value to the end of the array.
output
  :: (HasCallStack, Output o :> es)
  => o -- ^ The value.
  -> Eff es ()
output :: forall o (es :: [(Type -> Type) -> Type -> Type]).
(HasCallStack, Output o :> es) =>
o -> Eff es ()
output !o
o = (Env es -> IO ()) -> Eff es ()
forall (es :: [(Type -> Type) -> Type -> Type]) a.
(Env es -> IO a) -> Eff es a
unsafeEff ((Env es -> IO ()) -> Eff es ()) -> (Env es -> IO ()) -> Eff es ()
forall a b. (a -> b) -> a -> b
$ \Env es
es -> do
  Output MVar (OutputData o)
v <- Env es -> IO (EffectRep (DispatchOf (Output o)) (Output o))
forall (e :: (Type -> Type) -> Type -> Type)
       (es :: [(Type -> Type) -> Type -> Type]).
(HasCallStack, e :> es) =>
Env es -> IO (EffectRep (DispatchOf e) e)
getEnv Env es
es
  MVar (OutputData o) -> (OutputData o -> IO (OutputData o)) -> IO ()
forall a. MVar a -> (a -> IO a) -> IO ()
S.modifyMVar_ MVar (OutputData o)
v ((OutputData o -> IO (OutputData o)) -> IO ())
-> (OutputData o -> IO (OutputData o)) -> IO ()
forall a b. (a -> b) -> a -> b
$ \(OutputData Int
size MutableArray RealWorld o
arr0) -> do
    let len0 :: Int
len0 = MutableArray RealWorld o -> Int
forall s a. MutableArray s a -> Int
sizeofMutableArray MutableArray RealWorld o
arr0
    MutableArray RealWorld o
arr <- case Int
size Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
`compare` Int
len0 of
      Ordering
GT -> [Char] -> IO (MutableArray RealWorld o)
forall a. HasCallStack => [Char] -> a
error ([Char] -> IO (MutableArray RealWorld o))
-> [Char] -> IO (MutableArray RealWorld o)
forall a b. (a -> b) -> a -> b
$ [Char]
"size (" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
size [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
") > len0 (" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
len0 [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
")"
      Ordering
LT -> MutableArray RealWorld o -> IO (MutableArray RealWorld o)
forall a. a -> IO a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure MutableArray RealWorld o
arr0
      Ordering
EQ -> do
        let len :: Int
len = Int -> Int
growCapacity Int
len0
        MutableArray RealWorld o
arr <- Int -> o -> IO (MutableArray (PrimState IO) o)
forall (m :: Type -> Type) a.
PrimMonad m =>
Int -> a -> m (MutableArray (PrimState m) a)
newArray Int
len o
forall a. HasCallStack => a
undefinedValue
        MutableArray (PrimState IO) o
-> Int -> MutableArray (PrimState IO) o -> Int -> Int -> IO ()
forall (m :: Type -> Type) a.
PrimMonad m =>
MutableArray (PrimState m) a
-> Int -> MutableArray (PrimState m) a -> Int -> Int -> m ()
copyMutableArray MutableArray RealWorld o
MutableArray (PrimState IO) o
arr Int
0 MutableArray RealWorld o
MutableArray (PrimState IO) o
arr0 Int
0 Int
size
        MutableArray RealWorld o -> IO (MutableArray RealWorld o)
forall a. a -> IO a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure MutableArray RealWorld o
arr
    MutableArray (PrimState IO) o -> Int -> o -> IO ()
forall (m :: Type -> Type) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> a -> m ()
writeArray MutableArray RealWorld o
MutableArray (PrimState IO) o
arr Int
size o
o
    OutputData o -> IO (OutputData o)
forall a. a -> IO a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (OutputData o -> IO (OutputData o))
-> OutputData o -> IO (OutputData o)
forall a b. (a -> b) -> a -> b
$ Int -> MutableArray RealWorld o -> OutputData o
forall o. Int -> MutableArray RealWorld o -> OutputData o
OutputData (Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) MutableArray RealWorld o
arr

----------------------------------------
-- Helpers

runOutputImpl
  :: HasCallStack
  => (OutputData o -> IO acc)
  -> Eff (Output o : es) a
  -> Eff es (a, acc)
runOutputImpl :: forall o acc (es :: [(Type -> Type) -> Type -> Type]) a.
HasCallStack =>
(OutputData o -> IO acc)
-> Eff (Output o : es) a -> Eff es (a, acc)
runOutputImpl OutputData o -> IO acc
f Eff (Output o : es) a
action = do
  MVar (OutputData o)
v <- IO (MVar (OutputData o)) -> Eff es (MVar (OutputData o))
forall a (es :: [(Type -> Type) -> Type -> Type]). IO a -> Eff es a
unsafeEff_ (IO (MVar (OutputData o)) -> Eff es (MVar (OutputData o)))
-> IO (MVar (OutputData o)) -> Eff es (MVar (OutputData o))
forall a b. (a -> b) -> a -> b
$ OutputData o -> IO (MVar (OutputData o))
forall a. a -> IO (MVar a)
S.newMVar (OutputData o -> IO (MVar (OutputData o)))
-> (MutableArray RealWorld o -> OutputData o)
-> MutableArray RealWorld o
-> IO (MVar (OutputData o))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> MutableArray RealWorld o -> OutputData o
forall o. Int -> MutableArray RealWorld o -> OutputData o
OutputData Int
0 (MutableArray RealWorld o -> IO (MVar (OutputData o)))
-> IO (MutableArray RealWorld o) -> IO (MVar (OutputData o))
forall (m :: Type -> Type) a b. Monad m => (a -> m b) -> m a -> m b
=<< Int -> o -> IO (MutableArray (PrimState IO) o)
forall (m :: Type -> Type) a.
PrimMonad m =>
Int -> a -> m (MutableArray (PrimState m) a)
newArray Int
0 o
forall a. HasCallStack => a
undefinedValue
  a
a <- StaticRep (Output o) -> Eff (Output o : es) a -> Eff es a
forall (e :: (Type -> Type) -> Type -> Type)
       (sideEffects :: SideEffects)
       (es :: [(Type -> Type) -> Type -> Type]) a.
(HasCallStack, DispatchOf e ~ 'Static sideEffects,
 MaybeIOE sideEffects es) =>
StaticRep e -> Eff (e : es) a -> Eff es a
evalStaticRep (MVar (OutputData o) -> StaticRep (Output o)
forall o. MVar (OutputData o) -> StaticRep (Output o)
Output MVar (OutputData o)
v) Eff (Output o : es) a
action
  acc
acc <- IO acc -> Eff es acc
forall a (es :: [(Type -> Type) -> Type -> Type]). IO a -> Eff es a
unsafeEff_ (IO acc -> Eff es acc) -> IO acc -> Eff es acc
forall a b. (a -> b) -> a -> b
$ OutputData o -> IO acc
f (OutputData o -> IO acc) -> IO (OutputData o) -> IO acc
forall (m :: Type -> Type) a b. Monad m => (a -> m b) -> m a -> m b
=<< MVar (OutputData o) -> IO (OutputData o)
forall a. MVar a -> IO a
S.readMVar MVar (OutputData o)
v
  (a, acc) -> Eff es (a, acc)
forall a. a -> Eff es a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (a
a, acc
acc)

undefinedValue :: HasCallStack => a
undefinedValue :: forall a. HasCallStack => a
undefinedValue = [Char] -> a
forall a. HasCallStack => [Char] -> a
error [Char]
"Undefined value"