-- | Utilities related to Monad and Applicative classes
--   Mostly for backwards compatibility.

module MonadUtils
        ( Applicative(..)
        , (<$>)

        , MonadFix(..)
        , MonadIO(..)

        , zipWith3M, zipWith3M_, zipWith4M, zipWithAndUnzipM
        , mapAndUnzipM, mapAndUnzip3M, mapAndUnzip4M, mapAndUnzip5M
        , mapAccumLM
        , mapSndM
        , concatMapM
        , mapMaybeM
        , fmapMaybeM, fmapEitherM
        , anyM, allM, orM
        , foldlM, foldlM_, foldrM
        , maybeMapM
        , whenM, unlessM
        , filterOutM
        ) where

-------------------------------------------------------------------------------
-- Imports
-------------------------------------------------------------------------------

import GhcPrelude

import Control.Applicative
import Control.Monad
import Control.Monad.Fix
import Control.Monad.IO.Class
import Data.Foldable (sequenceA_, foldlM, foldrM)
import Data.List (unzip4, unzip5, zipWith4)

-------------------------------------------------------------------------------
-- Common functions
--  These are used throughout the compiler
-------------------------------------------------------------------------------

{-

Note [Inline @zipWithNM@ functions]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

The inline principle for 'zipWith3M', 'zipWith4M' and 'zipWith3M_' is the same
as for 'zipWithM' and 'zipWithM_' in "Control.Monad", see
Note [Fusion for zipN/zipWithN] in GHC/List.hs for more details.

The 'zipWithM'/'zipWithM_' functions are inlined so that the `zipWith` and
`sequenceA` functions with which they are defined have an opportunity to fuse.

Furthermore, 'zipWith3M'/'zipWith4M' and 'zipWith3M_' have been explicitly
rewritten in a non-recursive way similarly to 'zipWithM'/'zipWithM_', and for
more than just uniformity: after [D5241](https://phabricator.haskell.org/D5241)
for issue #14037, all @zipN@/@zipWithN@ functions fuse, meaning
'zipWith3M'/'zipWIth4M' and 'zipWith3M_'@ now behave like 'zipWithM' and
'zipWithM_', respectively, with regards to fusion.

As such, since there are not any differences between 2-ary 'zipWithM'/
'zipWithM_' and their n-ary counterparts below aside from the number of
arguments, the `INLINE` pragma should be replicated in the @zipWithNM@
functions below as well.

-}

zipWith3M :: Monad m => (a -> b -> c -> m d) -> [a] -> [b] -> [c] -> m [d]
{-# INLINE zipWith3M #-}
-- Inline so that fusion with 'zipWith3' and 'sequenceA' has a chance to fire.
-- See Note [Inline @zipWithNM@ functions] above.
zipWith3M :: (a -> b -> c -> m d) -> [a] -> [b] -> [c] -> m [d]
zipWith3M a -> b -> c -> m d
f [a]
xs [b]
ys [c]
zs = [m d] -> m [d]
forall (t :: * -> *) (f :: * -> *) a.
(Traversable t, Applicative f) =>
t (f a) -> f (t a)
sequenceA ((a -> b -> c -> m d) -> [a] -> [b] -> [c] -> [m d]
forall a b c d. (a -> b -> c -> d) -> [a] -> [b] -> [c] -> [d]
zipWith3 a -> b -> c -> m d
f [a]
xs [b]
ys [c]
zs)

zipWith3M_ :: Monad m => (a -> b -> c -> m d) -> [a] -> [b] -> [c] -> m ()
{-# INLINE zipWith3M_ #-}
-- Inline so that fusion with 'zipWith4' and 'sequenceA' has a chance to fire.
-- See  Note [Inline @zipWithNM@ functions] above.
zipWith3M_ :: (a -> b -> c -> m d) -> [a] -> [b] -> [c] -> m ()
zipWith3M_ a -> b -> c -> m d
f [a]
xs [b]
ys [c]
zs = [m d] -> m ()
forall (t :: * -> *) (f :: * -> *) a.
(Foldable t, Applicative f) =>
t (f a) -> f ()
sequenceA_ ((a -> b -> c -> m d) -> [a] -> [b] -> [c] -> [m d]
forall a b c d. (a -> b -> c -> d) -> [a] -> [b] -> [c] -> [d]
zipWith3 a -> b -> c -> m d
f [a]
xs [b]
ys [c]
zs)

zipWith4M :: Monad m => (a -> b -> c -> d -> m e)
          -> [a] -> [b] -> [c] -> [d] -> m [e]
{-# INLINE zipWith4M #-}
-- Inline so that fusion with 'zipWith5' and 'sequenceA' has a chance to fire.
-- See  Note [Inline @zipWithNM@ functions] above.
zipWith4M :: (a -> b -> c -> d -> m e) -> [a] -> [b] -> [c] -> [d] -> m [e]
zipWith4M a -> b -> c -> d -> m e
f [a]
xs [b]
ys [c]
ws [d]
zs = [m e] -> m [e]
forall (t :: * -> *) (f :: * -> *) a.
(Traversable t, Applicative f) =>
t (f a) -> f (t a)
sequenceA ((a -> b -> c -> d -> m e) -> [a] -> [b] -> [c] -> [d] -> [m e]
forall a b c d e.
(a -> b -> c -> d -> e) -> [a] -> [b] -> [c] -> [d] -> [e]
zipWith4 a -> b -> c -> d -> m e
f [a]
xs [b]
ys [c]
ws [d]
zs)

zipWithAndUnzipM :: Monad m
                 => (a -> b -> m (c, d)) -> [a] -> [b] -> m ([c], [d])
{-# INLINABLE zipWithAndUnzipM #-}
-- See Note [flatten_many performance] in TcFlatten for why this
-- pragma is essential.
zipWithAndUnzipM :: (a -> b -> m (c, d)) -> [a] -> [b] -> m ([c], [d])
zipWithAndUnzipM a -> b -> m (c, d)
f (a
x:[a]
xs) (b
y:[b]
ys)
  = do { (c
c, d
d) <- a -> b -> m (c, d)
f a
x b
y
       ; ([c]
cs, [d]
ds) <- (a -> b -> m (c, d)) -> [a] -> [b] -> m ([c], [d])
forall (m :: * -> *) a b c d.
Monad m =>
(a -> b -> m (c, d)) -> [a] -> [b] -> m ([c], [d])
zipWithAndUnzipM a -> b -> m (c, d)
f [a]
xs [b]
ys
       ; ([c], [d]) -> m ([c], [d])
forall (m :: * -> *) a. Monad m => a -> m a
return (c
cc -> [c] -> [c]
forall a. a -> [a] -> [a]
:[c]
cs, d
dd -> [d] -> [d]
forall a. a -> [a] -> [a]
:[d]
ds) }
zipWithAndUnzipM a -> b -> m (c, d)
_ [a]
_ [b]
_ = ([c], [d]) -> m ([c], [d])
forall (m :: * -> *) a. Monad m => a -> m a
return ([], [])

{-

Note [Inline @mapAndUnzipNM@ functions]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

The inline principle is the same as 'mapAndUnzipM' in "Control.Monad".
The 'mapAndUnzipM' function is inlined so that the `unzip` and `traverse`
functions with which it is defined have an opportunity to fuse, see
Note [Inline @unzipN@ functions] in Data/OldList.hs for more details.

Furthermore, the @mapAndUnzipNM@ functions have been explicitly rewritten in a
non-recursive way similarly to 'mapAndUnzipM', and for more than just
uniformity: after [D5249](https://phabricator.haskell.org/D5249) for Trac
ticket #14037, all @unzipN@ functions fuse, meaning 'mapAndUnzip3M',
'mapAndUnzip4M' and 'mapAndUnzip5M' now behave like 'mapAndUnzipM' with regards
to fusion.

As such, since there are not any differences between 2-ary 'mapAndUnzipM' and
its n-ary counterparts below aside from the number of arguments, the `INLINE`
pragma should be replicated in the @mapAndUnzipNM@ functions below as well.

-}

-- | mapAndUnzipM for triples
mapAndUnzip3M :: Monad m => (a -> m (b,c,d)) -> [a] -> m ([b],[c],[d])
{-# INLINE mapAndUnzip3M #-}
-- Inline so that fusion with 'unzip3' and 'traverse' has a chance to fire.
-- See Note [Inline @mapAndUnzipNM@ functions] above.
mapAndUnzip3M :: (a -> m (b, c, d)) -> [a] -> m ([b], [c], [d])
mapAndUnzip3M a -> m (b, c, d)
f [a]
xs =  [(b, c, d)] -> ([b], [c], [d])
forall a b c. [(a, b, c)] -> ([a], [b], [c])
unzip3 ([(b, c, d)] -> ([b], [c], [d]))
-> m [(b, c, d)] -> m ([b], [c], [d])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (a -> m (b, c, d)) -> [a] -> m [(b, c, d)]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse a -> m (b, c, d)
f [a]
xs

mapAndUnzip4M :: Monad m => (a -> m (b,c,d,e)) -> [a] -> m ([b],[c],[d],[e])
{-# INLINE mapAndUnzip4M #-}
-- Inline so that fusion with 'unzip4' and 'traverse' has a chance to fire.
-- See Note [Inline @mapAndUnzipNM@ functions] above.
mapAndUnzip4M :: (a -> m (b, c, d, e)) -> [a] -> m ([b], [c], [d], [e])
mapAndUnzip4M a -> m (b, c, d, e)
f [a]
xs =  [(b, c, d, e)] -> ([b], [c], [d], [e])
forall a b c d. [(a, b, c, d)] -> ([a], [b], [c], [d])
unzip4 ([(b, c, d, e)] -> ([b], [c], [d], [e]))
-> m [(b, c, d, e)] -> m ([b], [c], [d], [e])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (a -> m (b, c, d, e)) -> [a] -> m [(b, c, d, e)]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse a -> m (b, c, d, e)
f [a]
xs

mapAndUnzip5M :: Monad m => (a -> m (b,c,d,e,f)) -> [a] -> m ([b],[c],[d],[e],[f])
{-# INLINE mapAndUnzip5M #-}
-- Inline so that fusion with 'unzip5' and 'traverse' has a chance to fire.
-- See Note [Inline @mapAndUnzipNM@ functions] above.
mapAndUnzip5M :: (a -> m (b, c, d, e, f)) -> [a] -> m ([b], [c], [d], [e], [f])
mapAndUnzip5M a -> m (b, c, d, e, f)
f [a]
xs =  [(b, c, d, e, f)] -> ([b], [c], [d], [e], [f])
forall a b c d e. [(a, b, c, d, e)] -> ([a], [b], [c], [d], [e])
unzip5 ([(b, c, d, e, f)] -> ([b], [c], [d], [e], [f]))
-> m [(b, c, d, e, f)] -> m ([b], [c], [d], [e], [f])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (a -> m (b, c, d, e, f)) -> [a] -> m [(b, c, d, e, f)]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse a -> m (b, c, d, e, f)
f [a]
xs

-- | Monadic version of mapAccumL
mapAccumLM :: Monad m
            => (acc -> x -> m (acc, y)) -- ^ combining function
            -> acc                      -- ^ initial state
            -> [x]                      -- ^ inputs
            -> m (acc, [y])             -- ^ final state, outputs
mapAccumLM :: (acc -> x -> m (acc, y)) -> acc -> [x] -> m (acc, [y])
mapAccumLM acc -> x -> m (acc, y)
_ acc
s []     = (acc, [y]) -> m (acc, [y])
forall (m :: * -> *) a. Monad m => a -> m a
return (acc
s, [])
mapAccumLM acc -> x -> m (acc, y)
f acc
s (x
x:[x]
xs) = do
    (acc
s1, y
x')  <- acc -> x -> m (acc, y)
f acc
s x
x
    (acc
s2, [y]
xs') <- (acc -> x -> m (acc, y)) -> acc -> [x] -> m (acc, [y])
forall (m :: * -> *) acc x y.
Monad m =>
(acc -> x -> m (acc, y)) -> acc -> [x] -> m (acc, [y])
mapAccumLM acc -> x -> m (acc, y)
f acc
s1 [x]
xs
    (acc, [y]) -> m (acc, [y])
forall (m :: * -> *) a. Monad m => a -> m a
return    (acc
s2, y
x' y -> [y] -> [y]
forall a. a -> [a] -> [a]
: [y]
xs')

-- | Monadic version of mapSnd
mapSndM :: Monad m => (b -> m c) -> [(a,b)] -> m [(a,c)]
mapSndM :: (b -> m c) -> [(a, b)] -> m [(a, c)]
mapSndM b -> m c
_ []         = [(a, c)] -> m [(a, c)]
forall (m :: * -> *) a. Monad m => a -> m a
return []
mapSndM b -> m c
f ((a
a,b
b):[(a, b)]
xs) = do { c
c <- b -> m c
f b
b; [(a, c)]
rs <- (b -> m c) -> [(a, b)] -> m [(a, c)]
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> [(a, b)] -> m [(a, c)]
mapSndM b -> m c
f [(a, b)]
xs; [(a, c)] -> m [(a, c)]
forall (m :: * -> *) a. Monad m => a -> m a
return ((a
a,c
c)(a, c) -> [(a, c)] -> [(a, c)]
forall a. a -> [a] -> [a]
:[(a, c)]
rs) }

-- | Monadic version of concatMap
concatMapM :: Monad m => (a -> m [b]) -> [a] -> m [b]
concatMapM :: (a -> m [b]) -> [a] -> m [b]
concatMapM a -> m [b]
f [a]
xs = ([[b]] -> [b]) -> m [[b]] -> m [b]
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM [[b]] -> [b]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ((a -> m [b]) -> [a] -> m [[b]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM a -> m [b]
f [a]
xs)

-- | Applicative version of mapMaybe
mapMaybeM :: Applicative m => (a -> m (Maybe b)) -> [a] -> m [b]
mapMaybeM :: (a -> m (Maybe b)) -> [a] -> m [b]
mapMaybeM a -> m (Maybe b)
f = (a -> m [b] -> m [b]) -> m [b] -> [a] -> m [b]
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr a -> m [b] -> m [b]
g ([b] -> m [b]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [])
  where g :: a -> m [b] -> m [b]
g a
a = (Maybe b -> [b] -> [b]) -> m (Maybe b) -> m [b] -> m [b]
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 (([b] -> [b]) -> (b -> [b] -> [b]) -> Maybe b -> [b] -> [b]
forall b a. b -> (a -> b) -> Maybe a -> b
maybe [b] -> [b]
forall a. a -> a
id (:)) (a -> m (Maybe b)
f a
a)

-- | Monadic version of fmap
fmapMaybeM :: (Monad m) => (a -> m b) -> Maybe a -> m (Maybe b)
fmapMaybeM :: (a -> m b) -> Maybe a -> m (Maybe b)
fmapMaybeM a -> m b
_ Maybe a
Nothing  = Maybe b -> m (Maybe b)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe b
forall a. Maybe a
Nothing
fmapMaybeM a -> m b
f (Just a
x) = a -> m b
f a
x m b -> (b -> m (Maybe b)) -> m (Maybe b)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (Maybe b -> m (Maybe b)
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe b -> m (Maybe b)) -> (b -> Maybe b) -> b -> m (Maybe b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. b -> Maybe b
forall a. a -> Maybe a
Just)

-- | Monadic version of fmap
fmapEitherM :: Monad m => (a -> m b) -> (c -> m d) -> Either a c -> m (Either b d)
fmapEitherM :: (a -> m b) -> (c -> m d) -> Either a c -> m (Either b d)
fmapEitherM a -> m b
fl c -> m d
_ (Left  a
a) = a -> m b
fl a
a m b -> (b -> m (Either b d)) -> m (Either b d)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (Either b d -> m (Either b d)
forall (m :: * -> *) a. Monad m => a -> m a
return (Either b d -> m (Either b d))
-> (b -> Either b d) -> b -> m (Either b d)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. b -> Either b d
forall a b. a -> Either a b
Left)
fmapEitherM a -> m b
_ c -> m d
fr (Right c
b) = c -> m d
fr c
b m d -> (d -> m (Either b d)) -> m (Either b d)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (Either b d -> m (Either b d)
forall (m :: * -> *) a. Monad m => a -> m a
return (Either b d -> m (Either b d))
-> (d -> Either b d) -> d -> m (Either b d)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. d -> Either b d
forall a b. b -> Either a b
Right)

-- | Monadic version of 'any', aborts the computation at the first @True@ value
anyM :: Monad m => (a -> m Bool) -> [a] -> m Bool
anyM :: (a -> m Bool) -> [a] -> m Bool
anyM a -> m Bool
_ []     = Bool -> m Bool
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False
anyM a -> m Bool
f (a
x:[a]
xs) = do Bool
b <- a -> m Bool
f a
x
                   if Bool
b then Bool -> m Bool
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True
                        else (a -> m Bool) -> [a] -> m Bool
forall (m :: * -> *) a. Monad m => (a -> m Bool) -> [a] -> m Bool
anyM a -> m Bool
f [a]
xs

-- | Monad version of 'all', aborts the computation at the first @False@ value
allM :: Monad m => (a -> m Bool) -> [a] -> m Bool
allM :: (a -> m Bool) -> [a] -> m Bool
allM a -> m Bool
_ []     = Bool -> m Bool
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True
allM a -> m Bool
f (a
b:[a]
bs) = (a -> m Bool
f a
b) m Bool -> (Bool -> m Bool) -> m Bool
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (\Bool
bv -> if Bool
bv then (a -> m Bool) -> [a] -> m Bool
forall (m :: * -> *) a. Monad m => (a -> m Bool) -> [a] -> m Bool
allM a -> m Bool
f [a]
bs else Bool -> m Bool
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False)

-- | Monadic version of or
orM :: Monad m => m Bool -> m Bool -> m Bool
orM :: m Bool -> m Bool -> m Bool
orM m Bool
m1 m Bool
m2 = m Bool
m1 m Bool -> (Bool -> m Bool) -> m Bool
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \Bool
x -> if Bool
x then Bool -> m Bool
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True else m Bool
m2

-- | Monadic version of foldl that discards its result
foldlM_ :: (Monad m, Foldable t) => (a -> b -> m a) -> a -> t b -> m ()
foldlM_ :: (a -> b -> m a) -> a -> t b -> m ()
foldlM_ = (a -> b -> m a) -> a -> t b -> m ()
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m ()
foldM_

-- | Monadic version of fmap specialised for Maybe
maybeMapM :: Monad m => (a -> m b) -> (Maybe a -> m (Maybe b))
maybeMapM :: (a -> m b) -> Maybe a -> m (Maybe b)
maybeMapM a -> m b
_ Maybe a
Nothing  = Maybe b -> m (Maybe b)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe b
forall a. Maybe a
Nothing
maybeMapM a -> m b
m (Just a
x) = (b -> Maybe b) -> m b -> m (Maybe b)
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM b -> Maybe b
forall a. a -> Maybe a
Just (m b -> m (Maybe b)) -> m b -> m (Maybe b)
forall a b. (a -> b) -> a -> b
$ a -> m b
m a
x

-- | Monadic version of @when@, taking the condition in the monad
whenM :: Monad m => m Bool -> m () -> m ()
whenM :: m Bool -> m () -> m ()
whenM m Bool
mb m ()
thing = do { Bool
b <- m Bool
mb
                    ; Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
b m ()
thing }

-- | Monadic version of @unless@, taking the condition in the monad
unlessM :: Monad m => m Bool -> m () -> m ()
unlessM :: m Bool -> m () -> m ()
unlessM m Bool
condM m ()
acc = do { Bool
cond <- m Bool
condM
                       ; Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless Bool
cond m ()
acc }

-- | Like 'filterM', only it reverses the sense of the test.
filterOutM :: (Applicative m) => (a -> m Bool) -> [a] -> m [a]
filterOutM :: (a -> m Bool) -> [a] -> m [a]
filterOutM a -> m Bool
p =
  (a -> m [a] -> m [a]) -> m [a] -> [a] -> m [a]
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\ a
x -> (Bool -> [a] -> [a]) -> m Bool -> m [a] -> m [a]
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 (\ Bool
flg -> if Bool
flg then [a] -> [a]
forall a. a -> a
id else (a
xa -> [a] -> [a]
forall a. a -> [a] -> [a]
:)) (a -> m Bool
p a
x)) ([a] -> m [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [])