Commit fb669f51 authored by Tobias Decking's avatar Tobias Decking Committed by Ben Gamari

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
parent 1ef90f99
......@@ -149,6 +149,10 @@ Template Haskell
a representational one. There is really no reason to care about the
type of the underlying object.
- The functions ``zipWith3`` and ``zip3`` in ``Prelude`` can now fuse,
together with ``zipWith4`` to ``zipWith7`` as well as their
tuple counterparts in ``Data.List``.
Build system
~~~~~~~~~~~~
......
......@@ -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])
......
......@@ -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])
......
module Main where
import Data.List
expensive :: [Word]
expensive = [1 .. 10000]
cheap :: [Word]
cheap = repeat 2
test_zipWith :: IO ()
test_zipWith = do
let zw3 = sum $ zipWith3 (\a b c -> a*b*c) expensive cheap cheap
zw4 = sum $ zipWith4 (\a b c d -> a*b*c*d) expensive cheap cheap cheap
zw5 = sum $ zipWith5 (\a b c d e ->
a*b*c*d*e) expensive cheap cheap cheap cheap
zw6 = sum $ zipWith6 (\a b c d e f ->
a*b*c*d*e*f) expensive cheap cheap cheap cheap cheap
zw7 = sum $ zipWith7 (\a b c d e f g ->
a*b*c*d*e*f*g) expensive cheap cheap cheap cheap cheap cheap
putStrLn ("zipWith3: " ++ show zw3)
putStrLn ("zipWith4: " ++ show zw4)
putStrLn ("zipWith5: " ++ show zw5)
putStrLn ("zipWith6: " ++ show zw6)
putStrLn ("zipWith7: " ++ show zw7)
test_zip3 :: IO ()
test_zip3 = do
let z3 = foldr (\(x,y,z) acc -> x*y*z+acc) 0 (zip3 expensive cheap cheap)
putStrLn ("zip3: " ++ show z3)
main :: IO ()
main = do
test_zip3
test_zipWith
zip3: 200020000
zipWith3: 200020000
zipWith4: 400040000
zipWith5: 800080000
zipWith6: 1600160000
zipWith7: 3200320000
......@@ -351,6 +351,13 @@ test('T15226a',
compile_and_run,
['-O'])
test('T15263',
[stats_num_field('bytes allocated',
[(wordsize(64), 1382184, 4)]),
only_ways(['normal'])],
compile_and_run,
['-O'])
test('T15426',
[collect_stats('bytes allocated', 20),
only_ways(['normal'])],
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment