Generalise SetOperations benchmark creation to accomodate Data.Map.partitionKeys which has different types for first, second argument and for result

Author: sergv

containers-tests/benchmarks/SetOperations/SetOperations.hs

@@ -1,7 +1,7 @@
 {-# LANGUAGE BangPatterns #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 
-module SetOperations (benchmark) where
+module SetOperations (benchmark, benchmark') where
 
 import Test.Tasty.Bench (bench, defaultMain, whnf)
 import Data.List (partition, sortBy)
@@ -14,13 +14,24 @@ import Data.Tuple as Tuple
 -- * Flag if we should benchmark the operations with reversed arguments.
 -- * A list of operations.
 benchmark :: forall container. (Show container, Eq container) => ([Int] -> container) -> Bool -> [(String, container -> container -> container)] -> IO ()
-benchmark fromList swap methods = do
+benchmark fromList swap methods =
+  benchmark' fromList fromList swap methods
+
+benchmark'
+  :: forall container1 container2 container3.
+     (Show container1, Eq container1, Show container2, Eq container2, Show container3, Eq container3)
+  => ([Int] -> container1)
+  -> ([Int] -> container2)
+  -> Bool
+  -> [(String, container1 -> container2 -> container3)]
+  -> IO ()
+benchmark' fromList1 fromList2 swap methods = do
 
   defaultMain $ [ bench (method_str++"-"++input_str ++ "_" ++ data_sizes) $
                         whnf (method input1) input2
 
                 | (method_str, method) <- methods
-                , (input_str, data_sizes, (input1, input2)) <- sortBenchs (base_inputs ++ swapped_input)
+                , (input_str, data_sizes, (input1, input2)) <- sortBenchs all_inputs
                 ]
 
   where
@@ -36,12 +47,23 @@ benchmark fromList swap methods = do
     s {-small-} = n `div` 10
     t {-tiny-} = round $ sqrt $ fromIntegral n
 
-    base_inputs :: [(String,String,(container,container))]
-    base_inputs = [ ("disj", "nn", disj_nn), ("disj","ns", disj_ns), ("disj","nt", disj_nt)
-                                            , ("common","nn", common_nn), ("common","ns", common_ns), ("common","nt", common_nt)
-                                            , ("mix","nn", mix_nn), ("mix","ns", mix_ns), ("mix","nt", mix_nt)
-                                            , ("block","nn", block_nn), ("block","ns", block_ns)
-                                            ]
+    all_inputs :: [(String, String, (container1, container2))]
+    all_inputs = map (\(a, b, c) -> (a, b, fromLists c)) $ base_inputs ++ swapped_input
+
+    base_inputs :: [(String, String, ([Int], [Int]))]
+    base_inputs =
+      [ ("disj",   "nn", disj_nn)
+      , ("disj",   "ns", disj_ns)
+      , ("disj",   "nt", disj_nt)
+      , ("common", "nn", common_nn)
+      , ("common", "ns", common_ns)
+      , ("common", "nt", common_nt)
+      , ("mix",    "nn", mix_nn)
+      , ("mix",    "ns", mix_ns)
+      , ("mix",    "nt", mix_nt)
+      , ("block",  "nn", block_nn)
+      , ("block",  "ns", block_ns)
+      ]
 
     -- Input with set arguments swapped.
     swapped_input
@@ -53,19 +75,19 @@ benchmark fromList swap methods = do
         (name, reverse data_sizes ++ "_swap", Tuple.swap input_data)
 
     -- Data variants
-    all_n = fromList [1..n]
+    all_n = [1..n]
 
-    !disj_nn = seqPair $ (all_n, fromList [n+1..n+n])
-    !disj_ns = seqPair $ (all_n, fromList [n+1..n+s])
-    !disj_nt = seqPair $ (all_n, fromList [n+1..n+t])
-    !common_nn = seqPair $ (all_n, fromList [2,4..n])
-    !common_ns = seqPair $ (all_n, fromList [0,1+n`div`s..n])
-    !common_nt = seqPair $ (all_n, fromList [0,1+n`div`t..n])
-    !mix_nn = seqPair $ fromLists $ partition ((/= 0) . (`mod` 2)) [1..n+n]
-    !mix_ns = seqPair $ fromLists $ partition ((/= 0) . (`mod` (1 + n`div`s))) [1..s+n]
-    !mix_nt = seqPair $ fromLists $ partition ((/= 0) . (`mod` (1 + n`div`t))) [1..t+n]
-    !block_nn = seqPair $ fromLists $ partition ((>= t) . (`mod` (t * 2))) [1..n+n]
-    !block_ns = seqPair $ fromLists $ partition ((>= t) . (`mod` (t * (1 + n`div`s)))) [1..s+n]
+    !disj_nn   = (all_n, [n+1..n+n])
+    !disj_ns   = (all_n, [n+1..n+s])
+    !disj_nt   = (all_n, [n+1..n+t])
+    !common_nn = (all_n, [2,4..n])
+    !common_ns = (all_n, [0,1+n`div`s..n])
+    !common_nt = (all_n, [0,1+n`div`t..n])
+    !mix_nn    = partition ((/= 0) . (`mod` 2)) [1..n+n]
+    !mix_ns    = partition ((/= 0) . (`mod` (1 + n`div`s))) [1..s+n]
+    !mix_nt    = partition ((/= 0) . (`mod` (1 + n`div`t))) [1..t+n]
+    !block_nn  = partition ((>= t) . (`mod` (t * 2))) [1..n+n]
+    !block_ns  = partition ((>= t) . (`mod` (t * (1 + n`div`s)))) [1..s+n]
 
-    fromLists (xs, ys) = (fromList xs, fromList ys)
+    fromLists (xs, ys) = seqPair (fromList1 xs, fromList2 ys)
     seqPair pair@(xs, ys) = xs `seq` ys `seq` pair