content: -- SPDX-FileCopyrightText: Copyright Preetham Gujjula
  -- SPDX-License-Identifier: BSD-3-Clause
  {-# LANGUAGE CPP #-}
  {-# LANGUAGE ImpredicativeTypes #-}
  
  module Bench.ApplyMerge (benchmarks) where
  
  import qualified ApplyMerge.DoublyLinkedList
  import qualified ApplyMerge.IntMap
  import qualified ApplyMerge.IntSet
  import qualified ApplyMerge.MergeAll
  import Data.Function ((&))
  #if !MIN_VERSION_base(4,20,0)
  import Data.Foldable (foldl')
  #endif
  import Data.List.Ordered (minus)
  import Test.Tasty.Bench (Benchmark, bench, bgroup, nf)
  
  benchmarks :: Benchmark
  benchmarks =
    bgroup
      "ApplyMerge"
      [ funcToBenchmark
          "linear shape: applyMerge const [1..] [1..]"
          const,
        funcToBenchmark
          "double linear shape: applyMerge min [1..] [1..]"
          min,
        funcToBenchmark
          "box shape: applyMerge max [1..] [1..]"
          max,
        funcToBenchmark
          "triangular shape: applyMerge (+) [1..] [1..]"
          (+),
        funcToBenchmark
          "skewed triangular shape: applyMerge (\\x y -> 4 * x + y) [1..] [1..]"
          (\x y -> 4 * x + y),
        funcToBenchmark
          "hyperbolic shape: applyMerge (*) [1..] [1..]"
          (*),
        funcToBenchmark
          "skewed hyperbolic shape: applyMerge (\\x y -> x^3 * y) [1..] [1..]"
          ( \x y ->
              floor @Double $
                100 * (3 * log (fromIntegral x + fromIntegral y))
          ),
        funcToBenchmark
          "circular shape: applyMerge (\\x y -> x*x + y*y) [1..] [1..]"
          (\x y -> x * x + y * y),
        funcToBenchmark
          "elliptical shape: applyMerge (\\x y -> 4*x*x + y*y) [1..] [1..]"
          (\x y -> 4 * x * x + y * y),
        collapseToBenchmark "primes" $ \applyMerge n ->
          let zero :: Int
              zero = (n `quot` maxBound)
  
              primes :: [Int]
              primes = 2 : 3 : 5 : ([7 ..] `minus` composites)
  
              composites :: [Int]
              composites = applyMerge (\p j -> p * (p + j)) primes [zero ..]
           in foldl' (+) 0 (takeWhile (<= n) primes)
      ]
  
  type ApplyMerge = forall a b c. (Ord c) => (a -> b -> c) -> [a] -> [b] -> [c]
  
  funcToBenchmark :: String -> (Int -> Int -> Int) -> Benchmark
  funcToBenchmark name f = collapseToBenchmark name (funcToCollapse f)
  
  funcToCollapse :: (Int -> Int -> Int) -> ApplyMerge -> Int -> Int
  funcToCollapse f applyMerge n =
    let one = (n `quot` maxBound) + 1
     in applyMerge f [one ..] [one ..]
          & take n
          & foldl' (+) 0
  
  collapseToBenchmark :: String -> (ApplyMerge -> Int -> Int) -> Benchmark
  collapseToBenchmark name collapse = bgroup name (map mkBench [1 .. 6])
    where
      mkBench :: Int -> Benchmark
      mkBench i =
        let limit :: Int
            limit = 10 ^ i
  
            -- For some reason, GHC 8.0 complains if we set the signature
            --   applyMerges :: [(String, ApplyMerge)]
            -- We can remove this hack once we drop support for GHC 8.0
            applyMerges ::
              [ ( String,
                  forall a b c. (Ord c) => (a -> b -> c) -> [a] -> [b] -> [c]
                )
              ]
            applyMerges =
              [ ("DoublyLinkedList", ApplyMerge.DoublyLinkedList.applyMerge),
                ("IntMap", ApplyMerge.IntMap.applyMerge),
                ("IntSet", ApplyMerge.IntSet.applyMerge),
                ("MergeAll", ApplyMerge.MergeAll.applyMerge),
                ( "MergeAll (flipped)",
                  flip . ApplyMerge.MergeAll.applyMerge . flip
                )
              ]
         in bgroup (show limit) $
              flip map applyMerges $ \(applyMergeName, applyMerge) ->
                bench applyMergeName (nf (collapse applyMerge) limit)
  