Commit c4a7520e authored by snoyberg's avatar snoyberg Committed by Ben Gamari

Provide an optimized replicateM_ implementation #11795

In my testing, the worker/wrapper transformation applied here
significantly decreases the number of allocations performed when using
replicateM_. Additionally, this version of the function behaves
correctly for negative numbers (namely, it will behave the same as
replicateM_ 0, which is what previous versions of base have done).

Reviewers: bgamari, simonpj, hvr, austin

Reviewed By: bgamari, simonpj, austin

Subscribers: nomeata, simonpj, mpickering, thomie

Differential Revision: https://phabricator.haskell.org/D2086

GHC Trac Issues: #11795
parent 485608d3
...@@ -80,8 +80,8 @@ import Data.Functor ( void, (<$>) ) ...@@ -80,8 +80,8 @@ import Data.Functor ( void, (<$>) )
import Data.Traversable ( forM, mapM, traverse, sequence, sequenceA ) import Data.Traversable ( forM, mapM, traverse, sequence, sequenceA )
import GHC.Base hiding ( mapM, sequence ) import GHC.Base hiding ( mapM, sequence )
import GHC.Enum ( pred )
import GHC.List ( zipWith, unzip ) import GHC.List ( zipWith, unzip )
import GHC.Num ( (-) )
-- ----------------------------------------------------------------------------- -- -----------------------------------------------------------------------------
-- Functions mandated by the Prelude -- Functions mandated by the Prelude
...@@ -169,22 +169,55 @@ foldM_ :: (Foldable t, Monad m) => (b -> a -> m b) -> b -> t a -> m () ...@@ -169,22 +169,55 @@ foldM_ :: (Foldable t, Monad m) => (b -> a -> m b) -> b -> t a -> m ()
{-# SPECIALISE foldM_ :: (a -> b -> Maybe a) -> a -> [b] -> Maybe () #-} {-# SPECIALISE foldM_ :: (a -> b -> Maybe a) -> a -> [b] -> Maybe () #-}
foldM_ f a xs = foldlM f a xs >> return () foldM_ f a xs = foldlM f a xs >> return ()
{-
Note [Worker/wrapper transform on replicateM/replicateM_
--------------------------------------------------------
The implementations of replicateM and replicateM_ both leverage the
worker/wrapper transform. The simpler implementation of replicateM_, as an
example, would be:
replicateM_ 0 _ = pure ()
replicateM_ n f = f *> replicateM_ (n - 1) f
However, the self-recrusive nature of this implementation inhibits inlining,
which means we never get to specialise to the action (`f` in the code above).
By contrast, the implementation below with a local loop makes it possible to
inline the entire definition (as hapens for foldr, for example) thereby
specialising for the particular action.
For further information, see this Trac comment, which includes side-by-side
Core.
https://ghc.haskell.org/trac/ghc/ticket/11795#comment:6
-}
-- | @'replicateM' n act@ performs the action @n@ times, -- | @'replicateM' n act@ performs the action @n@ times,
-- gathering the results. -- gathering the results.
replicateM :: (Applicative m) => Int -> m a -> m [a] replicateM :: (Applicative m) => Int -> m a -> m [a]
{-# INLINEABLE replicateM #-} {-# INLINEABLE replicateM #-}
{-# SPECIALISE replicateM :: Int -> IO a -> IO [a] #-} {-# SPECIALISE replicateM :: Int -> IO a -> IO [a] #-}
{-# SPECIALISE replicateM :: Int -> Maybe a -> Maybe [a] #-} {-# SPECIALISE replicateM :: Int -> Maybe a -> Maybe [a] #-}
replicateM 0 _ = pure [] replicateM cnt0 f =
replicateM n x = liftA2 (:) x (replicateM (pred n) x) loop cnt0
where
loop cnt
| cnt <= 0 = pure []
| otherwise = liftA2 (:) f (loop (cnt - 1))
-- | Like 'replicateM', but discards the result. -- | Like 'replicateM', but discards the result.
replicateM_ :: (Applicative m) => Int -> m a -> m () replicateM_ :: (Applicative m) => Int -> m a -> m ()
{-# INLINEABLE replicateM_ #-} {-# INLINEABLE replicateM_ #-}
{-# SPECIALISE replicateM_ :: Int -> IO a -> IO () #-} {-# SPECIALISE replicateM_ :: Int -> IO a -> IO () #-}
{-# SPECIALISE replicateM_ :: Int -> Maybe a -> Maybe () #-} {-# SPECIALISE replicateM_ :: Int -> Maybe a -> Maybe () #-}
replicateM_ 0 _ = pure () replicateM_ cnt0 f =
replicateM_ n x = x *> replicateM_ (pred n) x loop cnt0
where
loop cnt
| cnt <= 0 = pure ()
| otherwise = f *> loop (cnt - 1)
-- | The reverse of 'when'. -- | The reverse of 'when'.
unless :: (Applicative f) => Bool -> f () -> f () unless :: (Applicative f) => Bool -> f () -> f ()
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment