Provide an optimized replicateM_ implementation #11795
authorMichael Snoyman <michael@snoyman.com>
Sun, 10 Apr 2016 16:52:47 +0000 (18:52 +0200)
committerBen Gamari <ben@smart-cactus.org>
Sun, 10 Apr 2016 18:09:03 +0000 (20:09 +0200)
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

index 6957ad4..9d858bd 100644 (file)
@@ -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 ()