{-# LANGUAGE CPP #-}
{-# LANGUAGE ConstraintKinds #-}

-- |
-- Module      :   Grisette.Internal.Core.Data.Class.TryMerge
-- Copyright   :   (c) Sirui Lu 2023-2024
-- License     :   BSD-3-Clause (see the LICENSE file)
--
-- Maintainer  :   siruilu@cs.washington.edu
-- Stability   :   Experimental
-- Portability :   GHC only
module Grisette.Internal.Core.Data.Class.TryMerge
  ( TryMerge (..),
    tryMerge,
    MonadTryMerge,
    mrgSingle,
    mrgSingleWithStrategy,
  )
where

import Control.Monad.Cont (ContT (ContT))
import Control.Monad.Except (ExceptT (ExceptT))
import Control.Monad.Identity
  ( Identity,
    IdentityT (IdentityT),
  )
import qualified Control.Monad.RWS.Lazy as RWSLazy
import qualified Control.Monad.RWS.Strict as RWSStrict
import Control.Monad.Reader (ReaderT (ReaderT))
import qualified Control.Monad.State.Lazy as StateLazy
import qualified Control.Monad.State.Strict as StateStrict
import Control.Monad.Trans.Maybe (MaybeT (MaybeT))
import qualified Control.Monad.Writer.Lazy as WriterLazy
import qualified Control.Monad.Writer.Strict as WriterStrict
import Data.Functor.Sum (Sum (InL, InR))
import qualified Data.Monoid as Monoid
import Grisette.Internal.Core.Data.Class.Mergeable
  ( Mergeable (rootStrategy),
    Mergeable1 (liftRootStrategy),
    Mergeable2 (liftRootStrategy2),
    Mergeable3 (liftRootStrategy3),
    MergingStrategy,
  )

-- $setup
-- >>> import Grisette.Core
-- >>> import Grisette.SymPrim

-- | A class for containers that may or may not be merged.
--
-- If the container is capable of multi-path execution, then the
-- `tryMergeWithStrategy` function should merge the paths according to the
-- supplied strategy.
--
-- If the container is not capable of multi-path execution, then the
-- `tryMergeWithStrategy` function should be equivalent to `id`.
--
-- Note that this will not necessarily do a recursive merge for the elements.
class TryMerge m where
  tryMergeWithStrategy :: MergingStrategy a -> m a -> m a

-- | Try to merge the container with the root strategy.
tryMerge :: (TryMerge m, Mergeable a) => m a -> m a
tryMerge :: forall (m :: * -> *) a. (TryMerge m, Mergeable a) => m a -> m a
tryMerge = MergingStrategy a -> m a -> m a
forall a. MergingStrategy a -> m a -> m a
forall (m :: * -> *) a.
TryMerge m =>
MergingStrategy a -> m a -> m a
tryMergeWithStrategy MergingStrategy a
forall a. Mergeable a => MergingStrategy a
rootStrategy
{-# INLINE tryMerge #-}

-- | Wrap a value in the applicative functor and capture the 'Mergeable'
-- knowledge.
--
-- >>> mrgSingleWithStrategy rootStrategy "a" :: UnionM SymInteger
-- {a}
--
-- __Note:__ Be careful to call this directly from your code.
-- The supplied merge strategy should be consistent with the type's root merge
-- strategy, or some internal invariants would be broken and the program can
-- crash.
--
-- This function is to be called when the 'Mergeable' constraint can not be
-- resolved, e.g., the merge strategy for the contained type is given with
-- 'Mergeable1'. In other cases, 'mrgPure' is usually a better alternative.
mrgSingleWithStrategy ::
  (TryMerge m, Applicative m) =>
  MergingStrategy a ->
  a ->
  m a
mrgSingleWithStrategy :: forall (m :: * -> *) a.
(TryMerge m, Applicative m) =>
MergingStrategy a -> a -> m a
mrgSingleWithStrategy MergingStrategy a
strategy = MergingStrategy a -> m a -> m a
forall a. MergingStrategy a -> m a -> m a
forall (m :: * -> *) a.
TryMerge m =>
MergingStrategy a -> m a -> m a
tryMergeWithStrategy MergingStrategy a
strategy (m a -> m a) -> (a -> m a) -> a -> m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> m a
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
{-# INLINE mrgSingleWithStrategy #-}

-- | Wrap a value in the applicative functor and propagate the type's root merge
-- strategy.
--
-- Equivalent to @'mrgSingleWithStrategy' 'rootStrategy'@.
--
-- >>> mrgSingle "a" :: UnionM SymInteger
-- {a}
mrgSingle :: (TryMerge m, Applicative m, Mergeable a) => a -> m a
mrgSingle :: forall (m :: * -> *) a.
(TryMerge m, Applicative m, Mergeable a) =>
a -> m a
mrgSingle = MergingStrategy a -> a -> m a
forall (m :: * -> *) a.
(TryMerge m, Applicative m) =>
MergingStrategy a -> a -> m a
mrgSingleWithStrategy MergingStrategy a
forall a. Mergeable a => MergingStrategy a
rootStrategy
{-# INLINE mrgSingle #-}

instance (TryMerge m) => TryMerge (MaybeT m) where
  tryMergeWithStrategy :: forall a. MergingStrategy a -> MaybeT m a -> MaybeT m a
tryMergeWithStrategy MergingStrategy a
strategy (MaybeT m (Maybe a)
ma) =
    m (Maybe a) -> MaybeT m a
forall (m :: * -> *) a. m (Maybe a) -> MaybeT m a
MaybeT (m (Maybe a) -> MaybeT m a) -> m (Maybe a) -> MaybeT m a
forall a b. (a -> b) -> a -> b
$ MergingStrategy (Maybe a) -> m (Maybe a) -> m (Maybe a)
forall a. MergingStrategy a -> m a -> m a
forall (m :: * -> *) a.
TryMerge m =>
MergingStrategy a -> m a -> m a
tryMergeWithStrategy (MergingStrategy a -> MergingStrategy (Maybe a)
forall a. MergingStrategy a -> MergingStrategy (Maybe a)
forall (u :: * -> *) a.
Mergeable1 u =>
MergingStrategy a -> MergingStrategy (u a)
liftRootStrategy MergingStrategy a
strategy) m (Maybe a)
ma
  {-# INLINE tryMergeWithStrategy #-}

instance (Mergeable e, TryMerge m) => TryMerge (ExceptT e m) where
  tryMergeWithStrategy :: forall a. MergingStrategy a -> ExceptT e m a -> ExceptT e m a
tryMergeWithStrategy MergingStrategy a
strategy (ExceptT m (Either e a)
ma) =
    m (Either e a) -> ExceptT e m a
forall e (m :: * -> *) a. m (Either e a) -> ExceptT e m a
ExceptT (m (Either e a) -> ExceptT e m a)
-> m (Either e a) -> ExceptT e m a
forall a b. (a -> b) -> a -> b
$ MergingStrategy (Either e a) -> m (Either e a) -> m (Either e a)
forall a. MergingStrategy a -> m a -> m a
forall (m :: * -> *) a.
TryMerge m =>
MergingStrategy a -> m a -> m a
tryMergeWithStrategy (MergingStrategy a -> MergingStrategy (Either e a)
forall a. MergingStrategy a -> MergingStrategy (Either e a)
forall (u :: * -> *) a.
Mergeable1 u =>
MergingStrategy a -> MergingStrategy (u a)
liftRootStrategy MergingStrategy a
strategy) m (Either e a)
ma
  {-# INLINE tryMergeWithStrategy #-}

instance (TryMerge m) => TryMerge (ReaderT r m) where
  tryMergeWithStrategy :: forall a. MergingStrategy a -> ReaderT r m a -> ReaderT r m a
tryMergeWithStrategy MergingStrategy a
strategy (ReaderT r -> m a
f) =
    (r -> m a) -> ReaderT r m a
forall r (m :: * -> *) a. (r -> m a) -> ReaderT r m a
ReaderT ((r -> m a) -> ReaderT r m a) -> (r -> m a) -> ReaderT r m a
forall a b. (a -> b) -> a -> b
$ \r
v -> MergingStrategy a -> m a -> m a
forall a. MergingStrategy a -> m a -> m a
forall (m :: * -> *) a.
TryMerge m =>
MergingStrategy a -> m a -> m a
tryMergeWithStrategy MergingStrategy a
strategy (m a -> m a) -> m a -> m a
forall a b. (a -> b) -> a -> b
$ r -> m a
f r
v
  {-# INLINE tryMergeWithStrategy #-}

instance (Mergeable s, TryMerge m) => TryMerge (StateLazy.StateT s m) where
  tryMergeWithStrategy :: forall a. MergingStrategy a -> StateT s m a -> StateT s m a
tryMergeWithStrategy MergingStrategy a
strategy (StateLazy.StateT s -> m (a, s)
f) =
    (s -> m (a, s)) -> StateT s m a
forall s (m :: * -> *) a. (s -> m (a, s)) -> StateT s m a
StateLazy.StateT ((s -> m (a, s)) -> StateT s m a)
-> (s -> m (a, s)) -> StateT s m a
forall a b. (a -> b) -> a -> b
$
      \s
s -> MergingStrategy (a, s) -> m (a, s) -> m (a, s)
forall a. MergingStrategy a -> m a -> m a
forall (m :: * -> *) a.
TryMerge m =>
MergingStrategy a -> m a -> m a
tryMergeWithStrategy (MergingStrategy a -> MergingStrategy s -> MergingStrategy (a, s)
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 a
strategy MergingStrategy s
forall a. Mergeable a => MergingStrategy a
rootStrategy) (s -> m (a, s)
f s
s)
  {-# INLINE tryMergeWithStrategy #-}

instance (Mergeable s, TryMerge m) => TryMerge (StateStrict.StateT s m) where
  tryMergeWithStrategy :: forall a. MergingStrategy a -> StateT s m a -> StateT s m a
tryMergeWithStrategy MergingStrategy a
strategy (StateStrict.StateT s -> m (a, s)
f) =
    (s -> m (a, s)) -> StateT s m a
forall s (m :: * -> *) a. (s -> m (a, s)) -> StateT s m a
StateStrict.StateT ((s -> m (a, s)) -> StateT s m a)
-> (s -> m (a, s)) -> StateT s m a
forall a b. (a -> b) -> a -> b
$
      \s
s -> MergingStrategy (a, s) -> m (a, s) -> m (a, s)
forall a. MergingStrategy a -> m a -> m a
forall (m :: * -> *) a.
TryMerge m =>
MergingStrategy a -> m a -> m a
tryMergeWithStrategy (MergingStrategy a -> MergingStrategy s -> MergingStrategy (a, s)
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 a
strategy MergingStrategy s
forall a. Mergeable a => MergingStrategy a
rootStrategy) (s -> m (a, s)
f s
s)
  {-# INLINE tryMergeWithStrategy #-}

instance
  (Monoid w, Mergeable w, TryMerge m) =>
  TryMerge (WriterLazy.WriterT w m)
  where
  tryMergeWithStrategy :: forall a. MergingStrategy a -> WriterT w m a -> WriterT w m a
tryMergeWithStrategy MergingStrategy a
strategy (WriterLazy.WriterT m (a, w)
f) =
    m (a, w) -> WriterT w m a
forall w (m :: * -> *) a. m (a, w) -> WriterT w m a
WriterLazy.WriterT (m (a, w) -> WriterT w m a) -> m (a, w) -> WriterT w m a
forall a b. (a -> b) -> a -> b
$
      MergingStrategy (a, w) -> m (a, w) -> m (a, w)
forall a. MergingStrategy a -> m a -> m a
forall (m :: * -> *) a.
TryMerge m =>
MergingStrategy a -> m a -> m a
tryMergeWithStrategy (MergingStrategy a -> MergingStrategy w -> MergingStrategy (a, w)
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 a
strategy MergingStrategy w
forall a. Mergeable a => MergingStrategy a
rootStrategy) m (a, w)
f
  {-# INLINE tryMergeWithStrategy #-}

instance
  (Monoid w, Mergeable w, TryMerge m) =>
  TryMerge (WriterStrict.WriterT w m)
  where
  tryMergeWithStrategy :: forall a. MergingStrategy a -> WriterT w m a -> WriterT w m a
tryMergeWithStrategy MergingStrategy a
strategy (WriterStrict.WriterT m (a, w)
f) =
    m (a, w) -> WriterT w m a
forall w (m :: * -> *) a. m (a, w) -> WriterT w m a
WriterStrict.WriterT (m (a, w) -> WriterT w m a) -> m (a, w) -> WriterT w m a
forall a b. (a -> b) -> a -> b
$
      MergingStrategy (a, w) -> m (a, w) -> m (a, w)
forall a. MergingStrategy a -> m a -> m a
forall (m :: * -> *) a.
TryMerge m =>
MergingStrategy a -> m a -> m a
tryMergeWithStrategy (MergingStrategy a -> MergingStrategy w -> MergingStrategy (a, w)
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 a
strategy MergingStrategy w
forall a. Mergeable a => MergingStrategy a
rootStrategy) m (a, w)
f
  {-# INLINE tryMergeWithStrategy #-}

instance
  (Monoid w, Mergeable w, Mergeable s, TryMerge m) =>
  TryMerge (RWSStrict.RWST r w s m)
  where
  tryMergeWithStrategy :: forall a. MergingStrategy a -> RWST r w s m a -> RWST r w s m a
tryMergeWithStrategy MergingStrategy a
strategy (RWSStrict.RWST r -> s -> m (a, s, w)
f) =
    (r -> s -> m (a, s, w)) -> RWST r w s m a
forall r w s (m :: * -> *) a.
(r -> s -> m (a, s, w)) -> RWST r w s m a
RWSStrict.RWST ((r -> s -> m (a, s, w)) -> RWST r w s m a)
-> (r -> s -> m (a, s, w)) -> RWST r w s m a
forall a b. (a -> b) -> a -> b
$
      \r
r s
s ->
        MergingStrategy (a, s, w) -> m (a, s, w) -> m (a, s, w)
forall a. MergingStrategy a -> m a -> m a
forall (m :: * -> *) a.
TryMerge m =>
MergingStrategy a -> m a -> m a
tryMergeWithStrategy
          (MergingStrategy a
-> MergingStrategy s
-> MergingStrategy w
-> MergingStrategy (a, s, w)
forall a b c.
MergingStrategy a
-> MergingStrategy b
-> MergingStrategy c
-> MergingStrategy (a, b, c)
forall (u :: * -> * -> * -> *) a b c.
Mergeable3 u =>
MergingStrategy a
-> MergingStrategy b
-> MergingStrategy c
-> MergingStrategy (u a b c)
liftRootStrategy3 MergingStrategy a
strategy MergingStrategy s
forall a. Mergeable a => MergingStrategy a
rootStrategy MergingStrategy w
forall a. Mergeable a => MergingStrategy a
rootStrategy)
          (r -> s -> m (a, s, w)
f r
r s
s)
  {-# INLINE tryMergeWithStrategy #-}

instance
  (Monoid w, Mergeable w, Mergeable s, TryMerge m) =>
  TryMerge (RWSLazy.RWST r w s m)
  where
  tryMergeWithStrategy :: forall a. MergingStrategy a -> RWST r w s m a -> RWST r w s m a
tryMergeWithStrategy MergingStrategy a
strategy (RWSLazy.RWST r -> s -> m (a, s, w)
f) =
    (r -> s -> m (a, s, w)) -> RWST r w s m a
forall r w s (m :: * -> *) a.
(r -> s -> m (a, s, w)) -> RWST r w s m a
RWSLazy.RWST ((r -> s -> m (a, s, w)) -> RWST r w s m a)
-> (r -> s -> m (a, s, w)) -> RWST r w s m a
forall a b. (a -> b) -> a -> b
$
      \r
r s
s ->
        MergingStrategy (a, s, w) -> m (a, s, w) -> m (a, s, w)
forall a. MergingStrategy a -> m a -> m a
forall (m :: * -> *) a.
TryMerge m =>
MergingStrategy a -> m a -> m a
tryMergeWithStrategy
          (MergingStrategy a
-> MergingStrategy s
-> MergingStrategy w
-> MergingStrategy (a, s, w)
forall a b c.
MergingStrategy a
-> MergingStrategy b
-> MergingStrategy c
-> MergingStrategy (a, b, c)
forall (u :: * -> * -> * -> *) a b c.
Mergeable3 u =>
MergingStrategy a
-> MergingStrategy b
-> MergingStrategy c
-> MergingStrategy (u a b c)
liftRootStrategy3 MergingStrategy a
strategy MergingStrategy s
forall a. Mergeable a => MergingStrategy a
rootStrategy MergingStrategy w
forall a. Mergeable a => MergingStrategy a
rootStrategy)
          (r -> s -> m (a, s, w)
f r
r s
s)
  {-# INLINE tryMergeWithStrategy #-}

instance (TryMerge m) => TryMerge (IdentityT m) where
  tryMergeWithStrategy :: forall a. MergingStrategy a -> IdentityT m a -> IdentityT m a
tryMergeWithStrategy MergingStrategy a
strategy (IdentityT m a
ma) =
    m a -> IdentityT m a
forall {k} (f :: k -> *) (a :: k). f a -> IdentityT f a
IdentityT (m a -> IdentityT m a) -> m a -> IdentityT m a
forall a b. (a -> b) -> a -> b
$ MergingStrategy a -> m a -> m a
forall a. MergingStrategy a -> m a -> m a
forall (m :: * -> *) a.
TryMerge m =>
MergingStrategy a -> m a -> m a
tryMergeWithStrategy MergingStrategy a
strategy m a
ma
  {-# INLINE tryMergeWithStrategy #-}

instance (TryMerge m, Mergeable r) => TryMerge (ContT r m) where
  tryMergeWithStrategy :: forall a. MergingStrategy a -> ContT r m a -> ContT r m a
tryMergeWithStrategy MergingStrategy a
_ (ContT (a -> m r) -> m r
ma) =
    ((a -> m r) -> m r) -> ContT r m a
forall {k} (r :: k) (m :: k -> *) a.
((a -> m r) -> m r) -> ContT r m a
ContT (((a -> m r) -> m r) -> ContT r m a)
-> ((a -> m r) -> m r) -> ContT r m a
forall a b. (a -> b) -> a -> b
$ \a -> m r
c -> MergingStrategy r -> m r -> m r
forall a. MergingStrategy a -> m a -> m a
forall (m :: * -> *) a.
TryMerge m =>
MergingStrategy a -> m a -> m a
tryMergeWithStrategy MergingStrategy r
forall a. Mergeable a => MergingStrategy a
rootStrategy ((a -> m r) -> m r
ma a -> m r
c)
  {-# INLINE tryMergeWithStrategy #-}

-- | Alias for a monad type that has 'TryMerge'.
type MonadTryMerge f = (TryMerge f, Monad f)

#define TRYMERGE_ID(T) \
  instance TryMerge (T) where { \
    tryMergeWithStrategy _ = id; {-# INLINE tryMergeWithStrategy #-} \
  }

#if 1
TRYMERGE_ID(Either a)
TRYMERGE_ID(Maybe)
TRYMERGE_ID(Identity)
TRYMERGE_ID([])
TRYMERGE_ID((,) a)
TRYMERGE_ID((,,) a b)
TRYMERGE_ID((,,,) a b c)
TRYMERGE_ID((,,,,) a b c d)
TRYMERGE_ID((,,,,,) a b c d e)
TRYMERGE_ID((,,,,,,) a b c d e f)
TRYMERGE_ID((,,,,,,,) a b c d e f g)
TRYMERGE_ID((,,,,,,,,) a b c d e f g h)
#endif

instance (TryMerge f, TryMerge g) => TryMerge (Sum f g) where
  tryMergeWithStrategy :: forall a. MergingStrategy a -> Sum f g a -> Sum f g a
tryMergeWithStrategy MergingStrategy a
strategy (InL f a
fa) =
    f a -> Sum f g a
forall {k} (f :: k -> *) (g :: k -> *) (a :: k). f a -> Sum f g a
InL (f a -> Sum f g a) -> f a -> Sum f g a
forall a b. (a -> b) -> a -> b
$ MergingStrategy a -> f a -> f a
forall a. MergingStrategy a -> f a -> f a
forall (m :: * -> *) a.
TryMerge m =>
MergingStrategy a -> m a -> m a
tryMergeWithStrategy MergingStrategy a
strategy f a
fa
  tryMergeWithStrategy MergingStrategy a
strategy (InR g a
fa) =
    g a -> Sum f g a
forall {k} (f :: k -> *) (g :: k -> *) (a :: k). g a -> Sum f g a
InR (g a -> Sum f g a) -> g a -> Sum f g a
forall a b. (a -> b) -> a -> b
$ MergingStrategy a -> g a -> g a
forall a. MergingStrategy a -> g a -> g a
forall (m :: * -> *) a.
TryMerge m =>
MergingStrategy a -> m a -> m a
tryMergeWithStrategy MergingStrategy a
strategy g a
fa

instance TryMerge Monoid.Sum where
  tryMergeWithStrategy :: forall a. MergingStrategy a -> Sum a -> Sum a
tryMergeWithStrategy MergingStrategy a
_ = Sum a -> Sum a
forall a. a -> a
id
  {-# INLINE tryMergeWithStrategy #-}