{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE CPP #-}

-- SPDX-License-Identifier: MPL-2.0

{- |
Copyright   :  (c) 2024 Sayo Koyoneda
License     :  MPL-2.0 (see the file LICENSE)
Maintainer  :  ymdfield@outlook.jp

Effects for parallel computations.
-}
module Data.Effect.Concurrent.Parallel where

#if ( __GLASGOW_HASKELL__ < 906 )
import Control.Applicative (liftA2)
#endif
import Control.Applicative (Alternative (empty, (<|>)))
import Data.Tuple (swap)

-- | An `Applicative`-based effect for executing computations in parallel.
data Parallel f a where
    -- | Executes two actions in parallel and blocks until both are complete.
    -- Finally, aggregates the execution results based on the specified function.
    LiftP2
        :: (a -> b -> c)
        -- ^ A function that aggregates the two execution results.
        -> f a
        -- ^ The first action to be executed in parallel.
        -> f b
        -- ^ The second action to be executed in parallel.
        -> Parallel f c

-- | An effect that blocks a computation indefinitely.
data Halt (a :: Type) where
    -- | Blocks a computation indefinitely.
    Halt :: Halt a

{- |
An effect that adopts the result of the computation that finishes first among
two computations and cancels the other.
-}
data Race f (a :: Type) where
    -- | Adopts the result of the computation that finishes first among two
    --   computations and cancels the other.
    Race :: f a -> f a -> Race f a

makeEffect [''Halt] [''Parallel, ''Race]

{- |
A wrapper that allows using the `Parallel` effect in the form of `Applicative` /
 `Alternative` instances.
-}
newtype Concurrently f a = Concurrently {forall {k} (f :: k -> *) (a :: k). Concurrently f a -> f a
runConcurrently :: f a}
    deriving ((forall a b. (a -> b) -> Concurrently f a -> Concurrently f b)
-> (forall a b. a -> Concurrently f b -> Concurrently f a)
-> Functor (Concurrently f)
forall a b. a -> Concurrently f b -> Concurrently f a
forall a b. (a -> b) -> Concurrently f a -> Concurrently f b
forall (f :: * -> *) a b.
Functor f =>
a -> Concurrently f b -> Concurrently f a
forall (f :: * -> *) a b.
Functor f =>
(a -> b) -> Concurrently f a -> Concurrently f b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall (f :: * -> *) a b.
Functor f =>
(a -> b) -> Concurrently f a -> Concurrently f b
fmap :: forall a b. (a -> b) -> Concurrently f a -> Concurrently f b
$c<$ :: forall (f :: * -> *) a b.
Functor f =>
a -> Concurrently f b -> Concurrently f a
<$ :: forall a b. a -> Concurrently f b -> Concurrently f a
Functor)

instance (Parallel <<: f, Applicative f) => Applicative (Concurrently f) where
    pure :: forall a. a -> Concurrently f a
pure = f a -> Concurrently f a
forall {k} (f :: k -> *) (a :: k). f a -> Concurrently f a
Concurrently (f a -> Concurrently f a) -> (a -> f a) -> a -> Concurrently f a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> f a
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
    {-# INLINE pure #-}

    liftA2 :: forall a b c.
(a -> b -> c)
-> Concurrently f a -> Concurrently f b -> Concurrently f c
liftA2 a -> b -> c
f (Concurrently f a
a) (Concurrently f b
b) = f c -> Concurrently f c
forall {k} (f :: k -> *) (a :: k). f a -> Concurrently f a
Concurrently (f c -> Concurrently f c) -> f c -> Concurrently f c
forall a b. (a -> b) -> a -> b
$ (a -> b -> c) -> f a -> f b -> f c
forall a b c (f :: * -> *).
SendHOE Parallel f =>
(a -> b -> c) -> f a -> f b -> f c
liftP2 a -> b -> c
f f a
a f b
b
    {-# INLINE liftA2 #-}

instance (Race <<: f, Halt <: f, Parallel <<: f, Applicative f) => Alternative (Concurrently f) where
    empty :: forall a. Concurrently f a
empty = f a -> Concurrently f a
forall {k} (f :: k -> *) (a :: k). f a -> Concurrently f a
Concurrently f a
forall a (f :: * -> *). SendFOE Halt f => f a
halt
    {-# INLINE empty #-}

    (Concurrently f a
a) <|> :: forall a. Concurrently f a -> Concurrently f a -> Concurrently f a
<|> (Concurrently f a
b) = f a -> Concurrently f a
forall {k} (f :: k -> *) (a :: k). f a -> Concurrently f a
Concurrently (f a -> Concurrently f a) -> f a -> Concurrently f a
forall a b. (a -> b) -> a -> b
$ f a -> f a -> f a
forall a (f :: * -> *). SendHOE Race f => f a -> f a -> f a
race f a
a f a
b
    {-# INLINE (<|>) #-}

{- |
Executes three actions in parallel and blocks until all are complete.
Finally, aggregates the execution results based on the specified function.
-}
liftP3
    :: (Parallel <<: f, Applicative f)
    => (a -> b -> c -> d)
    -- ^ A function that aggregates the three execution results.
    -> f a
    -- ^ The first action to be executed in parallel.
    -> f b
    -- ^ The second action to be executed in parallel.
    -> f c
    -- ^ The third action to be executed in parallel.
    -> f d
liftP3 :: forall (f :: * -> *) a b c d.
(Parallel <<: f, Applicative f) =>
(a -> b -> c -> d) -> f a -> f b -> f c -> f d
liftP3 a -> b -> c -> d
f f a
a f b
b = ((c -> d) -> c -> d) -> f (c -> d) -> f c -> f d
forall a b c (f :: * -> *).
SendHOE Parallel f =>
(a -> b -> c) -> f a -> f b -> f c
liftP2 (c -> d) -> c -> d
forall a b. (a -> b) -> a -> b
($) ((a -> b -> c -> d) -> f a -> f b -> f (c -> d)
forall a b c (f :: * -> *).
SendHOE Parallel f =>
(a -> b -> c) -> f a -> f b -> f c
liftP2 a -> b -> c -> d
f f a
a f b
b)
{-# INLINE liftP3 #-}

-- | An effect that realizes polling and cancellation of actions running in parallel.
data Poll f a where
    -- | Performs polling on an action running in parallel in the form of a fold.
    --
    -- First, the parallel execution of two actions begins.
    --
    -- When the execution of the first action completes, polling on the second
    -- action is performed at that point, and the result is passed to the
    -- folding function. If the function returns `Left`, the folding terminates
    -- and it becomes the final result. If the second action is not yet
    -- complete, it is canceled. If the function returns `Right`, the folding
    -- continues, and the same process repeats.
    Poldl
        :: (a -> Maybe b -> f (Either r a))
        -- ^ A function for folding.
        -> f a
        -- ^ The first action to be executed in parallel.
        -> f b
        -- ^ The second action to be executed in parallel; the target of polling.
        -> Poll f r

makeEffectH [''Poll]

-- | Executes two actions in parallel. If the first action completes before the second, the second action is canceled.
cancels
    :: (Poll <<: f, Applicative f)
    => f a
    -- ^ The action that controls the cancellation.
    -> f b
    -- ^ The action to be canceled.
    -> f (a, Maybe b)
cancels :: forall (f :: * -> *) a b.
(Poll <<: f, Applicative f) =>
f a -> f b -> f (a, Maybe b)
cancels = (a -> Maybe b -> f (Either (a, Maybe b) a))
-> f a -> f b -> f (a, Maybe b)
forall a b r (f :: * -> *).
SendHOE Poll f =>
(a -> Maybe b -> f (Either r a)) -> f a -> f b -> f r
poldl ((a -> Maybe b -> f (Either (a, Maybe b) a))
 -> f a -> f b -> f (a, Maybe b))
-> (a -> Maybe b -> f (Either (a, Maybe b) a))
-> f a
-> f b
-> f (a, Maybe b)
forall a b. (a -> b) -> a -> b
$ ((a, Maybe b) -> f (Either (a, Maybe b) a))
-> a -> Maybe b -> f (Either (a, Maybe b) a)
forall a b c. ((a, b) -> c) -> a -> b -> c
curry (((a, Maybe b) -> f (Either (a, Maybe b) a))
 -> a -> Maybe b -> f (Either (a, Maybe b) a))
-> ((a, Maybe b) -> f (Either (a, Maybe b) a))
-> a
-> Maybe b
-> f (Either (a, Maybe b) a)
forall a b. (a -> b) -> a -> b
$ Either (a, Maybe b) a -> f (Either (a, Maybe b) a)
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either (a, Maybe b) a -> f (Either (a, Maybe b) a))
-> ((a, Maybe b) -> Either (a, Maybe b) a)
-> (a, Maybe b)
-> f (Either (a, Maybe b) a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a, Maybe b) -> Either (a, Maybe b) a
forall a b. a -> Either a b
Left
{-# INLINE cancels #-}

-- | Executes two actions in parallel. If the second action completes before the first, the first action is canceled.
cancelBy
    :: (Poll <<: f, Applicative f)
    => f a
    -- ^ The action to be canceled.
    -> f b
    -- ^ The action that controls the cancellation.
    -> f (Maybe a, b)
cancelBy :: forall (f :: * -> *) a b.
(Poll <<: f, Applicative f) =>
f a -> f b -> f (Maybe a, b)
cancelBy = (f b -> f a -> f (Maybe a, b)) -> f a -> f b -> f (Maybe a, b)
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((f b -> f a -> f (Maybe a, b)) -> f a -> f b -> f (Maybe a, b))
-> (f b -> f a -> f (Maybe a, b)) -> f a -> f b -> f (Maybe a, b)
forall a b. (a -> b) -> a -> b
$ (b -> Maybe a -> f (Either (Maybe a, b) b))
-> f b -> f a -> f (Maybe a, b)
forall a b r (f :: * -> *).
SendHOE Poll f =>
(a -> Maybe b -> f (Either r a)) -> f a -> f b -> f r
poldl ((b -> Maybe a -> f (Either (Maybe a, b) b))
 -> f b -> f a -> f (Maybe a, b))
-> (b -> Maybe a -> f (Either (Maybe a, b) b))
-> f b
-> f a
-> f (Maybe a, b)
forall a b. (a -> b) -> a -> b
$ ((b, Maybe a) -> f (Either (Maybe a, b) b))
-> b -> Maybe a -> f (Either (Maybe a, b) b)
forall a b c. ((a, b) -> c) -> a -> b -> c
curry (((b, Maybe a) -> f (Either (Maybe a, b) b))
 -> b -> Maybe a -> f (Either (Maybe a, b) b))
-> ((b, Maybe a) -> f (Either (Maybe a, b) b))
-> b
-> Maybe a
-> f (Either (Maybe a, b) b)
forall a b. (a -> b) -> a -> b
$ Either (Maybe a, b) b -> f (Either (Maybe a, b) b)
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either (Maybe a, b) b -> f (Either (Maybe a, b) b))
-> ((b, Maybe a) -> Either (Maybe a, b) b)
-> (b, Maybe a)
-> f (Either (Maybe a, b) b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Maybe a, b) -> Either (Maybe a, b) b
forall a b. a -> Either a b
Left ((Maybe a, b) -> Either (Maybe a, b) b)
-> ((b, Maybe a) -> (Maybe a, b))
-> (b, Maybe a)
-> Either (Maybe a, b) b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (b, Maybe a) -> (Maybe a, b)
forall a b. (a, b) -> (b, a)
swap
{-# INLINE cancelBy #-}

-- | An effect for parallel computations based on a `Traversable` container @t@.
data For (t :: Type -> Type) f a where
    -- | Executes in parallel the actions stored within a `Traversable` container @t@.
    For :: t (f a) -> For t f (t a)

makeEffectH_ [''For]
makeHFunctor' ''For \(t :< _) -> [t|Functor $t|]

-- | Converts the `Traversable` container-based parallel computation effect t`For` into the `Applicative`-based parallel computation effect `Parallel`.
forToParallel :: (Parallel <<: f, Traversable t, Applicative f) => For t f ~> f
forToParallel :: forall (f :: * -> *) (t :: * -> *).
(Parallel <<: f, Traversable t, Applicative f) =>
For t f ~> f
forToParallel (For t (f a)
iters) = Concurrently f x -> f x
forall {k} (f :: k -> *) (a :: k). Concurrently f a -> f a
runConcurrently (Concurrently f x -> f x) -> Concurrently f x -> f x
forall a b. (a -> b) -> a -> b
$ (f a -> Concurrently f a) -> t (f a) -> Concurrently f (t a)
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) -> t a -> f (t b)
traverse f a -> Concurrently f a
forall {k} (f :: k -> *) (a :: k). f a -> Concurrently f a
Concurrently t (f a)
iters
{-# INLINE forToParallel #-}