{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE Trustworthy #-}

-- |
-- Module      :   Grisette.Lib.Control.Traversable
-- Copyright   :   (c) Sirui Lu 2021-2023
-- License     :   BSD-3-Clause (see the LICENSE file)
--
-- Maintainer  :   siruilu@cs.washington.edu
-- Stability   :   Experimental
-- Portability :   GHC only
module Grisette.Lib.Data.Traversable
  ( -- * The 'Traversable' class
    mrgTraverse,
    mrgSequenceA,
    mrgMapM,
    mrgSequence,

    -- * Utility functions
    mrgFor,
    mrgForM,
    mrgMapAccumM,
    mrgForAccumM,
  )
where

import Control.Monad.State (StateT (StateT, runStateT))
import Grisette.Internal.Core.Data.Class.Mergeable
  ( Mergeable (rootStrategy),
    Mergeable1,
    Mergeable2 (liftRootStrategy2),
    rootStrategy1,
  )
import Grisette.Internal.Core.Data.Class.TryMerge
  ( MonadTryMerge,
    TryMerge (tryMergeWithStrategy),
    tryMerge,
  )
import Grisette.Lib.Control.Applicative (mrgPure)

-- | 'Data.Traversable.traverse' with 'Grisette.Core.MergingStrategy' knowledge
-- propagation.
mrgTraverse ::
  forall a b t f.
  ( Mergeable b,
    Mergeable1 t,
    TryMerge f,
    Applicative f,
    Traversable t
  ) =>
  (a -> f b) ->
  t a ->
  f (t b)
mrgTraverse :: forall a b (t :: * -> *) (f :: * -> *).
(Mergeable b, Mergeable1 t, TryMerge f, Applicative f,
 Traversable t) =>
(a -> f b) -> t a -> f (t b)
mrgTraverse a -> f b
f = MergingStrategy (t b) -> f (t b) -> f (t b)
forall a. MergingStrategy a -> f a -> f a
forall (m :: * -> *) a.
TryMerge m =>
MergingStrategy a -> m a -> m a
tryMergeWithStrategy MergingStrategy (t b)
forall a (u :: * -> *).
(Mergeable a, Mergeable1 u) =>
MergingStrategy (u a)
rootStrategy1 (f (t b) -> f (t b)) -> (t a -> f (t b)) -> t a -> f (t b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> f b) -> t a -> f (t 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) -> t a -> f (t b)
traverse (f b -> f b
forall (m :: * -> *) a. (TryMerge m, Mergeable a) => m a -> m a
tryMerge (f b -> f b) -> (a -> f b) -> a -> f b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> f b
f)
{-# INLINE mrgTraverse #-}

-- | 'Data.Traversable.sequenceA' with 'Grisette.Core.MergingStrategy' knowledge
-- propagation.
mrgSequenceA ::
  forall a t f.
  ( Mergeable a,
    Mergeable1 t,
    Applicative f,
    TryMerge f,
    Traversable t
  ) =>
  t (f a) ->
  f (t a)
mrgSequenceA :: forall a (t :: * -> *) (f :: * -> *).
(Mergeable a, Mergeable1 t, Applicative f, TryMerge f,
 Traversable t) =>
t (f a) -> f (t a)
mrgSequenceA = (f a -> f a) -> t (f a) -> f (t a)
forall a b (t :: * -> *) (f :: * -> *).
(Mergeable b, Mergeable1 t, TryMerge f, Applicative f,
 Traversable t) =>
(a -> f b) -> t a -> f (t b)
mrgTraverse f a -> f a
forall a. a -> a
id
{-# INLINE mrgSequenceA #-}

-- | 'Data.Traversable.mapM' with 'Grisette.Core.MergingStrategy' knowledge
-- propagation.
mrgMapM ::
  forall a b t f.
  ( Mergeable b,
    Mergeable1 t,
    MonadTryMerge f,
    Traversable t
  ) =>
  (a -> f b) ->
  t a ->
  f (t b)
mrgMapM :: forall a b (t :: * -> *) (f :: * -> *).
(Mergeable b, Mergeable1 t, MonadTryMerge f, Traversable t) =>
(a -> f b) -> t a -> f (t b)
mrgMapM = (a -> f b) -> t a -> f (t b)
forall a b (t :: * -> *) (f :: * -> *).
(Mergeable b, Mergeable1 t, TryMerge f, Applicative f,
 Traversable t) =>
(a -> f b) -> t a -> f (t b)
mrgTraverse
{-# INLINE mrgMapM #-}

-- | 'Data.Traversable.sequence' with 'Grisette.Core.MergingStrategy' knowledge
-- propagation.
mrgSequence ::
  forall a t f.
  ( Mergeable a,
    Mergeable1 t,
    MonadTryMerge f,
    Traversable t
  ) =>
  t (f a) ->
  f (t a)
mrgSequence :: forall a (t :: * -> *) (f :: * -> *).
(Mergeable a, Mergeable1 t, MonadTryMerge f, Traversable t) =>
t (f a) -> f (t a)
mrgSequence = t (f a) -> f (t a)
forall a (t :: * -> *) (f :: * -> *).
(Mergeable a, Mergeable1 t, Applicative f, TryMerge f,
 Traversable t) =>
t (f a) -> f (t a)
mrgSequenceA
{-# INLINE mrgSequence #-}

-- | 'Data.Traversable.for' with 'Grisette.Core.MergingStrategy' knowledge
-- propagation.
mrgFor ::
  ( Mergeable b,
    Mergeable1 t,
    Traversable t,
    TryMerge m,
    Applicative m
  ) =>
  t a ->
  (a -> m b) ->
  m (t b)
mrgFor :: forall b (t :: * -> *) (m :: * -> *) a.
(Mergeable b, Mergeable1 t, Traversable t, TryMerge m,
 Applicative m) =>
t a -> (a -> m b) -> m (t b)
mrgFor = ((a -> m b) -> t a -> m (t b)) -> t a -> (a -> m b) -> m (t b)
forall a b c. (a -> b -> c) -> b -> a -> c
flip (a -> m b) -> t a -> m (t b)
forall a b (t :: * -> *) (f :: * -> *).
(Mergeable b, Mergeable1 t, TryMerge f, Applicative f,
 Traversable t) =>
(a -> f b) -> t a -> f (t b)
mrgTraverse
{-# INLINE mrgFor #-}

-- | 'Data.Traversable.forM' with 'Grisette.Core.MergingStrategy' knowledge
-- propagation.
mrgForM ::
  ( Mergeable b,
    Mergeable1 t,
    Traversable t,
    MonadTryMerge m
  ) =>
  t a ->
  (a -> m b) ->
  m (t b)
mrgForM :: forall b (t :: * -> *) (m :: * -> *) a.
(Mergeable b, Mergeable1 t, Traversable t, MonadTryMerge m) =>
t a -> (a -> m b) -> m (t b)
mrgForM = ((a -> m b) -> t a -> m (t b)) -> t a -> (a -> m b) -> m (t b)
forall a b c. (a -> b -> c) -> b -> a -> c
flip (a -> m b) -> t a -> m (t b)
forall a b (t :: * -> *) (f :: * -> *).
(Mergeable b, Mergeable1 t, MonadTryMerge f, Traversable t) =>
(a -> f b) -> t a -> f (t b)
mrgMapM
{-# INLINE mrgForM #-}

-- | 'Data.Traversable.mapAccumM' with 'Grisette.Core.MergingStrategy' knowledge
-- propagation.
mrgMapAccumM ::
  (MonadTryMerge m, Traversable t, Mergeable s, Mergeable b, Mergeable1 t) =>
  (s -> a -> m (s, b)) ->
  s ->
  t a ->
  m (s, t b)
mrgMapAccumM :: forall (m :: * -> *) (t :: * -> *) s b a.
(MonadTryMerge m, Traversable t, Mergeable s, Mergeable b,
 Mergeable1 t) =>
(s -> a -> m (s, b)) -> s -> t a -> m (s, t b)
mrgMapAccumM s -> a -> m (s, b)
f s
s t a
t =
  MergingStrategy (s, t b) -> m (s, t b) -> m (s, t b)
forall a. MergingStrategy a -> m a -> m a
forall (m :: * -> *) a.
TryMerge m =>
MergingStrategy a -> m a -> m a
tryMergeWithStrategy (MergingStrategy s
-> MergingStrategy (t b) -> MergingStrategy (s, t b)
forall a b.
MergingStrategy a -> MergingStrategy b -> MergingStrategy (a, b)
forall (u :: * -> * -> *) a b.
Mergeable2 u =>
MergingStrategy a -> MergingStrategy b -> MergingStrategy (u a b)
liftRootStrategy2 MergingStrategy s
forall a. Mergeable a => MergingStrategy a
rootStrategy MergingStrategy (t b)
forall a (u :: * -> *).
(Mergeable a, Mergeable1 u) =>
MergingStrategy (u a)
rootStrategy1) (m (s, t b) -> m (s, t b)) -> m (s, t b) -> m (s, t b)
forall a b. (a -> b) -> a -> b
$ do
    (tb, s) <- (StateT s m (t b) -> s -> m (t b, s))
-> s -> StateT s m (t b) -> m (t b, s)
forall a b c. (a -> b -> c) -> b -> a -> c
flip StateT s m (t b) -> s -> m (t b, s)
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT s
s (StateT s m (t b) -> m (t b, s)) -> StateT s m (t b) -> m (t b, s)
forall a b. (a -> b) -> a -> b
$ do
      (a -> StateT s m b) -> t a -> StateT s m (t b)
forall a b (t :: * -> *) (f :: * -> *).
(Mergeable b, Mergeable1 t, MonadTryMerge f, Traversable t) =>
(a -> f b) -> t a -> f (t b)
mrgMapM
        ( \a
a -> (s -> m (b, s)) -> StateT s m b
forall s (m :: * -> *) a. (s -> m (a, s)) -> StateT s m a
StateT ((s -> m (b, s)) -> StateT s m b)
-> (s -> m (b, s)) -> StateT s m b
forall a b. (a -> b) -> a -> b
$ \s
s -> do
            (sr, br) <- s -> a -> m (s, b)
f s
s a
a
            mrgPure (br, sr)
        )
        t a
t
    return (s, tb)
{-# INLINE mrgMapAccumM #-}

-- | 'Data.Traversable.forAccumM' and 'Grisette.Core.MergingStrategy' knowledge
-- propagation.
mrgForAccumM ::
  (MonadTryMerge m, Traversable t, Mergeable s, Mergeable b, Mergeable1 t) =>
  s ->
  t a ->
  (s -> a -> m (s, b)) ->
  m (s, t b)
mrgForAccumM :: forall (m :: * -> *) (t :: * -> *) s b a.
(MonadTryMerge m, Traversable t, Mergeable s, Mergeable b,
 Mergeable1 t) =>
s -> t a -> (s -> a -> m (s, b)) -> m (s, t b)
mrgForAccumM s
s t a
t s -> a -> m (s, b)
f = (s -> a -> m (s, b)) -> s -> t a -> m (s, t b)
forall (m :: * -> *) (t :: * -> *) s b a.
(MonadTryMerge m, Traversable t, Mergeable s, Mergeable b,
 Mergeable1 t) =>
(s -> a -> m (s, b)) -> s -> t a -> m (s, t b)
mrgMapAccumM s -> a -> m (s, b)
f s
s t a
t
{-# INLINE mrgForAccumM #-}