MonadUtils.hs 8.45 KB
Newer Older
1
-- | Utilities related to Monad and Applicative classes
2
--   Mostly for backwards compatibility.
3 4 5 6

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

8 9
        , MonadFix(..)
        , MonadIO(..)
10

11 12
        , zipWith3M, zipWith3M_, zipWith4M, zipWithAndUnzipM
        , mapAndUnzipM, mapAndUnzip3M, mapAndUnzip4M, mapAndUnzip5M
13 14 15
        , mapAccumLM
        , mapSndM
        , concatMapM
16
        , mapMaybeM
17
        , fmapMaybeM, fmapEitherM
18
        , anyM, allM, orM
19
        , foldlM, foldlM_, foldrM
20
        , maybeMapM
21
        , whenM, unlessM
22
        , filterOutM
23 24
        ) where

25
-------------------------------------------------------------------------------
26
-- Imports
27
-------------------------------------------------------------------------------
28

29 30
import GhcPrelude

31
import Control.Applicative
32 33
import Control.Monad
import Control.Monad.Fix
34
import Control.Monad.IO.Class
35
import Data.Foldable (sequenceA_, foldr)
36
import Data.List (unzip4, unzip5, zipWith4)
37

38
-------------------------------------------------------------------------------
39
-- Common functions
40
--  These are used throughout the compiler
41
-------------------------------------------------------------------------------
42

43 44 45 46 47 48 49 50 51 52 53 54 55 56 57
{-

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)
58
for issue #14037, all @zipN@/@zipWithN@ functions fuse, meaning
59 60 61 62 63 64 65 66 67 68
'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.

-}

