[commit: ghc] master: Provide an optimized replicateM_ implementation #11795 (c4a7520)

[git at git.haskell.org]

Repository : ssh://git@git.haskell.org/ghc

On branch  : master
Link       : http://ghc.haskell.org/trac/ghc/changeset/c4a7520ef3a0b5e0e33d66ae1d628af93e0d7590/ghc

>---------------------------------------------------------------

commit c4a7520ef3a0b5e0e33d66ae1d628af93e0d7590
Author: Michael Snoyman <michael at snoyman.com>
Date:   Sun Apr 10 18:52:47 2016 +0200

    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


>---------------------------------------------------------------

c4a7520ef3a0b5e0e33d66ae1d628af93e0d7590
 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 ()