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

3 files changed, 50 insertions, 0 deletions

diff --git a/testsuite/tests/perf/should_run/T15263.hs b/testsuite/tests/perf/should_run/T15263.hs
new file mode 100644
index 0000000000..3ba914217d
--- /dev/null
+++ b/testsuite/tests/perf/should_run/T15263.hs
@@ -0,0 +1,37 @@
+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
diff --git a/testsuite/tests/perf/should_run/T15263.stdout b/testsuite/tests/perf/should_run/T15263.stdout
new file mode 100644
index 0000000000..6b0cb99a94
--- /dev/null
+++ b/testsuite/tests/perf/should_run/T15263.stdout
@@ -0,0 +1,6 @@
+zip3: 200020000
+zipWith3: 200020000
+zipWith4: 400040000
+zipWith5: 800080000
+zipWith6: 1600160000
+zipWith7: 3200320000
diff --git a/testsuite/tests/perf/should_run/all.T b/testsuite/tests/perf/should_run/all.T
index 0b70398e46..d700fd5c56 100644
--- a/testsuite/tests/perf/should_run/all.T
+++ b/testsuite/tests/perf/should_run/all.T
@@ -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'])],