69
zipWith3M :: Monad m => (a -> b -> c -> m d) -> [a] -> [b] -> [c] -> m [d]
70 71 72 73
{-# INLINE zipWith3M #-}
-- Inline so that fusion with 'zipWith3' and 'sequenceA' has a chance to fire.
-- See Note [Inline @zipWithNM@ functions] above.
zipWith3M f xs ys zs = sequenceA (zipWith3 f xs ys zs)
74

eir@cis.upenn.edu's avatar
eir@cis.upenn.edu committed
75
zipWith3M_ :: Monad m => (a -> b -> c -> m d) -> [a] -> [b] -> [c] -> m ()
76 77 78 79
{-# INLINE zipWith3M_ #-}
-- Inline so that fusion with 'zipWith4' and 'sequenceA' has a chance to fire.
-- See  Note [Inline @zipWithNM@ functions] above.
zipWith3M_ f xs ys zs = sequenceA_ (zipWith3 f xs ys zs)
eir@cis.upenn.edu's avatar
eir@cis.upenn.edu committed
80

81 82
zipWith4M :: Monad m => (a -> b -> c -> d -> m e)
          -> [a] -> [b] -> [c] -> [d] -> m [e]
83 84 85 86
{-# INLINE zipWith4M #-}
-- Inline so that fusion with 'zipWith5' and 'sequenceA' has a chance to fire.
-- See  Note [Inline @zipWithNM@ functions] above.
zipWith4M f xs ys ws zs = sequenceA (zipWith4 f xs ys ws zs)
87

eir@cis.upenn.edu's avatar
eir@cis.upenn.edu committed
88 89
zipWithAndUnzipM :: Monad m
                 => (a -> b -> m (c, d)) -> [a] -> [b] -> m ([c], [d])
90
{-# INLINABLE zipWithAndUnzipM #-}
91 92
-- See Note [flatten_many performance] in TcFlatten for why this
-- pragma is essential.
eir@cis.upenn.edu's avatar
eir@cis.upenn.edu committed
93 94 95 96 97 98
zipWithAndUnzipM f (x:xs) (y:ys)
  = do { (c, d) <- f x y
       ; (cs, ds) <- zipWithAndUnzipM f xs ys
       ; return (c:cs, d:ds) }
zipWithAndUnzipM _ _ _ = return ([], [])

99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121
{-

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.

-}

122 123
-- | mapAndUnzipM for triples
mapAndUnzip3M :: Monad m => (a -> m (b,c,d)) -> [a] -> m ([b],[c],[d])
124 125 126 127
{-# INLINE mapAndUnzip3M #-}
-- Inline so that fusion with 'unzip3' and 'traverse' has a chance to fire.
-- See Note [Inline @mapAndUnzipNM@ functions] above.
mapAndUnzip3M f xs =  unzip3 <$> traverse f xs
128 129

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

135
mapAndUnzip5M :: Monad m => (a -> m (b,c,d,e,f)) -> [a] -> m ([b],[c],[d],[e],[f])
136 137 138 139
{-# INLINE mapAndUnzip5M #-}
-- Inline so that fusion with 'unzip5' and 'traverse' has a chance to fire.
-- See Note [Inline @mapAndUnzipNM@ functions] above.
mapAndUnzip5M f xs =  unzip5 <$> traverse f xs
140

141 142
-- | Monadic version of mapAccumL
mapAccumLM :: Monad m
143
            => (acc -> x -> m (acc, y)) -- ^ combining function
144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161
            -> acc                      -- ^ initial state
            -> [x]                      -- ^ inputs
            -> m (acc, [y])             -- ^ final state, outputs
mapAccumLM _ s []     = return (s, [])
mapAccumLM f s (x:xs) = do
    (s1, x')  <- f s x
    (s2, xs') <- mapAccumLM f s1 xs
    return    (s2, x' : xs')

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

-- | Monadic version of concatMap
concatMapM :: Monad m => (a -> m [b]) -> [a] -> m [b]
concatMapM f xs = liftM concat (mapM f xs)

Simon Jakobi's avatar
Simon Jakobi committed
162 163 164 165
-- | Applicative version of mapMaybe
mapMaybeM :: Applicative m => (a -> m (Maybe b)) -> [a] -> m [b]
mapMaybeM f = foldr g (pure [])
  where g a = liftA2 (maybe id (:)) (f a)
166

167 168 169 170 171 172 173 174 175 176
-- | Monadic version of fmap
fmapMaybeM :: (Monad m) => (a -> m b) -> Maybe a -> m (Maybe b)
fmapMaybeM _ Nothing  = return Nothing
fmapMaybeM f (Just x) = f x >>= (return . Just)

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

batterseapower's avatar
batterseapower committed
177
-- | Monadic version of 'any', aborts the computation at the first @True@ value
178 179 180
anyM :: Monad m => (a -> m Bool) -> [a] -> m Bool
anyM _ []     = return False
anyM f (x:xs) = do b <- f x
181
                   if b then return True
182 183
                        else anyM f xs

batterseapower's avatar
batterseapower committed
184 185 186 187 188
-- | Monad version of 'all', aborts the computation at the first @False@ value
allM :: Monad m => (a -> m Bool) -> [a] -> m Bool
allM _ []     = return True
allM f (b:bs) = (f b) >>= (\bv -> if bv then allM f bs else return False)

189 190 191 192
-- | Monadic version of or
orM :: Monad m => m Bool -> m Bool -> m Bool
orM m1 m2 = m1 >>= \x -> if x then return True else m2

193
-- | Monadic version of foldl
194
foldlM :: (Monad m, Foldable t) => (a -> b -> m a) -> a -> t b -> m a
195 196
foldlM = foldM

197
-- | Monadic version of foldl that discards its result
198
foldlM_ :: (Monad m, Foldable t) => (a -> b -> m a) -> a -> t b -> m ()
199 200
foldlM_ = foldM_

201
-- | Monadic version of foldr
202 203
foldrM        :: (Monad m, Foldable t) => (b -> a -> m a) -> a -> t b -> m a
foldrM k z x = foldr (\x r -> r >>= k x) (pure z) x
204 205 206 207 208

-- | Monadic version of fmap specialised for Maybe
maybeMapM :: Monad m => (a -> m b) -> (Maybe a -> m (Maybe b))
maybeMapM _ Nothing  = return Nothing
maybeMapM m (Just x) = liftM Just $ m x
209 210 211 212 213

-- | Monadic version of @when@, taking the condition in the monad
whenM :: Monad m => m Bool -> m () -> m ()
whenM mb thing = do { b <- mb
                    ; when b thing }
214 215 216 217 218

-- | Monadic version of @unless@, taking the condition in the monad
unlessM :: Monad m => m Bool -> m () -> m ()
unlessM condM acc = do { cond <- condM
                       ; unless cond acc }
219 220 221 222 223

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