From de521be79f20d36d9b0d9792d8e3c44b7a584ecb Mon Sep 17 00:00:00 2001 From: Alexandre Date: Thu, 21 Feb 2019 08:19:43 +0000 Subject: [#14037] Allow fusion for zip7 and related Reviewers: bgamari, simonpj, hvr Reviewed By: simonpj Subscribers: AndreasK, simonpj, osa1, dfeuer, rwbarton, carter GHC Trac Issues: #14037 Differential Revision: https://phabricator.haskell.org/D5249 --- compiler/utils/MonadUtils.hs | 110 +++++++++++++++++++++++++++------------- libraries/base/Control/Monad.hs | 2 + libraries/base/Data/OldList.hs | 28 ++++++++++ libraries/base/GHC/List.hs | 4 ++ 4 files changed, 109 insertions(+), 35 deletions(-) diff --git a/compiler/utils/MonadUtils.hs b/compiler/utils/MonadUtils.hs index 8f40f88ba9..f4320ecb4d 100644 --- a/compiler/utils/MonadUtils.hs +++ b/compiler/utils/MonadUtils.hs @@ -34,6 +34,8 @@ import Control.Applicative import Control.Monad import Control.Monad.Fix import Control.Monad.IO.Class +import Data.Foldable (sequenceA_) +import Data.List (unzip4, unzip5, zipWith4) ------------------------------------------------------------------------------- -- Lift combinators @@ -61,32 +63,50 @@ liftIO4 = (((.).(.)).((.).(.))) liftIO -- These are used throughout the compiler ------------------------------------------------------------------------------- +{- + +Note [Inline @zipWithNM@ functions] +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +The inline principle for 'zipWith3M', 'zipWith4M' and 'zipWith3M_' is the same +as for 'zipWithM' and 'zipWithM_' in "Control.Monad", see +Note [Fusion for zipN/zipWithN] in GHC/List.hs for more details. + +The 'zipWithM'/'zipWithM_' functions are inlined so that the `zipWith` and +`sequenceA` functions with which they are defined have an opportunity to fuse. + +Furthermore, 'zipWith3M'/'zipWith4M' and 'zipWith3M_' have been explicitly +rewritten in a non-recursive way similarly to 'zipWithM'/'zipWithM_', and for +more than just uniformity: after [D5241](https://phabricator.haskell.org/D5241) +for Trac ticket #14037, all @zipN@/@zipWithN@ functions fuse, meaning +'zipWith3M'/'zipWIth4M' and 'zipWith3M_'@ now behave like 'zipWithM' and +'zipWithM_', respectively, with regards to fusion. + +As such, since there are not any differences between 2-ary 'zipWithM'/ +'zipWithM_' and their n-ary counterparts below aside from the number of +arguments, the `INLINE` pragma should be replicated in the @zipWithNM@ +functions below as well. + +-} + zipWith3M :: Monad m => (a -> b -> c -> m d) -> [a] -> [b] -> [c] -> m [d] -zipWith3M _ [] _ _ = return [] -zipWith3M _ _ [] _ = return [] -zipWith3M _ _ _ [] = return [] -zipWith3M f (x:xs) (y:ys) (z:zs) - = do { r <- f x y z - ; rs <- zipWith3M f xs ys zs - ; return $ r:rs - } +{-# INLINE zipWith3M #-} +-- Inline so that fusion with 'zipWith3' and 'sequenceA' has a chance to fire. +-- See Note [Inline @zipWithNM@ functions] above. +zipWith3M f xs ys zs = sequenceA (zipWith3 f xs ys zs) zipWith3M_ :: Monad m => (a -> b -> c -> m d) -> [a] -> [b] -> [c] -> m () -zipWith3M_ f as bs cs = do { _ <- zipWith3M f as bs cs - ; return () } +{-# INLINE zipWith3M_ #-} +-- Inline so that fusion with 'zipWith4' and 'sequenceA' has a chance to fire. +-- See Note [Inline @zipWithNM@ functions] above. +zipWith3M_ f xs ys zs = sequenceA_ (zipWith3 f xs ys zs) zipWith4M :: Monad m => (a -> b -> c -> d -> m e) -> [a] -> [b] -> [c] -> [d] -> m [e] -zipWith4M _ [] _ _ _ = return [] -zipWith4M _ _ [] _ _ = return [] -zipWith4M _ _ _ [] _ = return [] -zipWith4M _ _ _ _ [] = return [] -zipWith4M f (x:xs) (y:ys) (z:zs) (a:as) - = do { r <- f x y z a - ; rs <- zipWith4M f xs ys zs as - ; return $ r:rs - } - +{-# INLINE zipWith4M #-} +-- Inline so that fusion with 'zipWith5' and 'sequenceA' has a chance to fire. +-- See Note [Inline @zipWithNM@ functions] above. +zipWith4M f xs ys ws zs = sequenceA (zipWith4 f xs ys ws zs) zipWithAndUnzipM :: Monad m => (a -> b -> m (c, d)) -> [a] -> [b] -> m ([c], [d]) @@ -99,27 +119,47 @@ zipWithAndUnzipM f (x:xs) (y:ys) ; return (c:cs, d:ds) } zipWithAndUnzipM _ _ _ = return ([], []) +{- + +Note [Inline @mapAndUnzipNM@ functions] +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +The inline principle is the same as 'mapAndUnzipM' in "Control.Monad". +The 'mapAndUnzipM' function is inlined so that the `unzip` and `traverse` +functions with which it is defined have an opportunity to fuse, see +Note [Inline @unzipN@ functions] in Data/OldList.hs for more details. + +Furthermore, the @mapAndUnzipNM@ functions have been explicitly rewritten in a +non-recursive way similarly to 'mapAndUnzipM', and for more than just +uniformity: after [D5249](https://phabricator.haskell.org/D5249) for Trac +ticket #14037, all @unzipN@ functions fuse, meaning 'mapAndUnzip3M', +'mapAndUnzip4M' and 'mapAndUnzip5M' now behave like 'mapAndUnzipM' with regards +to fusion. + +As such, since there are not any differences between 2-ary 'mapAndUnzipM' and +its n-ary counterparts below aside from the number of arguments, the `INLINE` +pragma should be replicated in the @mapAndUnzipNM@ functions below as well. + +-} + -- | mapAndUnzipM for triples mapAndUnzip3M :: Monad m => (a -> m (b,c,d)) -> [a] -> m ([b],[c],[d]) -mapAndUnzip3M _ [] = return ([],[],[]) -mapAndUnzip3M f (x:xs) = do - (r1, r2, r3) <- f x - (rs1, rs2, rs3) <- mapAndUnzip3M f xs - return (r1:rs1, r2:rs2, r3:rs3) +{-# INLINE mapAndUnzip3M #-} +-- Inline so that fusion with 'unzip3' and 'traverse' has a chance to fire. +-- See Note [Inline @mapAndUnzipNM@ functions] above. +mapAndUnzip3M f xs = unzip3 <$> traverse f xs mapAndUnzip4M :: Monad m => (a -> m (b,c,d,e)) -> [a] -> m ([b],[c],[d],[e]) -mapAndUnzip4M _ [] = return ([],[],[],[]) -mapAndUnzip4M f (x:xs) = do - (r1, r2, r3, r4) <- f x - (rs1, rs2, rs3, rs4) <- mapAndUnzip4M f xs - return (r1:rs1, r2:rs2, r3:rs3, r4:rs4) +{-# INLINE mapAndUnzip4M #-} +-- Inline so that fusion with 'unzip4' and 'traverse' has a chance to fire. +-- See Note [Inline @mapAndUnzipNM@ functions] above. +mapAndUnzip4M f xs = unzip4 <$> traverse f xs mapAndUnzip5M :: Monad m => (a -> m (b,c,d,e,f)) -> [a] -> m ([b],[c],[d],[e],[f]) -mapAndUnzip5M _ [] = return ([],[],[],[],[]) -mapAndUnzip5M f (x:xs) = do - (r1, r2, r3, r4, r5) <- f x - (rs1, rs2, rs3, rs4, rs5) <- mapAndUnzip5M f xs - return (r1:rs1, r2:rs2, r3:rs3, r4:rs4, r5:rs5) +{-# INLINE mapAndUnzip5M #-} +-- Inline so that fusion with 'unzip5' and 'traverse' has a chance to fire. +-- See Note [Inline @mapAndUnzipNM@ functions] above. +mapAndUnzip5M f xs = unzip5 <$> traverse f xs -- | Monadic version of mapAccumL mapAccumLM :: Monad m diff --git a/libraries/base/Control/Monad.hs b/libraries/base/Control/Monad.hs index 75bc2b2db3..fbdb99e5f4 100644 --- a/libraries/base/Control/Monad.hs +++ b/libraries/base/Control/Monad.hs @@ -191,6 +191,8 @@ forever a = let a' = a *> a' in a' -- data structures or a state monad. mapAndUnzipM :: (Applicative m) => (a -> m (b,c)) -> [a] -> m ([b], [c]) {-# INLINE mapAndUnzipM #-} +-- Inline so that fusion with 'unzip' and 'traverse' has a chance to fire. +-- See Note [Inline @unzipN@ functions] in GHC/OldList.hs. mapAndUnzipM f xs = unzip <$> traverse f xs -- | The 'zipWithM' function generalizes 'zipWith' to arbitrary applicative functors. diff --git a/libraries/base/Data/OldList.hs b/libraries/base/Data/OldList.hs index 559823c4bd..132ee14673 100644 --- a/libraries/base/Data/OldList.hs +++ b/libraries/base/Data/OldList.hs @@ -929,8 +929,27 @@ foldr7_left _ z _ _ _ _ _ _ _ _ = z #-} +{- + +Note [Inline @unzipN@ functions] +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +The inline principle for @unzip{4,5,6,7}@ is the same as 'unzip'/'unzip3' in +"GHC.List". +The 'unzip'/'unzip3' functions are inlined so that the `foldr` with which they +are defined has an opportunity to fuse. + +As such, since there are not any differences between 2/3-ary 'unzip' and its +n-ary counterparts below aside from the number of arguments, the `INLINE` +pragma should be replicated in the @unzipN@ functions below as well. + +-} + -- | The 'unzip4' function takes a list of quadruples and returns four -- lists, analogous to 'unzip'. +{-# INLINE unzip4 #-} +-- Inline so that fusion with `foldr` has an opportunity to fire. +-- See Note [Inline @unzipN@ functions] above. unzip4 :: [(a,b,c,d)] -> ([a],[b],[c],[d]) unzip4 = foldr (\(a,b,c,d) ~(as,bs,cs,ds) -> (a:as,b:bs,c:cs,d:ds)) @@ -938,6 +957,9 @@ unzip4 = foldr (\(a,b,c,d) ~(as,bs,cs,ds) -> -- | The 'unzip5' function takes a list of five-tuples and returns five -- lists, analogous to 'unzip'. +{-# INLINE unzip5 #-} +-- Inline so that fusion with `foldr` has an opportunity to fire. +-- See Note [Inline @unzipN@ functions] above. unzip5 :: [(a,b,c,d,e)] -> ([a],[b],[c],[d],[e]) unzip5 = foldr (\(a,b,c,d,e) ~(as,bs,cs,ds,es) -> (a:as,b:bs,c:cs,d:ds,e:es)) @@ -945,6 +967,9 @@ unzip5 = foldr (\(a,b,c,d,e) ~(as,bs,cs,ds,es) -> -- | The 'unzip6' function takes a list of six-tuples and returns six -- lists, analogous to 'unzip'. +{-# INLINE unzip6 #-} +-- Inline so that fusion with `foldr` has an opportunity to fire. +-- See Note [Inline @unzipN@ functions] above. unzip6 :: [(a,b,c,d,e,f)] -> ([a],[b],[c],[d],[e],[f]) unzip6 = foldr (\(a,b,c,d,e,f) ~(as,bs,cs,ds,es,fs) -> (a:as,b:bs,c:cs,d:ds,e:es,f:fs)) @@ -952,6 +977,9 @@ unzip6 = foldr (\(a,b,c,d,e,f) ~(as,bs,cs,ds,es,fs) -> -- | The 'unzip7' function takes a list of seven-tuples and returns -- seven lists, analogous to 'unzip'. +{-# INLINE unzip7 #-} +-- Inline so that fusion with `foldr` has an opportunity to fire. +-- See Note [Inline @unzipN@ functions] above. unzip7 :: [(a,b,c,d,e,f,g)] -> ([a],[b],[c],[d],[e],[f],[g]) unzip7 = foldr (\(a,b,c,d,e,f,g) ~(as,bs,cs,ds,es,fs,gs) -> (a:as,b:bs,c:cs,d:ds,e:es,f:fs,g:gs)) diff --git a/libraries/base/GHC/List.hs b/libraries/base/GHC/List.hs index 8f03ce3a08..d9b32ea9df 100644 --- a/libraries/base/GHC/List.hs +++ b/libraries/base/GHC/List.hs @@ -1128,12 +1128,16 @@ zipWith3FB cons func = \a b c r -> (func a b c) `cons` r -- and a list of second components. unzip :: [(a,b)] -> ([a],[b]) {-# INLINE unzip #-} +-- Inline so that fusion `foldr` has an opportunity to fire. +-- See Note [Inline @unzipN@ functions] in GHC/OldList.hs. unzip = foldr (\(a,b) ~(as,bs) -> (a:as,b:bs)) ([],[]) -- | The 'unzip3' function takes a list of triples and returns three -- lists, analogous to 'unzip'. unzip3 :: [(a,b,c)] -> ([a],[b],[c]) {-# INLINE unzip3 #-} +-- Inline so that fusion `foldr` has an opportunity to fire. +-- See Note [Inline @unzipN@ functions] in GHC/OldList.hs. unzip3 = foldr (\(a,b,c) ~(as,bs,cs) -> (a:as,b:bs,c:cs)) ([],[],[]) -- cgit v1.2.1