Add fusion rules for the zipWith functions in base (#15263)

This patch will allow `zip3` and `zipWith3` in `GHC.List` as well
as `zipWith4`, `zipWith5`, `zipWith6` and `zipWith7` in `Data.OldList` to fuse.

These rules are kept in a similar style as the rules for `zip` and `zipWith`.

Added a corresponding test case.

Test Plan: validate

Reviewers: hvr, bgamari, simonpj

Reviewed By: simonpj

Subscribers: simonpj, rockbmb, rwbarton, carter

GHC Trac Issues: #15263

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

2 files changed, 193 insertions, 0 deletions

diff --git a/libraries/base/Data/OldList.hs b/libraries/base/Data/OldList.hs
index ee2dfac982..820313599a 100644
--- a/libraries/base/Data/OldList.hs
+++ b/libraries/base/Data/OldList.hs
@@ -728,22 +728,34 @@ genericReplicate n x    =  genericTake n (repeat x)
 
 -- | The 'zip4' function takes four lists and returns a list of
 -- quadruples, analogous to 'zip'.
+-- It is capable of list fusion, but it is restricted to its
+-- first list argument and its resulting list.
+{-# INLINE zip4 #-}
 zip4                    :: [a] -> [b] -> [c] -> [d] -> [(a,b,c,d)]
 zip4                    =  zipWith4 (,,,)
 
 -- | The 'zip5' function takes five lists and returns a list of
 -- five-tuples, analogous to 'zip'.
+-- It is capable of list fusion, but it is restricted to its
+-- first list argument and its resulting list.
+{-# INLINE zip5 #-}
 zip5                    :: [a] -> [b] -> [c] -> [d] -> [e] -> [(a,b,c,d,e)]
 zip5                    =  zipWith5 (,,,,)
 
 -- | The 'zip6' function takes six lists and returns a list of six-tuples,
 -- analogous to 'zip'.
+-- It is capable of list fusion, but it is restricted to its
+-- first list argument and its resulting list.
+{-# INLINE zip6 #-}
 zip6                    :: [a] -> [b] -> [c] -> [d] -> [e] -> [f] ->
                               [(a,b,c,d,e,f)]
 zip6                    =  zipWith6 (,,,,,)
 
 -- | The 'zip7' function takes seven lists and returns a list of
 -- seven-tuples, analogous to 'zip'.
+-- It is capable of list fusion, but it is restricted to its
+-- first list argument and its resulting list.
+{-# INLINE zip7 #-}
 zip7                    :: [a] -> [b] -> [c] -> [d] -> [e] -> [f] ->
                               [g] -> [(a,b,c,d,e,f,g)]
 zip7                    =  zipWith7 (,,,,,,)
@@ -751,6 +763,9 @@ zip7                    =  zipWith7 (,,,,,,)
 -- | The 'zipWith4' function takes a function which combines four
 -- elements, as well as four lists and returns a list of their point-wise
 -- combination, analogous to 'zipWith'.
+-- It is capable of list fusion, but it is restricted to its
+-- first list argument and its resulting list.
+{-# NOINLINE [1] zipWith4 #-}
 zipWith4                :: (a->b->c->d->e) -> [a]->[b]->[c]->[d]->[e]
 zipWith4 z (a:as) (b:bs) (c:cs) (d:ds)
                         =  z a b c d : zipWith4 z as bs cs ds
@@ -759,6 +774,9 @@ zipWith4 _ _ _ _ _      =  []
 -- | The 'zipWith5' function takes a function which combines five
 -- elements, as well as five lists and returns a list of their point-wise
 -- combination, analogous to 'zipWith'.
+-- It is capable of list fusion, but it is restricted to its
+-- first list argument and its resulting list.
+{-# NOINLINE [1] zipWith5 #-}
 zipWith5                :: (a->b->c->d->e->f) ->
                            [a]->[b]->[c]->[d]->[e]->[f]
 zipWith5 z (a:as) (b:bs) (c:cs) (d:ds) (e:es)
@@ -768,6 +786,9 @@ zipWith5 _ _ _ _ _ _    = []
 -- | The 'zipWith6' function takes a function which combines six
 -- elements, as well as six lists and returns a list of their point-wise
 -- combination, analogous to 'zipWith'.
+-- It is capable of list fusion, but it is restricted to its
+-- first list argument and its resulting list.
+{-# NOINLINE [1] zipWith6 #-}
 zipWith6                :: (a->b->c->d->e->f->g) ->
                            [a]->[b]->[c]->[d]->[e]->[f]->[g]
 zipWith6 z (a:as) (b:bs) (c:cs) (d:ds) (e:es) (f:fs)
@@ -777,12 +798,133 @@ zipWith6 _ _ _ _ _ _ _  = []
 -- | The 'zipWith7' function takes a function which combines seven
 -- elements, as well as seven lists and returns a list of their point-wise
 -- combination, analogous to 'zipWith'.
+-- It is capable of list fusion, but it is restricted to its
+-- first list argument and its resulting list.
+{-# NOINLINE [1] zipWith7 #-}
 zipWith7                :: (a->b->c->d->e->f->g->h) ->
                            [a]->[b]->[c]->[d]->[e]->[f]->[g]->[h]
 zipWith7 z (a:as) (b:bs) (c:cs) (d:ds) (e:es) (f:fs) (g:gs)
                    =  z a b c d e f g : zipWith7 z as bs cs ds es fs gs
 zipWith7 _ _ _ _ _ _ _ _ = []
 
+{-
+Functions and rules for fusion of zipWith4, zipWith5, zipWith6 and zipWith7.
+The principle is the same as for zip and zipWith in GHC.List:
+Turn zipWithX into a version in which the first argument and the result
+can be fused. Turn it back into the original function if no fusion happens.
+-}
+
+{-# INLINE [0] zipWith4FB #-} -- See Note [Inline FB functions]
+zipWith4FB :: (e->xs->xs') -> (a->b->c->d->e) ->
+              a->b->c->d->xs->xs'
+zipWith4FB cons func = \a b c d r -> (func a b c d) `cons` r
+
+{-# INLINE [0] zipWith5FB #-} -- See Note [Inline FB functions]
+zipWith5FB :: (f->xs->xs') -> (a->b->c->d->e->f) ->
+              a->b->c->d->e->xs->xs'
+zipWith5FB cons func = \a b c d e r -> (func a b c d e) `cons` r
+
+{-# INLINE [0] zipWith6FB #-} -- See Note [Inline FB functions]
+zipWith6FB :: (g->xs->xs') -> (a->b->c->d->e->f->g) ->
+              a->b->c->d->e->f->xs->xs'
+zipWith6FB cons func = \a b c d e f r -> (func a b c d e f) `cons` r
+
+{-# INLINE [0] zipWith7FB #-} -- See Note [Inline FB functions]
+zipWith7FB :: (h->xs->xs') -> (a->b->c->d->e->f->g->h) ->
+              a->b->c->d->e->f->g->xs->xs'
+zipWith7FB cons func = \a b c d e f g r -> (func a b c d e f g) `cons` r
+
+{-# INLINE [0] foldr4 #-}
+foldr4 :: (a->b->c->d->e->e) ->
+          e->[a]->[b]->[c]->[d]->e
+foldr4 k z = go
+  where
+    go (a:as) (b:bs) (c:cs) (d:ds) = k a b c d (go as bs cs ds)
+    go _      _      _      _      = z
+
+{-# INLINE [0] foldr5 #-}
+foldr5 :: (a->b->c->d->e->f->f) ->
+          f->[a]->[b]->[c]->[d]->[e]->f
+foldr5 k z = go
+  where
+    go (a:as) (b:bs) (c:cs) (d:ds) (e:es) = k a b c d e (go as bs cs ds es)
+    go _      _      _      _      _      = z
+
+{-# INLINE [0] foldr6 #-}
+foldr6 :: (a->b->c->d->e->f->g->g) ->
+          g->[a]->[b]->[c]->[d]->[e]->[f]->g
+foldr6 k z = go
+  where
+    go (a:as) (b:bs) (c:cs) (d:ds) (e:es) (f:fs) = k a b c d e f (
+        go as bs cs ds es fs)
+    go _      _      _      _      _      _      = z
+
+{-# INLINE [0] foldr7 #-}
+foldr7 :: (a->b->c->d->e->f->g->h->h) ->
+          h->[a]->[b]->[c]->[d]->[e]->[f]->[g]->h
+foldr7 k z = go
+  where
+    go (a:as) (b:bs) (c:cs) (d:ds) (e:es) (f:fs) (g:gs) = k a b c d e f g (
+        go as bs cs ds es fs gs)
+    go _      _      _      _      _      _      _      = z
+
+foldr4_left :: (a->b->c->d->e->f)->
+               f->a->([b]->[c]->[d]->e)->
+               [b]->[c]->[d]->f
+foldr4_left k _z a r (b:bs) (c:cs) (d:ds) = k a b c d (r bs cs ds)
+foldr4_left _  z _ _ _      _      _      = z
+
+foldr5_left :: (a->b->c->d->e->f->g)->
+               g->a->([b]->[c]->[d]->[e]->f)->
+               [b]->[c]->[d]->[e]->g
+foldr5_left k _z a r (b:bs) (c:cs) (d:ds) (e:es) = k a b c d e (r bs cs ds es)
+foldr5_left _  z _ _ _      _      _      _      = z
+
+foldr6_left :: (a->b->c->d->e->f->g->h)->
+               h->a->([b]->[c]->[d]->[e]->[f]->g)->
+               [b]->[c]->[d]->[e]->[f]->h
+foldr6_left k _z a r (b:bs) (c:cs) (d:ds) (e:es) (f:fs) =
+    k a b c d e f (r bs cs ds es fs)
+foldr6_left _  z _ _ _      _      _      _      _      = z
+
+foldr7_left :: (a->b->c->d->e->f->g->h->i)->
+               i->a->([b]->[c]->[d]->[e]->[f]->[g]->h)->
+               [b]->[c]->[d]->[e]->[f]->[g]->i
+foldr7_left k _z a r (b:bs) (c:cs) (d:ds) (e:es) (f:fs) (g:gs) =
+    k a b c d e f g (r bs cs ds es fs gs)
+foldr7_left _  z _ _ _      _      _      _      _      _      = z
+
+{-# RULES
+
+"foldr4/left"   forall k z (g::forall b.(a->b->b)->b->b).
+                  foldr4 k z (build g) = g (foldr4_left k z) (\_ _ _ -> z)
+"foldr5/left"   forall k z (g::forall b.(a->b->b)->b->b).
+                  foldr5 k z (build g) = g (foldr5_left k z) (\_ _ _ _ -> z)
+"foldr6/left"   forall k z (g::forall b.(a->b->b)->b->b).
+                  foldr6 k z (build g) = g (foldr6_left k z) (\_ _ _ _ _ -> z)
+"foldr7/left"   forall k z (g::forall b.(a->b->b)->b->b).
+                  foldr7 k z (build g) = g (foldr7_left k z) (\_ _ _ _ _ _ -> z)
+
+"zipWith4" [~1] forall f as bs cs ds.
+                  zipWith4 f as bs cs ds = build (\c n ->
+                        foldr4 (zipWith4FB c f) n as bs cs ds)
+"zipWith5" [~1] forall f as bs cs ds es.
+                  zipWith5 f as bs cs ds es = build (\c n ->
+                        foldr5 (zipWith5FB c f) n as bs cs ds es)
+"zipWith6" [~1] forall f as bs cs ds es fs.
+                  zipWith6 f as bs cs ds es fs = build (\c n ->
+                        foldr6 (zipWith6FB c f) n as bs cs ds es fs)
+"zipWith7" [~1] forall f as bs cs ds es fs gs.
+                  zipWith7 f as bs cs ds es fs gs = build (\c n ->
+                        foldr7 (zipWith7FB c f) n as bs cs ds es fs gs)
+
+"zipWith4List"  [1]  forall f.   foldr4 (zipWith4FB (:) f) [] = zipWith4 f
+"zipWith5List"  [1]  forall f.   foldr5 (zipWith5FB (:) f) [] = zipWith5 f
+"zipWith6List"  [1]  forall f.   foldr6 (zipWith6FB (:) f) [] = zipWith6 f
+"zipWith7List"  [1]  forall f.   foldr7 (zipWith7FB (:) f) [] = zipWith7 f
+
+ #-}
+
 -- | The 'unzip4' function takes a list of quadruples and returns four
 -- lists, analogous to 'unzip'.
 unzip4                  :: [(a,b,c,d)] -> ([a],[b],[c],[d])
diff --git a/libraries/base/GHC/List.hs b/libraries/base/GHC/List.hs
index 92b5952cbe..63144ce291 100644
--- a/libraries/base/GHC/List.hs
+++ b/libraries/base/GHC/List.hs
@@ -926,6 +926,28 @@ foldr2_left  k _z  x  r (y:ys) = k x y (r ys)
 "foldr2/left"   forall k z ys (g::forall b.(a->b->b)->b->b) .
                   foldr2 k z (build g) ys = g (foldr2_left  k z) (\_ -> z) ys
  #-}
+
+foldr3 :: (a -> b -> c -> d -> d) -> d -> [a] -> [b] -> [c] -> d
+foldr3 k z = go
+  where
+    go  []    _      _      = z
+    go  _     []     _      = z
+    go  _     _      []     = z
+    go (a:as) (b:bs) (c:cs) = k a b c (go as bs cs)
+{-# INLINE [0] foldr3 #-}
+
+
+foldr3_left :: (a -> b -> c -> d -> e) -> e -> a ->
+               ([b] -> [c] -> d) -> [b] -> [c] -> e
+foldr3_left k _z a r (b:bs) (c:cs) = k a b c (r bs cs)
+foldr3_left _  z _ _  _      _     = z
+
+-- foldr3 k n xs ys zs = foldr (foldr3_left k n) (\_ _ -> n) xs ys zs
+{-# RULES
+"foldr3/left"   forall k z (g::forall b.(a->b->b)->b->b).
+                  foldr3 k z (build g) = g (foldr3_left k z) (\_ _ -> z)
+ #-}
+
 -- There used to be a foldr2/right rule, allowing foldr2 to fuse with a build
 -- form on the right. However, this causes trouble if the right list ends in
 -- a bottom that is only avoided by the left list ending at that spot. That is,
@@ -959,6 +981,9 @@ foldr2_left  k _z  x  r (y:ys) = k x y (r ys)
 --
 -- > zip [] _|_ = []
 -- > zip _|_ [] = _|_
+--
+-- 'zip' is capable of list fusion, but it is restricted to its
+-- first list argument and its resulting list.
 {-# NOINLINE [1] zip #-}
 zip :: [a] -> [b] -> [(a,b)]
 zip []     _bs    = []
@@ -977,12 +1002,23 @@ zipFB c = \x y r -> (x,y) `c` r
 ----------------------------------------------
 -- | 'zip3' takes three lists and returns a list of triples, analogous to
 -- 'zip'.
+-- It is capable of list fusion, but it is restricted to its
+-- first list argument and its resulting list.
+{-# NOINLINE [1] zip3 #-}
 zip3 :: [a] -> [b] -> [c] -> [(a,b,c)]
 -- Specification
 -- zip3 =  zipWith3 (,,)
 zip3 (a:as) (b:bs) (c:cs) = (a,b,c) : zip3 as bs cs
 zip3 _      _      _      = []
 
+{-# INLINE [0] zip3FB #-} -- See Note [Inline FB functions]
+zip3FB :: ((a,b,c) -> xs -> xs') -> a -> b -> c -> xs -> xs'
+zip3FB cons = \a b c r -> (a,b,c) `cons` r
+
+{-# RULES
+"zip3"       [~1] forall as bs cs. zip3 as bs cs = build (\c n -> foldr3 (zip3FB c) n as bs cs)
+"zip3List"   [1]          foldr3 (zip3FB (:)) [] = zip3
+ #-}
 
 -- The zipWith family generalises the zip family by zipping with the
 -- function given as the first argument, instead of a tupling function.
@@ -996,6 +1032,9 @@ zip3 _      _      _      = []
 -- 'zipWith' is right-lazy:
 --
 -- > zipWith f [] _|_ = []
+--
+-- 'zipWith' is capable of list fusion, but it is restricted to its
+-- first list argument and its resulting list.
 {-# NOINLINE [1] zipWith #-}
 zipWith :: (a->b->c) -> [a]->[b]->[c]
 zipWith f = go
@@ -1018,12 +1057,24 @@ zipWithFB c f = \x y r -> (x `f` y) `c` r
 -- | The 'zipWith3' function takes a function which combines three
 -- elements, as well as three lists and returns a list of their point-wise
 -- combination, analogous to 'zipWith'.
+-- It is capable of list fusion, but it is restricted to its
+-- first list argument and its resulting list.
+{-# NOINLINE [1] zipWith3 #-}
 zipWith3                :: (a->b->c->d) -> [a]->[b]->[c]->[d]
 zipWith3 z = go
   where
     go (a:as) (b:bs) (c:cs) = z a b c : go as bs cs
     go _ _ _                = []
 
+{-# INLINE [0] zipWith3FB #-} -- See Note [Inline FB functions]
+zipWith3FB :: (d -> xs -> xs') -> (a -> b -> c -> d) -> a -> b -> c -> xs -> xs'
+zipWith3FB cons func = \a b c r -> (func a b c) `cons` r
+
+{-# RULES
+"zipWith3"      [~1] forall f as bs cs.   zipWith3 f as bs cs = build (\c n -> foldr3 (zipWith3FB c f) n as bs cs)
+"zipWith3List"  [1]  forall f.   foldr3 (zipWith3FB (:) f) [] = zipWith3 f
+ #-}
+
 -- | 'unzip' transforms a list of pairs into a list of first components
 -- and a list of second components.
 unzip    :: [(a,b)] -> ([a],[b])