Provide an optimized replicateM_ implementation #11795
[ghc.git] / libraries / base / Control / Monad.hs
1 {-# LANGUAGE Trustworthy #-}
2 {-# LANGUAGE NoImplicitPrelude #-}
3
4 -----------------------------------------------------------------------------
5 -- |
6 -- Module : Control.Monad
7 -- Copyright : (c) The University of Glasgow 2001
8 -- License : BSD-style (see the file libraries/base/LICENSE)
9 --
10 -- Maintainer : libraries@haskell.org
11 -- Stability : provisional
12 -- Portability : portable
13 --
14 -- The 'Functor', 'Monad' and 'MonadPlus' classes,
15 -- with some useful operations on monads.
16
17 module Control.Monad
18 (
19 -- * Functor and monad classes
20
21 Functor(fmap)
22 , Monad((>>=), (>>), return, fail)
23 , MonadPlus(mzero, mplus)
24 -- * Functions
25
26 -- ** Naming conventions
27 -- $naming
28
29 -- ** Basic @Monad@ functions
30
31 , mapM
32 , mapM_
33 , forM
34 , forM_
35 , sequence
36 , sequence_
37 , (=<<)
38 , (>=>)
39 , (<=<)
40 , forever
41 , void
42
43 -- ** Generalisations of list functions
44
45 , join
46 , msum
47 , mfilter
48 , filterM
49 , mapAndUnzipM
50 , zipWithM
51 , zipWithM_
52 , foldM
53 , foldM_
54 , replicateM
55 , replicateM_
56
57 -- ** Conditional execution of monadic expressions
58
59 , guard
60 , when
61 , unless
62
63 -- ** Monadic lifting operators
64
65 , liftM
66 , liftM2
67 , liftM3
68 , liftM4
69 , liftM5
70
71 , ap
72
73 -- ** Strict monadic functions
74
75 , (<$!>)
76 ) where
77
78 import Data.Foldable ( Foldable, sequence_, sequenceA_, msum, mapM_, foldlM, forM_ )
79 import Data.Functor ( void, (<$>) )
80 import Data.Traversable ( forM, mapM, traverse, sequence, sequenceA )
81
82 import GHC.Base hiding ( mapM, sequence )
83 import GHC.List ( zipWith, unzip )
84 import GHC.Num ( (-) )
85
86 -- -----------------------------------------------------------------------------
87 -- Functions mandated by the Prelude
88
89 -- | @'guard' b@ is @'pure' ()@ if @b@ is 'True',
90 -- and 'empty' if @b@ is 'False'.
91 guard :: (Alternative f) => Bool -> f ()
92 guard True = pure ()
93 guard False = empty
94
95 -- | This generalizes the list-based 'filter' function.
96
97 {-# INLINE filterM #-}
98 filterM :: (Applicative m) => (a -> m Bool) -> [a] -> m [a]
99 filterM p = foldr (\ x -> liftA2 (\ flg -> if flg then (x:) else id) (p x)) (pure [])
100
101 infixr 1 <=<, >=>
102
103 -- | Left-to-right Kleisli composition of monads.
104 (>=>) :: Monad m => (a -> m b) -> (b -> m c) -> (a -> m c)
105 f >=> g = \x -> f x >>= g
106
107 -- | Right-to-left Kleisli composition of monads. @('>=>')@, with the arguments flipped
108 (<=<) :: Monad m => (b -> m c) -> (a -> m b) -> (a -> m c)
109 (<=<) = flip (>=>)
110
111 -- | @'forever' act@ repeats the action infinitely.
112 forever :: (Applicative f) => f a -> f b
113 {-# INLINE forever #-}
114 forever a = let a' = a *> a' in a'
115 -- Use explicit sharing here, as it is prevents a space leak regardless of
116 -- optimizations.
117
118 -- -----------------------------------------------------------------------------
119 -- Other monad functions
120
121 -- | The 'mapAndUnzipM' function maps its first argument over a list, returning
122 -- the result as a pair of lists. This function is mainly used with complicated
123 -- data structures or a state-transforming monad.
124 mapAndUnzipM :: (Applicative m) => (a -> m (b,c)) -> [a] -> m ([b], [c])
125 {-# INLINE mapAndUnzipM #-}
126 mapAndUnzipM f xs = unzip <$> traverse f xs
127
128 -- | The 'zipWithM' function generalizes 'zipWith' to arbitrary applicative functors.
129 zipWithM :: (Applicative m) => (a -> b -> m c) -> [a] -> [b] -> m [c]
130 {-# INLINE zipWithM #-}
131 zipWithM f xs ys = sequenceA (zipWith f xs ys)
132
133 -- | 'zipWithM_' is the extension of 'zipWithM' which ignores the final result.
134 zipWithM_ :: (Applicative m) => (a -> b -> m c) -> [a] -> [b] -> m ()
135 {-# INLINE zipWithM_ #-}
136 zipWithM_ f xs ys = sequenceA_ (zipWith f xs ys)
137
138 {- | The 'foldM' function is analogous to 'foldl', except that its result is
139 encapsulated in a monad. Note that 'foldM' works from left-to-right over
140 the list arguments. This could be an issue where @('>>')@ and the `folded
141 function' are not commutative.
142
143
144 > foldM f a1 [x1, x2, ..., xm]
145
146 ==
147
148 > do
149 > a2 <- f a1 x1
150 > a3 <- f a2 x2
151 > ...
152 > f am xm
153
154 If right-to-left evaluation is required, the input list should be reversed.
155
156 Note: 'foldM' is the same as 'foldlM'
157 -}
158
159 foldM :: (Foldable t, Monad m) => (b -> a -> m b) -> b -> t a -> m b
160 {-# INLINEABLE foldM #-}
161 {-# SPECIALISE foldM :: (a -> b -> IO a) -> a -> [b] -> IO a #-}
162 {-# SPECIALISE foldM :: (a -> b -> Maybe a) -> a -> [b] -> Maybe a #-}
163 foldM = foldlM
164
165 -- | Like 'foldM', but discards the result.
166 foldM_ :: (Foldable t, Monad m) => (b -> a -> m b) -> b -> t a -> m ()
167 {-# INLINEABLE foldM_ #-}
168 {-# SPECIALISE foldM_ :: (a -> b -> IO a) -> a -> [b] -> IO () #-}
169 {-# SPECIALISE foldM_ :: (a -> b -> Maybe a) -> a -> [b] -> Maybe () #-}
170 foldM_ f a xs = foldlM f a xs >> return ()
171
172 {-
173 Note [Worker/wrapper transform on replicateM/replicateM_
174 --------------------------------------------------------
175
176 The implementations of replicateM and replicateM_ both leverage the
177 worker/wrapper transform. The simpler implementation of replicateM_, as an
178 example, would be:
179
180 replicateM_ 0 _ = pure ()
181 replicateM_ n f = f *> replicateM_ (n - 1) f
182
183 However, the self-recrusive nature of this implementation inhibits inlining,
184 which means we never get to specialise to the action (`f` in the code above).
185 By contrast, the implementation below with a local loop makes it possible to
186 inline the entire definition (as hapens for foldr, for example) thereby
187 specialising for the particular action.
188
189 For further information, see this Trac comment, which includes side-by-side
190 Core.
191
192 https://ghc.haskell.org/trac/ghc/ticket/11795#comment:6
193
194 -}
195
196 -- | @'replicateM' n act@ performs the action @n@ times,
197 -- gathering the results.
198 replicateM :: (Applicative m) => Int -> m a -> m [a]
199 {-# INLINEABLE replicateM #-}
200 {-# SPECIALISE replicateM :: Int -> IO a -> IO [a] #-}
201 {-# SPECIALISE replicateM :: Int -> Maybe a -> Maybe [a] #-}
202 replicateM cnt0 f =
203 loop cnt0
204 where
205 loop cnt
206 | cnt <= 0 = pure []
207 | otherwise = liftA2 (:) f (loop (cnt - 1))
208
209 -- | Like 'replicateM', but discards the result.
210 replicateM_ :: (Applicative m) => Int -> m a -> m ()
211 {-# INLINEABLE replicateM_ #-}
212 {-# SPECIALISE replicateM_ :: Int -> IO a -> IO () #-}
213 {-# SPECIALISE replicateM_ :: Int -> Maybe a -> Maybe () #-}
214 replicateM_ cnt0 f =
215 loop cnt0
216 where
217 loop cnt
218 | cnt <= 0 = pure ()
219 | otherwise = f *> loop (cnt - 1)
220
221
222 -- | The reverse of 'when'.
223 unless :: (Applicative f) => Bool -> f () -> f ()
224 {-# INLINEABLE unless #-}
225 {-# SPECIALISE unless :: Bool -> IO () -> IO () #-}
226 {-# SPECIALISE unless :: Bool -> Maybe () -> Maybe () #-}
227 unless p s = if p then pure () else s
228
229 infixl 4 <$!>
230
231 -- | Strict version of 'Data.Functor.<$>'.
232 --
233 -- @since 4.8.0.0
234 (<$!>) :: Monad m => (a -> b) -> m a -> m b
235 {-# INLINE (<$!>) #-}
236 f <$!> m = do
237 x <- m
238 let z = f x
239 z `seq` return z
240
241
242 -- -----------------------------------------------------------------------------
243 -- Other MonadPlus functions
244
245 -- | Direct 'MonadPlus' equivalent of 'filter'
246 -- @'filter'@ = @(mfilter:: (a -> Bool) -> [a] -> [a]@
247 -- applicable to any 'MonadPlus', for example
248 -- @mfilter odd (Just 1) == Just 1@
249 -- @mfilter odd (Just 2) == Nothing@
250
251 mfilter :: (MonadPlus m) => (a -> Bool) -> m a -> m a
252 {-# INLINEABLE mfilter #-}
253 mfilter p ma = do
254 a <- ma
255 if p a then return a else mzero
256
257 {- $naming
258
259 The functions in this library use the following naming conventions:
260
261 * A postfix \'@M@\' always stands for a function in the Kleisli category:
262 The monad type constructor @m@ is added to function results
263 (modulo currying) and nowhere else. So, for example,
264
265 > filter :: (a -> Bool) -> [a] -> [a]
266 > filterM :: (Monad m) => (a -> m Bool) -> [a] -> m [a]
267
268 * A postfix \'@_@\' changes the result type from @(m a)@ to @(m ())@.
269 Thus, for example:
270
271 > sequence :: Monad m => [m a] -> m [a]
272 > sequence_ :: Monad m => [m a] -> m ()
273
274 * A prefix \'@m@\' generalizes an existing function to a monadic form.
275 Thus, for example:
276
277 > sum :: Num a => [a] -> a
278 > msum :: MonadPlus m => [m a] -> m a
279
280 -}