From fb669f51b3f2cae79511ac3d1c43939d951b1f69 Mon Sep 17 00:00:00 2001 From: Tobias Decking Date: Thu, 6 Dec 2018 15:32:18 -0500 Subject: 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 --- libraries/base/Data/OldList.hs | 142 +++++++++++++++++++++++++++++++++++++++++ libraries/base/GHC/List.hs | 51 +++++++++++++++ 2 files changed, 193 insertions(+) (limited to 'libraries') 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]) -- cgit v1.2.1