content: {-# LANGUAGE AllowAmbiguousTypes #-}
  {-# LANGUAGE ApplicativeDo #-}
  {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE FlexibleInstances #-}
  {-# LANGUAGE MonoLocalBinds #-}
  {-# LANGUAGE RankNTypes #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE Trustworthy #-}
  {-# LANGUAGE TypeApplications #-}
  
  -- |
  -- Module      :   Grisette.Unified.Lib.Control.Monad
  -- Copyright   :   (c) Sirui Lu 2021-2024
  -- License     :   BSD-3-Clause (see the LICENSE file)
  --
  -- Maintainer  :   siruilu@cs.washington.edu
  -- Stability   :   Experimental
  -- Portability :   GHC only
  module Grisette.Unified.Lib.Control.Monad
    ( -- * Functor and Monad classes
      mrgFmap,
      (.<$),
      mrgReturnWithStrategy,
      mrgBindWithStrategy,
      mrgReturn,
      (.>>=),
      (.>>),
      mrgFail,
      mrgMzero,
      mrgMplus,
  
      -- * Functions
  
      -- ** Basic 'Monad' functions
      mrgMapM,
      mrgMapM_,
      mrgForM,
      mrgForM_,
      mrgSequence,
      mrgSequence_,
      (.=<<),
      (.>=>),
      (.<=<),
      mrgForever,
      mrgVoid,
  
      -- ** Generalisations of list functions
      mrgJoin,
      mrgMsum,
      mrgMfilter,
      symMfilter,
      mrgFilterM,
      symFilterM,
      mrgMapAndUnzipM,
      mrgZipWithM,
      mrgZipWithM_,
      mrgFoldM,
      mrgFoldM_,
      mrgReplicateM,
      symReplicateM,
      mrgReplicateM_,
      symReplicateM_,
  
      -- ** Conditional execution of monadic expressions
      mrgGuard,
      symGuard,
      mrgWhen,
      symWhen,
      mrgUnless,
      symUnless,
  
      -- ** Monadic lifting operators
      mrgLiftM,
      mrgLiftM2,
      mrgLiftM3,
      mrgLiftM4,
      mrgLiftM5,
      mrgAp,
  
      -- ** Strict monadic functions
      (.<$!>),
    )
  where
  
  import Control.Applicative (Alternative)
  import Control.Monad (MonadPlus (mplus, mzero), join)
  import Grisette.Internal.Core.Data.Class.LogicalOp (LogicalOp (symNot, (.||)))
  import Grisette.Internal.Core.Data.Class.Mergeable
    ( Mergeable (rootStrategy),
      MergingStrategy,
    )
  import Grisette.Internal.Core.Data.Class.SimpleMergeable (SymBranching)
  import Grisette.Internal.Core.Data.Class.TryMerge
    ( MonadTryMerge,
      TryMerge (tryMergeWithStrategy),
      tryMerge,
    )
  import Grisette.Internal.SymPrim.SymBool (SymBool)
  import Grisette.Lib.Data.Functor (mrgFmap, mrgUnzip, mrgVoid, (.<$))
  import Grisette.Lib.Data.Traversable
    ( mrgForM,
      mrgMapM,
      mrgSequence,
      mrgSequenceA,
      mrgTraverse,
    )
  import Grisette.Unified
    ( EvalModeBase,
      GetBool,
      UnifiedBranching,
      UnifiedSymOrd,
      mrgIf,
      (.<=),
    )
  import Grisette.Unified.Lib.Control.Applicative
    ( mrgEmpty,
      mrgLiftA2,
      mrgPure,
      (.*>),
      (.<$>),
      (.<*>),
    )
  import Grisette.Unified.Lib.Data.Foldable
    ( mrgFoldlM,
      mrgForM_,
      mrgMapM_,
      mrgMsum,
      mrgSequenceA_,
      mrgSequence_,
    )
  
  -- | 'return' with 'Grisette.Core.MergingStrategy' knowledge propagation.
  mrgReturnWithStrategy :: (MonadTryMerge u) => MergingStrategy a -> a -> u a
  mrgReturnWithStrategy s = tryMergeWithStrategy s . return
  {-# INLINE mrgReturnWithStrategy #-}
  
  -- | '>>=' with 'Grisette.Core.MergingStrategy' knowledge propagation.
  mrgBindWithStrategy ::
    (MonadTryMerge u) =>
    MergingStrategy a ->
    MergingStrategy b ->
    u a ->
    (a -> u b) ->
    u b
  mrgBindWithStrategy sa sb a f =
    tryMergeWithStrategy sb $ tryMergeWithStrategy sa a >>= f
  {-# INLINE mrgBindWithStrategy #-}
  
  -- | 'return' with 'Grisette.Core.MergingStrategy' knowledge propagation.
  mrgReturn :: (MonadTryMerge u, Mergeable a) => a -> u a
  mrgReturn = mrgReturnWithStrategy rootStrategy
  {-# INLINE mrgReturn #-}
  
  infixl 1 .>>=
  
  -- | '>>=' with 'Grisette.Core.MergingStrategy' knowledge propagation.
  (.>>=) ::
    (MonadTryMerge u, Mergeable a, Mergeable b) =>
    u a ->
    (a -> u b) ->
    u b
  (.>>=) = mrgBindWithStrategy rootStrategy rootStrategy
  {-# INLINE (.>>=) #-}
  
  infixl 1 .>>
  
  -- | '>>' with 'Grisette.Core.MergingStrategy' knowledge propagation.
  --
  -- This is usually more efficient than calling the original '>>' and merge the
  -- results.
  (.>>) :: (MonadTryMerge m, Mergeable a, Mergeable b) => m a -> m b -> m b
  a .>> f = tryMerge $ mrgVoid a >> f
  {-# INLINE (.>>) #-}
  
  -- | 'fail' with 'Grisette.Core.MergingStrategy' knowledge propagation.
  mrgFail :: (MonadTryMerge m, Mergeable a, MonadFail m) => String -> m a
  mrgFail = tryMerge . fail
  {-# INLINE mrgFail #-}
  
  -- | 'Control.Monad.mzero' with 'Grisette.Core.MergingStrategy' knowledge
  -- propagation.
  mrgMzero :: forall m a. (MonadTryMerge m, Mergeable a, MonadPlus m) => m a
  mrgMzero = tryMerge mzero
  {-# INLINE mrgMzero #-}
  
  -- | 'Control.Monad.mplus' with 'Grisette.Core.MergingStrategy' knowledge
  -- propagation.
  mrgMplus ::
    forall m a. (MonadTryMerge m, Mergeable a, MonadPlus m) => m a -> m a -> m a
  mrgMplus a b = tryMerge $ mplus (tryMerge a) (tryMerge b)
  {-# INLINE mrgMplus #-}
  
  infixr 1 .=<<
  
  -- | '=<<' with 'Grisette.Core.MergingStrategy' knowledge propagation.
  (.=<<) ::
    (MonadTryMerge m, Mergeable a, Mergeable b) => (a -> m b) -> m a -> m b
  f .=<< a = tryMerge $ f =<< tryMerge a
  {-# INLINE (.=<<) #-}
  
  infixr 1 .>=>
  
  -- | 'Control.Monad.>=>' with 'Grisette.Core.MergingStrategy' knowledge
  -- propagation.
  (.>=>) ::
    (MonadTryMerge m, Mergeable a, Mergeable b, Mergeable c) =>
    (a -> m b) ->
    (b -> m c) ->
    a ->
    m c
  f .>=> g = \a -> tryMerge $ tryMerge (f a) >>= g
  {-# INLINE (.>=>) #-}
  
  infixr 1 .<=<
  
  -- | 'Control.Monad.<=<' with 'Grisette.Core.MergingStrategy' knowledge
  -- propagation.
  (.<=<) ::
    (MonadTryMerge m, Mergeable a, Mergeable b, Mergeable c) =>
    (b -> m c) ->
    (a -> m b) ->
    a ->
    m c
  (.<=<) = flip (.>=>)
  {-# INLINE (.<=<) #-}
  
  -- | 'Control.Monad.forever' with 'Grisette.Core.MergingStrategy' knowledge
  -- propagation.
  mrgForever ::
    (Applicative m, TryMerge m, Mergeable b, Mergeable a) => m a -> m b
  mrgForever a = let a' = a .*> a' in a'
  {-# INLINE mrgForever #-}
  
  -- | 'Control.Monad.join' with 'Grisette.Core.MergingStrategy' knowledge
  -- propagation.
  mrgJoin :: (MonadTryMerge m, Mergeable a) => m (m a) -> m a
  mrgJoin a = tryMerge $ join a
  {-# INLINE mrgJoin #-}
  
  -- | 'Control.Monad.mfilter' with 'Grisette.Core.MergingStrategy' knowledge
  -- propagation.
  mrgMfilter ::
    (MonadTryMerge m, MonadPlus m, Mergeable a) =>
    (a -> Bool) ->
    m a ->
    m a
  mrgMfilter p ma = do
    a <- tryMerge ma
    if p a then mrgReturn a else mrgMzero
  {-# INLINE mrgMfilter #-}
  
  -- | 'Control.Monad.mfilter' with 'Grisette.Core.MergingStrategy' knowledge
  -- propagation and symbolic conditions.
  symMfilter ::
    forall mode m a.
    ( MonadTryMerge m,
      MonadPlus m,
      UnifiedBranching mode m,
      Mergeable a
    ) =>
    (a -> GetBool mode) ->
    m a ->
    m a
  symMfilter p ma = do
    a <- tryMerge ma
    mrgIf (p a) (mrgReturn a) mrgMzero
  {-# INLINE symMfilter #-}
  
  -- | 'Control.Monad.filterM' with 'Grisette.Core.MergingStrategy' knowledge
  -- propagation.
  mrgFilterM ::
    (TryMerge m, Applicative m, Mergeable a, Foldable t) =>
    (a -> m Bool) ->
    t a ->
    m [a]
  mrgFilterM p =
    foldr
      (\x lst -> (\flg -> if flg then (x :) else id) .<$> p x .<*> lst)
      (mrgPure [])
  {-# INLINE mrgFilterM #-}
  
  -- | 'Control.Monad.filterM' with 'Grisette.Core.MergingStrategy' knowledge
  -- propagation and symbolic conditions.
  symFilterM ::
    forall mode m t a.
    ( TryMerge m,
      UnifiedBranching mode m,
      MonadTryMerge m,
      EvalModeBase mode,
      Mergeable a,
      Foldable t
    ) =>
    (a -> m (GetBool mode)) ->
    t a ->
    m [a]
  symFilterM p =
    foldr
      ( \x lst -> do
          flag <- tryMerge $ p x
          mrgIf flag ((x :) <$> lst) lst
      )
      (mrgPure [])
  {-# INLINE symFilterM #-}
  
  -- | 'Control.Monad.mapAndUnzipM' with 'Grisette.Core.MergingStrategy' knowledge
  -- propagation.
  mrgMapAndUnzipM ::
    ( Applicative m,
      TryMerge m,
      Mergeable b,
      Mergeable c
    ) =>
    (a -> m (b, c)) ->
    [a] ->
    m ([b], [c])
  mrgMapAndUnzipM f xs = mrgUnzip .<$> mrgTraverse f xs
  {-# INLINE mrgMapAndUnzipM #-}
  
  -- | 'Control.Monad.zipWithM' with 'Grisette.Core.MergingStrategy' knowledge
  -- propagation.
  mrgZipWithM ::
    (Applicative m, TryMerge m, Mergeable c) =>
    (a -> b -> m c) ->
    [a] ->
    [b] ->
    m [c]
  mrgZipWithM f xs ys = mrgSequenceA (zipWith f xs ys)
  {-# INLINE mrgZipWithM #-}
  
  -- | 'Control.Monad.zipWithM_' with 'Grisette.Core.MergingStrategy' knowledge
  -- propagation.
  mrgZipWithM_ ::
    (Applicative m, TryMerge m, Mergeable c) =>
    (a -> b -> m c) ->
    [a] ->
    [b] ->
    m ()
  mrgZipWithM_ f xs ys = mrgSequenceA_ (zipWith f xs ys)
  {-# INLINE mrgZipWithM_ #-}
  
  -- | 'Control.Monad.foldM' with 'Grisette.Core.MergingStrategy' knowledge
  -- propagation.
  mrgFoldM ::
    (MonadTryMerge m, Mergeable b, Foldable t) =>
    (b -> a -> m b) ->
    b ->
    t a ->
    m b
  mrgFoldM = mrgFoldlM
  {-# INLINE mrgFoldM #-}
  
  -- | 'Control.Monad.foldM_' with 'Grisette.Core.MergingStrategy' knowledge
  -- propagation.
  mrgFoldM_ ::
    (MonadTryMerge m, Foldable t, Mergeable b) =>
    (b -> a -> m b) ->
    b ->
    t a ->
    m ()
  mrgFoldM_ f a xs = mrgFoldlM f a xs .>> mrgPure ()
  {-# INLINE mrgFoldM_ #-}
  
  -- | 'Control.Monad.replicateM' with 'Grisette.Core.MergingStrategy' knowledge
  -- propagation.
  mrgReplicateM ::
    (Applicative m, TryMerge m, Mergeable a) =>
    Int ->
    m a ->
    m [a]
  mrgReplicateM n = mrgSequenceA . replicate n
  {-# INLINE mrgReplicateM #-}
  
  -- | 'Control.Monad.replicateM' with 'Grisette.Core.MergingStrategy' knowledge
  -- propagation and symbolic number of elements.
  symReplicateM ::
    forall mode m a int.
    ( EvalModeBase mode,
      TryMerge m,
      Applicative m,
      Mergeable a,
      Num int,
      UnifiedBranching mode m,
      UnifiedSymOrd mode Int,
      UnifiedSymOrd mode int
    ) =>
    Int ->
    int ->
    m a ->
    m [a]
  symReplicateM maxCnt cnt0 f =
    loop maxCnt cnt0
    where
      loop concreteCnt cnt =
        mrgIf @mode
          (cnt .<= 0 .|| concreteCnt .<= 0)
          (mrgPure [])
          (mrgLiftA2 (:) f (loop (concreteCnt - 1) (cnt - 1)))
  {-# INLINE symReplicateM #-}
  
  -- | 'Control.Monad.replicateM_' with 'Grisette.Core.MergingStrategy' knowledge
  -- propagation.
  mrgReplicateM_ ::
    (Applicative m, TryMerge m, Mergeable a) =>
    Int ->
    m a ->
    m ()
  mrgReplicateM_ n = mrgSequenceA_ . replicate n
  {-# INLINE mrgReplicateM_ #-}
  
  -- | 'Control.Monad.replicateM_' with 'Grisette.Core.MergingStrategy' knowledge
  -- propagation and symbolic number of elements.
  symReplicateM_ ::
    forall mode m a int.
    ( EvalModeBase mode,
      TryMerge m,
      Applicative m,
      Mergeable a,
      Num int,
      UnifiedBranching mode m,
      UnifiedSymOrd mode Int,
      UnifiedSymOrd mode int
    ) =>
    Int ->
    int ->
    m a ->
    m ()
  symReplicateM_ maxCnt cnt0 f =
    loop maxCnt cnt0
    where
      loop concreteCnt cnt =
        mrgIf @mode
          (cnt .<= 0 .|| concreteCnt .<= 0)
          (mrgPure ())
          (f .*> (loop (concreteCnt - 1) (cnt - 1)))
  {-# INLINE symReplicateM_ #-}
  
  -- | 'Control.Monad.guard' with 'Grisette.Core.MergingStrategy' knowledge
  -- propagation.
  mrgGuard :: (Alternative m, TryMerge m) => Bool -> m ()
  mrgGuard True = mrgPure ()
  mrgGuard False = mrgEmpty
  {-# INLINE mrgGuard #-}
  
  -- | 'Control.Monad.guard' with 'Grisette.Core.MergingStrategy' knowledge
  -- propagation and symbolic conditions.
  symGuard :: (SymBranching m, TryMerge m, Alternative m) => SymBool -> m ()
  symGuard b = mrgIf b (mrgPure ()) mrgEmpty
  {-# INLINE symGuard #-}
  
  -- | 'Control.Monad.when' with 'Grisette.Core.MergingStrategy' knowledge
  -- propagation.
  mrgWhen :: (Applicative m, TryMerge m) => Bool -> m () -> m ()
  mrgWhen True a = tryMerge a
  mrgWhen False _ = mrgPure ()
  {-# INLINE mrgWhen #-}
  
  -- | 'Control.Monad.when' with 'Grisette.Core.MergingStrategy' knowledge
  -- propagation and symbolic conditions.
  symWhen ::
    (Applicative m, TryMerge m, SymBranching m) => SymBool -> m () -> m ()
  symWhen b a = mrgIf b a (mrgPure ())
  {-# INLINE symWhen #-}
  
  -- | 'Control.Monad.unless' with 'Grisette.Core.MergingStrategy' knowledge
  -- propagation.
  mrgUnless :: (Applicative m, TryMerge m) => Bool -> m () -> m ()
  mrgUnless b = mrgWhen (not b)
  {-# INLINE mrgUnless #-}
  
  -- | 'Control.Monad.unless' with 'Grisette.Core.MergingStrategy' knowledge
  -- propagation and symbolic conditions.
  symUnless ::
    (Applicative m, TryMerge m, SymBranching m) => SymBool -> m () -> m ()
  symUnless b = symWhen (symNot b)
  {-# INLINE symUnless #-}
  
  -- | 'Control.Monad.liftM' with 'Grisette.Core.MergingStrategy' knowledge
  -- propagation.
  mrgLiftM ::
    (MonadTryMerge m, Mergeable a, Mergeable b) => (a -> b) -> m a -> m b
  mrgLiftM f a = f .<$> a
  {-# INLINE mrgLiftM #-}
  
  -- | 'Control.Monad.liftM2' with 'Grisette.Core.MergingStrategy' knowledge
  -- propagation.
  mrgLiftM2 ::
    (MonadTryMerge m, Mergeable a, Mergeable b, Mergeable c) =>
    (a -> b -> c) ->
    m a ->
    m b ->
    m c
  mrgLiftM2 f a b = f .<$> a .<*> b
  {-# INLINE mrgLiftM2 #-}
  
  -- | 'Control.Monad.liftM3' with 'Grisette.Core.MergingStrategy' knowledge
  -- propagation.
  mrgLiftM3 ::
    (MonadTryMerge m, Mergeable a, Mergeable b, Mergeable c, Mergeable d) =>
    (a -> b -> c -> d) ->
    m a ->
    m b ->
    m c ->
    m d
  mrgLiftM3 f a b c = f .<$> a .<*> b .<*> c
  {-# INLINE mrgLiftM3 #-}
  
  -- | 'Control.Monad.liftM4' with 'Grisette.Core.MergingStrategy' knowledge
  -- propagation.
  mrgLiftM4 ::
    ( MonadTryMerge m,
      Mergeable a,
      Mergeable b,
      Mergeable c,
      Mergeable d,
      Mergeable e
    ) =>
    (a -> b -> c -> d -> e) ->
    m a ->
    m b ->
    m c ->
    m d ->
    m e
  mrgLiftM4 f a b c d = f .<$> a .<*> b .<*> c .<*> d
  {-# INLINE mrgLiftM4 #-}
  
  -- | 'Control.Monad.liftM5' with 'Grisette.Core.MergingStrategy' knowledge
  -- propagation.
  mrgLiftM5 ::
    ( MonadTryMerge m,
      Mergeable a,
      Mergeable b,
      Mergeable c,
      Mergeable d,
      Mergeable e,
      Mergeable f
    ) =>
    (a -> b -> c -> d -> e -> f) ->
    m a ->
    m b ->
    m c ->
    m d ->
    m e ->
    m f
  mrgLiftM5 f a b c d e = f .<$> a .<*> b .<*> c .<*> d .<*> e
  {-# INLINE mrgLiftM5 #-}
  
  -- | 'Control.Monad.<*>' with 'Grisette.Core.MergingStrategy' knowledge
  -- propagation.
  mrgAp ::
    (MonadTryMerge m, Mergeable a, Mergeable b) => m (a -> b) -> m a -> m b
  mrgAp = (.<*>)
  {-# INLINE mrgAp #-}
  
  infixl 4 .<$!>
  
  -- | 'Control.Monad.<$!>' with 'Grisette.Core.MergingStrategy' knowledge
  -- propagation. Merging is always strict so we can directly use
  -- 'Grisette.Internal.Unified.Lib.Data.Functor..<$>'.
  (.<$!>) ::
    (MonadTryMerge m, Mergeable a, Mergeable b) => (a -> b) -> m a -> m b
  f .<$!> a = f .<$> a
  {-# INLINE (.<$!>) #-}
  