From c4a7520ef3a0b5e0e33d66ae1d628af93e0d7590 Mon Sep 17 00:00:00 2001
From: Michael Snoyman
Date: Sun, 10 Apr 2016 18:52:47 +0200
Subject: [PATCH] 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
---
libraries/base/Control/Monad.hs | 43 ++++++++++++++++++++++++++++++++++++-----
1 file changed, 38 insertions(+), 5 deletions(-)
diff --git a/libraries/base/Control/Monad.hs b/libraries/base/Control/Monad.hs
index 6957ad4..9d858bd 100644
--- a/libraries/base/Control/Monad.hs
+++ b/libraries/base/Control/Monad.hs
@@ -80,8 +80,8 @@ import Data.Functor ( void, (<$>) )
import Data.Traversable ( forM, mapM, traverse, sequence, sequenceA )
import GHC.Base hiding ( mapM, sequence )
-import GHC.Enum ( pred )
import GHC.List ( zipWith, unzip )
+import GHC.Num ( (-) )
-- -----------------------------------------------------------------------------
-- Functions mandated by the Prelude
@@ -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 () #-}
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,
-- gathering the results.
replicateM :: (Applicative m) => Int -> m a -> m [a]
{-# INLINEABLE replicateM #-}
{-# SPECIALISE replicateM :: Int -> IO a -> IO [a] #-}
{-# SPECIALISE replicateM :: Int -> Maybe a -> Maybe [a] #-}
-replicateM 0 _ = pure []
-replicateM n x = liftA2 (:) x (replicateM (pred n) x)
+replicateM cnt0 f =
+ loop cnt0
+ where
+ loop cnt
+ | cnt <= 0 = pure []
+ | otherwise = liftA2 (:) f (loop (cnt - 1))
-- | Like 'replicateM', but discards the result.
replicateM_ :: (Applicative m) => Int -> m a -> m ()
{-# INLINEABLE replicateM_ #-}
{-# SPECIALISE replicateM_ :: Int -> IO a -> IO () #-}
{-# SPECIALISE replicateM_ :: Int -> Maybe a -> Maybe () #-}
-replicateM_ 0 _ = pure ()
-replicateM_ n x = x *> replicateM_ (pred n) x
+replicateM_ cnt0 f =
+ loop cnt0
+ where
+ loop cnt
+ | cnt <= 0 = pure ()
+ | otherwise = f *> loop (cnt - 1)
+
-- | The reverse of 'when'.
unless :: (Applicative f) => Bool -> f () -> f ()
--
1.9.1