{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}

-- |
-- Module      : Test.Hspec.BenchGolden.Runner
-- Description : Benchmark execution and golden file comparison
-- Copyright   : (c) Marco Zocca 2026
-- License     : MIT
-- Maintainer  : @ocramz
--
-- This module handles running benchmarks and comparing results against
-- golden files. It includes:
--
-- * Benchmark execution with warm-up iterations
-- * Golden file IO (reading and writing JSON statistics)
-- * Tolerance-based comparison with variance warnings
-- * Support for updating baselines via GOLDS_GYM_ACCEPT environment variable
-- * Evaluation strategies to control how values are forced (nf variants).
--
-- = Evaluation Strategies
--
-- Benchmarks require explicit evaluation strategies to prevent GHC from
-- optimizing away computations or sharing results across iterations:
--
-- * 'nf' - Force result to normal form (deep, full evaluation)
-- * 'nfIO' - Execute IO and force result to normal form
-- * 'nfAppIO' - Apply function, execute IO, force result to normal form
-- * 'io' - Plain IO without additional forcing
--
-- These are vendored from tasty-bench under the MIT license, (c) 2021 Andrew Lelechenko.

module Test.Hspec.BenchGolden.Runner
  ( -- * Running Benchmarks
    runBenchGolden
  , runBenchmark
  , runBenchmarkWithRawTimings

    -- * Parameter Sweeps
  , runSweepPoint
  , runSweep

    -- * Golden File Operations
  , readGoldenFile
  , writeGoldenFile
  , writeActualFile
  , getGoldenPath
  , getActualPath

    -- * Comparison
  , compareStats
  , checkVariance

    -- * Robust Statistics
  , calculateRobustStats
  , calculateTrimmedMean
  , calculateMAD
  , calculateIQR
  , detectOutliers

    -- * Environment
  , shouldUpdateGolden
  , shouldSkipBenchmarks
  , setAcceptGoldens
  , setSkipBenchmarks

    -- * Benchmarkable Constructors
  , io
  , nf
  , nfIO
  , nfAppIO
  ) where

import Control.DeepSeq (NFData, rnf)
import Control.Exception (evaluate)
import Control.Monad (when, replicateM_, forM)
import Data.Aeson (eitherDecodeFileStrict, encodeFile)
import Data.IORef (IORef, newIORef, readIORef, writeIORef)
import Data.List (sort)
import qualified Data.Text as T
import Data.Time (getCurrentTime)
import qualified Data.Vector.Unboxed as V
import Data.Word (Word64)
import qualified Statistics.Sample as Stats
import System.CPUTime (getCPUTime)
import System.Directory (createDirectoryIfMissing, doesFileExist)
import System.FilePath ((</>), (<.>))
import System.IO.Unsafe (unsafePerformIO)
import GHC.Exts (SPEC(..))

import Lens.Micro ((^.))

import Test.Hspec.BenchGolden.Arch (detectArchitecture, sanitizeForFilename)
import Test.Hspec.BenchGolden.CSV (writeSweepCSV)
import Test.Hspec.BenchGolden.Lenses (metricFor, varianceFor)
import Test.Hspec.BenchGolden.Types

-- | Benchmark a pure function applied to an argument, forcing the result to
-- normal form (NF) using 'rnf' from "Control.DeepSeq".
-- This ensures the entire result structure is evaluated.
--
-- Example:
--
-- @
-- benchGolden "fib 30" (nf fib 30)
-- @
nf :: NFData b => (a -> b) -> a -> BenchAction
nf :: forall b a. NFData b => (a -> b) -> a -> BenchAction
nf = (b -> ()) -> (a -> b) -> a -> BenchAction
forall a b c. (b -> c) -> (a -> b) -> a -> BenchAction
funcToBench b -> ()
forall a. NFData a => a -> ()
rnf
{-# INLINE nf #-}

-- | Benchmark an 'IO' action, forcing the result to normal form.
--
-- Example:
--
-- @
-- benchGolden "readFile" (nfIO $ readFile "data.txt")
-- @
nfIO :: NFData a => IO a -> BenchAction
nfIO :: forall a. NFData a => IO a -> BenchAction
nfIO = (a -> ()) -> IO a -> BenchAction
forall a b. (a -> b) -> IO a -> BenchAction
ioToBench a -> ()
forall a. NFData a => a -> ()
rnf
{-# INLINE nfIO #-}

-- | Benchmark a function that performs 'IO', forcing the result to normal form.
--
-- Example:
--
-- @
-- benchGolden "lookup in map" (nfAppIO lookupInDB "key")
-- @
nfAppIO :: NFData b => (a -> IO b) -> a -> BenchAction
nfAppIO :: forall b a. NFData b => (a -> IO b) -> a -> BenchAction
nfAppIO = (b -> ()) -> (a -> IO b) -> a -> BenchAction
forall a b c. (b -> c) -> (a -> IO b) -> a -> BenchAction
ioFuncToBench b -> ()
forall a. NFData a => a -> ()
rnf
{-# INLINE nfAppIO #-}

-- | Benchmark an 'IO' action, discarding the result.
-- This is for backward compatibility with code that uses @IO ()@ actions.
--
-- Example:
-- 
-- @
-- benchGolden "compute" (io $ do
--   result <- heavyComputation
--   evaluate result)
-- @
io :: IO () -> BenchAction
io :: IO () -> BenchAction
io IO ()
action = (Word64 -> IO ()) -> BenchAction
BenchAction (\Word64
n -> Int -> IO () -> IO ()
forall (m :: * -> *) a. Applicative m => Int -> m a -> m ()
replicateM_ (Word64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
n) IO ()
action)
{-# INLINE io #-}

-- Internal helpers

funcToBench :: forall a b c. (b -> c) -> (a -> b) -> a -> BenchAction
funcToBench :: forall a b c. (b -> c) -> (a -> b) -> a -> BenchAction
funcToBench b -> c
frc = ((Word64 -> IO ()) -> BenchAction
BenchAction ((Word64 -> IO ()) -> BenchAction)
-> (a -> Word64 -> IO ()) -> a -> BenchAction
forall b c a. (b -> c) -> (a -> b) -> a -> c
.) ((a -> Word64 -> IO ()) -> a -> BenchAction)
-> ((a -> b) -> a -> Word64 -> IO ())
-> (a -> b)
-> a
-> BenchAction
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SPEC -> (a -> b) -> a -> Word64 -> IO ()
funcToBenchLoop SPEC
SPEC
  where
    -- Here we rely on the fact that GHC (unless spurred by
    -- -fstatic-argument-transformation) is not smart enough:
    -- it does not notice that `f` and `x` arguments are loop invariant
    -- and could be floated, and the whole `f x` expression shared.
    -- If we create a closure with `f` and `x` bound in the environment,
    -- then GHC is smart enough to share computation of `f x`.
    funcToBenchLoop :: SPEC -> (a -> b) -> a -> Word64 -> IO ()
    funcToBenchLoop :: SPEC -> (a -> b) -> a -> Word64 -> IO ()
funcToBenchLoop !SPEC
_ a -> b
f a
x Word64
n
      | Word64
n Word64 -> Word64 -> Bool
forall a. Eq a => a -> a -> Bool
== Word64
0    = () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
      | Bool
otherwise = do
        c
_ <- c -> IO c
forall a. a -> IO a
evaluate (b -> c
frc (a -> b
f a
x))
        SPEC -> (a -> b) -> a -> Word64 -> IO ()
funcToBenchLoop SPEC
SPEC a -> b
f a
x (Word64
n Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
- Word64
1)
{-# INLINE funcToBench #-}

ioToBench :: forall a b. (a -> b) -> IO a -> BenchAction
ioToBench :: forall a b. (a -> b) -> IO a -> BenchAction
ioToBench a -> b
frc IO a
act = (Word64 -> IO ()) -> BenchAction
BenchAction (SPEC -> IO a -> Word64 -> IO ()
ioToBenchLoop SPEC
SPEC IO a
act)
  where
    ioToBenchLoop :: SPEC -> IO a -> Word64 -> IO ()
    ioToBenchLoop :: SPEC -> IO a -> Word64 -> IO ()
ioToBenchLoop !SPEC
_ IO a
action Word64
n
      | Word64
n Word64 -> Word64 -> Bool
forall a. Eq a => a -> a -> Bool
== Word64
0    = () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
      | Bool
otherwise = do
        a
x <- IO a
action
        b
_ <- b -> IO b
forall a. a -> IO a
evaluate (a -> b
frc a
x)
        SPEC -> IO a -> Word64 -> IO ()
ioToBenchLoop SPEC
SPEC IO a
action (Word64
n Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
- Word64
1)
{-# INLINE ioToBench #-}

ioFuncToBench :: forall a b c. (b -> c) -> (a -> IO b) -> a -> BenchAction
ioFuncToBench :: forall a b c. (b -> c) -> (a -> IO b) -> a -> BenchAction
ioFuncToBench b -> c
frc = ((Word64 -> IO ()) -> BenchAction
BenchAction ((Word64 -> IO ()) -> BenchAction)
-> (a -> Word64 -> IO ()) -> a -> BenchAction
forall b c a. (b -> c) -> (a -> b) -> a -> c
.) ((a -> Word64 -> IO ()) -> a -> BenchAction)
-> ((a -> IO b) -> a -> Word64 -> IO ())
-> (a -> IO b)
-> a
-> BenchAction
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SPEC -> (a -> IO b) -> a -> Word64 -> IO ()
ioFuncToBenchLoop SPEC
SPEC
  where
    ioFuncToBenchLoop :: SPEC -> (a -> IO b) -> a -> Word64 -> IO ()
    ioFuncToBenchLoop :: SPEC -> (a -> IO b) -> a -> Word64 -> IO ()
ioFuncToBenchLoop !SPEC
_ a -> IO b
f a
x Word64
n
      | Word64
n Word64 -> Word64 -> Bool
forall a. Eq a => a -> a -> Bool
== Word64
0    = () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
      | Bool
otherwise = do
        b
y <- a -> IO b
f a
x
        c
_ <- c -> IO c
forall a. a -> IO a
evaluate (b -> c
frc b
y)
        SPEC -> (a -> IO b) -> a -> Word64 -> IO ()
ioFuncToBenchLoop SPEC
SPEC a -> IO b
f a
x (Word64
n Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
- Word64
1)
{-# INLINE ioFuncToBench #-}

-- -----------------------------------------------------------------------------
-- Benchmark execution
-- -----------------------------------------------------------------------------


-- | Run a benchmark golden test.
--
-- This function:
--
-- 1. Runs warm-up iterations (discarded)
-- 2. Runs the actual benchmark
-- 3. Writes actual results to @.actual@ file
-- 4. If no golden exists, creates it (first run)
-- 5. Otherwise, compares against golden with tolerance
--
-- The result includes any warnings (e.g., variance changes).
runBenchGolden :: BenchGolden -> IO BenchResult
runBenchGolden :: BenchGolden -> IO BenchResult
runBenchGolden BenchGolden{[Char]
BenchConfig
BenchAction
benchName :: [Char]
benchAction :: BenchAction
benchConfig :: BenchConfig
benchConfig :: BenchGolden -> BenchConfig
benchAction :: BenchGolden -> BenchAction
benchName :: BenchGolden -> [Char]
..} = do
  -- Check for skip/update flags
  Bool
skip <- IO Bool
shouldSkipBenchmarks
  if Bool
skip
    then do
      -- Return a pass with dummy stats when skipped
      UTCTime
now <- IO UTCTime
getCurrentTime
      ArchConfig
arch <- IO ArchConfig
detectArchitecture
      let dummyStats :: GoldenStats
dummyStats = Double
-> Double
-> Double
-> Double
-> Double
-> [(Int, Double)]
-> Text
-> UTCTime
-> Double
-> Double
-> Double
-> [Double]
-> GoldenStats
GoldenStats Double
0 Double
0 Double
0 Double
0 Double
0 [] (ArchConfig -> Text
archId ArchConfig
arch) UTCTime
now Double
0 Double
0 Double
0 []
      BenchResult -> IO BenchResult
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (BenchResult -> IO BenchResult) -> BenchResult -> IO BenchResult
forall a b. (a -> b) -> a -> b
$ GoldenStats -> GoldenStats -> [Warning] -> BenchResult
Pass GoldenStats
dummyStats GoldenStats
dummyStats []
    else do
      -- Detect architecture for path construction
      ArchConfig
arch <- IO ArchConfig
detectArchitecture
      let archDir :: [Char]
archDir = Text -> [Char]
T.unpack (Text -> [Char]) -> Text -> [Char]
forall a b. (a -> b) -> a -> b
$ ArchConfig -> Text
archId ArchConfig
arch
          config :: BenchConfig
config  = BenchConfig
benchConfig

      -- Create output directory
      let dir :: [Char]
dir = BenchConfig -> [Char]
outputDir BenchConfig
config [Char] -> [Char] -> [Char]
</> [Char]
archDir
      Bool -> [Char] -> IO ()
createDirectoryIfMissing Bool
True [Char]
dir

      -- Run warm-up iterations
      Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (BenchConfig -> Int
warmupIterations BenchConfig
config Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
        BenchAction -> Word64 -> IO ()
runBenchAction BenchAction
benchAction (Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Word64) -> Int -> Word64
forall a b. (a -> b) -> a -> b
$ BenchConfig -> Int
warmupIterations BenchConfig
config)

      -- Run the actual benchmark
      GoldenStats
actualStats <- [Char]
-> BenchAction -> BenchConfig -> ArchConfig -> IO GoldenStats
runBenchmark [Char]
benchName BenchAction
benchAction BenchConfig
config ArchConfig
arch

      -- Write actual results
      [Char] -> [Char] -> [Char] -> GoldenStats -> IO ()
writeActualFile (BenchConfig -> [Char]
outputDir BenchConfig
config) [Char]
archDir [Char]
benchName GoldenStats
actualStats

      -- Check if we should force update
      Bool
update <- IO Bool
shouldUpdateGolden

      -- Read or create golden file
      let goldenPath :: [Char]
goldenPath = [Char] -> [Char] -> [Char] -> [Char]
getGoldenPath (BenchConfig -> [Char]
outputDir BenchConfig
config) [Char]
archDir [Char]
benchName
      Bool
goldenExists <- [Char] -> IO Bool
doesFileExist [Char]
goldenPath

      if Bool
update Bool -> Bool -> Bool
|| Bool -> Bool
not Bool
goldenExists
        then do
          -- First run or forced update: create/update golden
          [Char] -> [Char] -> [Char] -> GoldenStats -> IO ()
writeGoldenFile (BenchConfig -> [Char]
outputDir BenchConfig
config) [Char]
archDir [Char]
benchName GoldenStats
actualStats
          BenchResult -> IO BenchResult
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (BenchResult -> IO BenchResult) -> BenchResult -> IO BenchResult
forall a b. (a -> b) -> a -> b
$ GoldenStats -> BenchResult
FirstRun GoldenStats
actualStats
        else do
          -- Compare against existing golden
          Either [Char] GoldenStats
goldenResult <- [Char] -> IO (Either [Char] GoldenStats)
readGoldenFile [Char]
goldenPath
          case Either [Char] GoldenStats
goldenResult of
            Left [Char]
err -> [Char] -> IO BenchResult
forall a. HasCallStack => [Char] -> a
error ([Char] -> IO BenchResult) -> [Char] -> IO BenchResult
forall a b. (a -> b) -> a -> b
$ [Char]
"Failed to read golden file: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
err
            Right GoldenStats
goldenStats -> do
              let result :: BenchResult
result = BenchConfig -> GoldenStats -> GoldenStats -> BenchResult
compareStats BenchConfig
config GoldenStats
goldenStats GoldenStats
actualStats
              BenchResult -> IO BenchResult
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return BenchResult
result

-- | Run a benchmark and collect statistics.
--
-- Uses raw timing collection with proper inner iteration counts to ensure
-- the SPEC trick in nf/nfIO prevents thunk sharing.
runBenchmark :: String -> BenchAction -> BenchConfig -> ArchConfig -> IO GoldenStats
runBenchmark :: [Char]
-> BenchAction -> BenchConfig -> ArchConfig -> IO GoldenStats
runBenchmark [Char]
name BenchAction
action BenchConfig
config ArchConfig
arch =
  -- Always use the raw timing path since it correctly handles SPEC
  [Char]
-> BenchAction -> BenchConfig -> ArchConfig -> IO GoldenStats
runBenchmarkWithRawTimings [Char]
name BenchAction
action BenchConfig
config ArchConfig
arch

-- | Run a benchmark with raw timing collection for robust statistics.
--
-- This function times running all iterations in a single batch, then
-- divides to get per-iteration timing. The SPEC trick in nf/nfIO
-- prevents sharing within the batch.
--
-- We collect multiple samples by running the full batch multiple times,
-- ensuring accurate measurements even with GHC's -O2 optimizations.
runBenchmarkWithRawTimings :: String -> BenchAction -> BenchConfig -> ArchConfig -> IO GoldenStats
runBenchmarkWithRawTimings :: [Char]
-> BenchAction -> BenchConfig -> ArchConfig -> IO GoldenStats
runBenchmarkWithRawTimings [Char]
_name BenchAction
action BenchConfig
config ArchConfig
arch = do
  let iters :: Word64
iters = Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (BenchConfig -> Int
iterations BenchConfig
config) :: Word64
      numSamples :: Int
numSamples = Int
10 :: Int  -- Number of timing samples to collect
  
  -- Collect raw CPU timings (each sample runs all iterations)
  [Double]
rawTimings <- [Int] -> (Int -> IO Double) -> IO [Double]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [Int
1 .. Int
numSamples] ((Int -> IO Double) -> IO [Double])
-> (Int -> IO Double) -> IO [Double]
forall a b. (a -> b) -> a -> b
$ \Int
_ -> do
    Integer
startCpu <- IO Integer
getCPUTime
    BenchAction -> Word64 -> IO ()
runBenchAction BenchAction
action Word64
iters  -- SPEC trick prevents sharing within this call
    Integer
endCpu <- IO Integer
getCPUTime
    Double -> IO Double
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Double -> IO Double) -> Double -> IO Double
forall a b. (a -> b) -> a -> b
$ Integer -> Double
picosToMillis (Integer
endCpu Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
startCpu) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Word64 -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
iters
  
  let sortedTimings :: [Double]
sortedTimings = [Double] -> [Double]
forall a. Ord a => [a] -> [a]
sort [Double]
rawTimings
      vec :: Vector Double
vec = [Double] -> Vector Double
forall a. Unbox a => [a] -> Vector a
V.fromList [Double]
sortedTimings
      
      -- Standard statistics
      mean' :: Double
mean' = Vector Double -> Double
forall (v :: * -> *). Vector v Double => v Double -> Double
Stats.mean Vector Double
vec
      stddev' :: Double
stddev' = Vector Double -> Double
forall (v :: * -> *). Vector v Double => v Double -> Double
Stats.stdDev Vector Double
vec
      -- Median: middle element of sorted vector
      median' :: Double
median' = if Vector Double -> Bool
forall a. Unbox a => Vector a -> Bool
V.null Vector Double
vec
                then Double
0.0
                else Vector Double
vec Vector Double -> Int -> Double
forall a. Unbox a => Vector a -> Int -> a
V.! (Vector Double -> Int
forall a. Unbox a => Vector a -> Int
V.length Vector Double
vec Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2)
      min' :: Double
min' = Vector Double -> Double
forall a. (Unbox a, Ord a) => Vector a -> a
V.minimum Vector Double
vec
      max' :: Double
max' = Vector Double -> Double
forall a. (Unbox a, Ord a) => Vector a -> a
V.maximum Vector Double
vec
      
      -- Percentiles
      percentiles' :: [(Int, Double)]
percentiles' = [(Int
p, Int -> Vector Double -> Double
quantile Int
p Vector Double
vec) | Int
p <- [Int
50, Int
66, Int
75, Int
80, Int
90, Int
95, Int
98, Int
99, Int
100]]
      
      -- Robust statistics
      (Double
trimmedMean', Double
mad', Double
iqr', [Double]
outliers') = BenchConfig
-> Vector Double -> Double -> (Double, Double, Double, [Double])
calculateRobustStats BenchConfig
config Vector Double
vec Double
median'

  UTCTime
now <- IO UTCTime
getCurrentTime

  GoldenStats -> IO GoldenStats
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return GoldenStats
    { statsMean :: Double
statsMean        = Double
mean'
    , statsStddev :: Double
statsStddev      = Double
stddev'
    , statsMedian :: Double
statsMedian      = Double
median'
    , statsMin :: Double
statsMin         = Double
min'
    , statsMax :: Double
statsMax         = Double
max'
    , statsPercentiles :: [(Int, Double)]
statsPercentiles = [(Int, Double)]
percentiles'
    , statsArch :: Text
statsArch        = ArchConfig -> Text
archId ArchConfig
arch
    , statsTimestamp :: UTCTime
statsTimestamp   = UTCTime
now
    , statsTrimmedMean :: Double
statsTrimmedMean = Double
trimmedMean'
    , statsMAD :: Double
statsMAD         = Double
mad'
    , statsIQR :: Double
statsIQR         = Double
iqr'
    , statsOutliers :: [Double]
statsOutliers    = [Double]
outliers'
    }
  where
    picosToMillis :: Integer -> Double
    picosToMillis :: Integer -> Double
picosToMillis Integer
t = Integer -> Double
forall a b. (Real a, Fractional b) => a -> b
realToFrac Integer
t Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
10Double -> Int -> Double
forall a b. (Num a, Integral b) => a -> b -> a
^(Int
9 :: Int))
    
    quantile :: Int -> V.Vector Double -> Double
    quantile :: Int -> Vector Double -> Double
quantile Int
p Vector Double
v =
      let idx :: Int
idx = Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
ceiling ((Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Vector Double -> Int
forall a. Unbox a => Vector a -> Int
V.length Vector Double
v) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
100) Double -> Double -> Double
forall a. Num a => a -> a -> a
* Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
p :: Double) Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
          safeIdx :: Int
safeIdx = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
0 (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min (Vector Double -> Int
forall a. Unbox a => Vector a -> Int
V.length Vector Double
v Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int
idx)
      in Vector Double -> Int -> Double
forall a. Unbox a => Vector a -> Int -> a
V.unsafeIndex Vector Double
v Int
safeIdx

-- | Calculate robust statistics from raw timing data.
--
-- Returns: (trimmed mean, MAD, IQR, outliers)
calculateRobustStats :: BenchConfig -> V.Vector Double -> Double -> (Double, Double, Double, [Double])
calculateRobustStats :: BenchConfig
-> Vector Double -> Double -> (Double, Double, Double, [Double])
calculateRobustStats BenchConfig
config Vector Double
vec Double
median' =
  let trimmedMean' :: Double
trimmedMean' = Double -> Vector Double -> Double
calculateTrimmedMean (BenchConfig -> Double
trimPercent BenchConfig
config) Vector Double
vec
      mad' :: Double
mad' = Vector Double -> Double -> Double
calculateMAD Vector Double
vec Double
median'
      iqr' :: Double
iqr' = Vector Double -> Double
calculateIQR Vector Double
vec
      outliers' :: [Double]
outliers' = Double -> Vector Double -> Double -> Double -> [Double]
detectOutliers (BenchConfig -> Double
outlierThreshold BenchConfig
config) Vector Double
vec Double
median' Double
mad'
  in (Double
trimmedMean', Double
mad', Double
iqr', [Double]
outliers')

-- | Calculate trimmed mean by removing specified percentage from each tail.
calculateTrimmedMean :: Double -> V.Vector Double -> Double
calculateTrimmedMean :: Double -> Vector Double -> Double
calculateTrimmedMean Double
trimPct Vector Double
vec
  | Vector Double -> Bool
forall a. Unbox a => Vector a -> Bool
V.null Vector Double
vec = Double
0.0
  | Double
trimPct Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
0 Bool -> Bool -> Bool
|| Double
trimPct Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
50 = Vector Double -> Double
forall (v :: * -> *). Vector v Double => v Double -> Double
Stats.mean Vector Double
vec
  | Bool
otherwise =
      let n :: Int
n = Vector Double -> Int
forall a. Unbox a => Vector a -> Int
V.length Vector Double
vec
          trimCount :: Int
trimCount = Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
round (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
trimPct Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
100.0)
          trimmed :: Vector Double
trimmed = Int -> Int -> Vector Double -> Vector Double
forall a. Unbox a => Int -> Int -> Vector a -> Vector a
V.slice Int
trimCount (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
trimCount) Vector Double
vec
      in if Vector Double -> Bool
forall a. Unbox a => Vector a -> Bool
V.null Vector Double
trimmed then Vector Double -> Double
forall (v :: * -> *). Vector v Double => v Double -> Double
Stats.mean Vector Double
vec else Vector Double -> Double
forall (v :: * -> *). Vector v Double => v Double -> Double
Stats.mean Vector Double
trimmed

-- | Calculate Median Absolute Deviation (MAD).
--
-- MAD = median(|x_i - median(x)|)
calculateMAD :: V.Vector Double -> Double -> Double
calculateMAD :: Vector Double -> Double -> Double
calculateMAD Vector Double
vec Double
med
  | Vector Double -> Bool
forall a. Unbox a => Vector a -> Bool
V.null Vector Double
vec = Double
0.0
  | Bool
otherwise =
      let deviations :: [Double]
deviations = Vector Double -> [Double]
forall a. Unbox a => Vector a -> [a]
V.toList (Vector Double -> [Double]) -> Vector Double -> [Double]
forall a b. (a -> b) -> a -> b
$ (Double -> Double) -> Vector Double -> Vector Double
forall a b. (Unbox a, Unbox b) => (a -> b) -> Vector a -> Vector b
V.map (\Double
x -> Double -> Double
forall a. Num a => a -> a
abs (Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
med)) Vector Double
vec
          sortedDevs :: [Double]
sortedDevs = [Double] -> [Double]
forall a. Ord a => [a] -> [a]
sort [Double]
deviations
          n :: Int
n = [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
sortedDevs
      in [Double]
sortedDevs [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! (Int
n Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2)

-- | Calculate Interquartile Range (IQR = Q3 - Q1).
calculateIQR :: V.Vector Double -> Double
calculateIQR :: Vector Double -> Double
calculateIQR Vector Double
vec
  | Vector Double -> Int
forall a. Unbox a => Vector a -> Int
V.length Vector Double
vec Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
4 = Double
0.0
  | Bool
otherwise =
      let q1 :: Double
q1 = Int -> Vector Double -> Double
quantileAt Int
25 Vector Double
vec
          q3 :: Double
q3 = Int -> Vector Double -> Double
quantileAt Int
75 Vector Double
vec
      in Double
q3 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
q1
  where
    quantileAt :: Int -> V.Vector Double -> Double
    quantileAt :: Int -> Vector Double -> Double
quantileAt Int
p Vector Double
v =
      let n :: Int
n = Vector Double -> Int
forall a. Unbox a => Vector a -> Int
V.length Vector Double
v
          idx :: Int
idx = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
0 (Int -> Int) -> Int -> Int
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Int
forall a. Ord a => a -> a -> a
min (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) (Int -> Int) -> Int -> Int
forall a b. (a -> b) -> a -> b
$ Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
round (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n Double -> Double -> Double
forall a. Num a => a -> a -> a
* Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
p Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
100.0 :: Double)) Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
      in Vector Double -> Int -> Double
forall a. Unbox a => Vector a -> Int -> a
V.unsafeIndex Vector Double
v Int
idx

-- | Detect outliers using MAD-based threshold.
--
-- An observation is an outlier if: |x - median| > threshold * MAD
detectOutliers :: Double -> V.Vector Double -> Double -> Double -> [Double]
detectOutliers :: Double -> Vector Double -> Double -> Double -> [Double]
detectOutliers Double
threshold Vector Double
vec Double
med Double
mad
  | Vector Double -> Bool
forall a. Unbox a => Vector a -> Bool
V.null Vector Double
vec = []
  | Double
mad Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
0 = []  -- No variance, no outliers
  | Bool
otherwise =
      let isOutlier :: Double -> Bool
isOutlier Double
x = Double -> Double
forall a. Num a => a -> a
abs (Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
med) Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
threshold Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
mad
      in Vector Double -> [Double]
forall a. Unbox a => Vector a -> [a]
V.toList (Vector Double -> [Double]) -> Vector Double -> [Double]
forall a b. (a -> b) -> a -> b
$ (Double -> Bool) -> Vector Double -> Vector Double
forall a. Unbox a => (a -> Bool) -> Vector a -> Vector a
V.filter Double -> Bool
isOutlier Vector Double
vec

-- | Compare actual stats against golden stats.
--
-- Returns a 'BenchResult' indicating whether the benchmark passed,
-- regressed, or improved, along with any warnings.
--
-- = Hybrid Tolerance Strategy
--
-- The comparison uses BOTH percentage and absolute tolerance (when configured):
--
-- 1. Calculate percentage difference: @((actual - golden) / golden) * 100@
--
-- 2. Pass if @abs(percentDiff) <= tolerancePercent@ (percentage check)
--
-- 3. OR if @abs(actual - golden) <= absoluteToleranceMs@ (absolute check)
--
-- This prevents false failures for sub-millisecond operations where measurement
-- noise creates large percentage variations despite negligible absolute differences.
compareStats :: BenchConfig -> GoldenStats -> GoldenStats -> BenchResult
compareStats :: BenchConfig -> GoldenStats -> GoldenStats -> BenchResult
compareStats BenchConfig
config GoldenStats
golden GoldenStats
actual =
  let -- Use lens-based metric selection
      metric :: (Double -> Const Double Double)
-> GoldenStats -> Const Double GoldenStats
metric = BenchConfig -> Lens' GoldenStats Double
metricFor BenchConfig
config
      goldenValue :: Double
goldenValue = GoldenStats
golden GoldenStats
-> ((Double -> Const Double Double)
    -> GoldenStats -> Const Double GoldenStats)
-> Double
forall s a. s -> Getting a s a -> a
^. (Double -> Const Double Double)
-> GoldenStats -> Const Double GoldenStats
metric
      actualValue :: Double
actualValue = GoldenStats
actual GoldenStats
-> ((Double -> Const Double Double)
    -> GoldenStats -> Const Double GoldenStats)
-> Double
forall s a. s -> Getting a s a -> a
^. (Double -> Const Double Double)
-> GoldenStats -> Const Double GoldenStats
metric

      -- Calculate percentage difference
      meanDiff :: Double
meanDiff = if Double
goldenValue Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
0
                 then if Double
actualValue Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
0 then Double
0 else Double
100
                 else ((Double
actualValue Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
goldenValue) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
goldenValue) Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
100

      absDiff :: Double
absDiff = Double -> Double
forall a. Num a => a -> a
abs Double
meanDiff
      tolerance :: Double
tolerance = BenchConfig -> Double
tolerancePercent BenchConfig
config

      -- Calculate absolute time difference (in milliseconds)
      absTimeDiff :: Double
absTimeDiff = Double -> Double
forall a. Num a => a -> a
abs (Double
actualValue Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
goldenValue)

      -- Check if within absolute tolerance (hybrid tolerance strategy)
      withinAbsoluteTolerance :: Bool
withinAbsoluteTolerance = case BenchConfig -> Maybe Double
absoluteToleranceMs BenchConfig
config of
        Maybe Double
Nothing -> Bool
False
        Just Double
absThreshold -> Double
absTimeDiff Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
absThreshold

      -- Check variance if enabled
      baseWarnings :: [Warning]
baseWarnings = if BenchConfig -> Bool
warnOnVarianceChange BenchConfig
config
                     then BenchConfig -> GoldenStats -> GoldenStats -> [Warning]
checkVariance BenchConfig
config GoldenStats
golden GoldenStats
actual
                     else []
      
      -- Add outlier warnings if robust statistics enabled
      outlierWarnings :: [Warning]
outlierWarnings = if BenchConfig -> Bool
useRobustStatistics BenchConfig
config Bool -> Bool -> Bool
&& Bool -> Bool
not ([Double] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null (GoldenStats -> [Double]
statsOutliers GoldenStats
actual))
                        then [Int -> [Double] -> Warning
OutliersDetected ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([Double] -> Int) -> [Double] -> Int
forall a b. (a -> b) -> a -> b
$ GoldenStats -> [Double]
statsOutliers GoldenStats
actual) (GoldenStats -> [Double]
statsOutliers GoldenStats
actual)]
                        else []
      
      warnings :: [Warning]
warnings = [Warning]
baseWarnings [Warning] -> [Warning] -> [Warning]
forall a. [a] -> [a] -> [a]
++ [Warning]
outlierWarnings

  in if Double
absDiff Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
tolerance Bool -> Bool -> Bool
|| Bool
withinAbsoluteTolerance
     then GoldenStats -> GoldenStats -> [Warning] -> BenchResult
Pass GoldenStats
golden GoldenStats
actual [Warning]
warnings
     else if Double
meanDiff Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
0
          then GoldenStats
-> GoldenStats -> Double -> Double -> Maybe Double -> BenchResult
Regression GoldenStats
golden GoldenStats
actual Double
meanDiff Double
tolerance (BenchConfig -> Maybe Double
absoluteToleranceMs BenchConfig
config)
          else GoldenStats
-> GoldenStats -> Double -> Double -> Maybe Double -> BenchResult
Improvement GoldenStats
golden GoldenStats
actual (Double -> Double
forall a. Num a => a -> a
abs Double
meanDiff) Double
tolerance (BenchConfig -> Maybe Double
absoluteToleranceMs BenchConfig
config)

-- | Check for variance changes and generate warnings.
checkVariance :: BenchConfig -> GoldenStats -> GoldenStats -> [Warning]
checkVariance :: BenchConfig -> GoldenStats -> GoldenStats -> [Warning]
checkVariance BenchConfig
config GoldenStats
golden GoldenStats
actual =
  let -- Use lens-based variance metric selection
      vLens :: (Double -> Const Double Double)
-> GoldenStats -> Const Double GoldenStats
vLens = BenchConfig -> Lens' GoldenStats Double
varianceFor BenchConfig
config
      goldenVar :: Double
goldenVar = GoldenStats
golden GoldenStats
-> ((Double -> Const Double Double)
    -> GoldenStats -> Const Double GoldenStats)
-> Double
forall s a. s -> Getting a s a -> a
^. (Double -> Const Double Double)
-> GoldenStats -> Const Double GoldenStats
vLens
      actualVar :: Double
actualVar = GoldenStats
actual GoldenStats
-> ((Double -> Const Double Double)
    -> GoldenStats -> Const Double GoldenStats)
-> Double
forall s a. s -> Getting a s a -> a
^. (Double -> Const Double Double)
-> GoldenStats -> Const Double GoldenStats
vLens

      varDiff :: Double
varDiff = if Double
goldenVar Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
0
                then if Double
actualVar Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
0 then Double
0 else Double
100
                else ((Double
actualVar Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
goldenVar) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
goldenVar) Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
100

      _absVarDiff :: Double
_absVarDiff = Double -> Double
forall a. Num a => a -> a
abs Double
varDiff
      varTolerance :: Double
varTolerance = BenchConfig -> Double
varianceTolerancePercent BenchConfig
config

      -- Coefficient of variation (CV) for high variance detection
      -- Use appropriate measure based on mode
      cv :: Double
cv = if BenchConfig -> Bool
useRobustStatistics BenchConfig
config
           then if GoldenStats -> Double
statsMedian GoldenStats
actual Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
0 then Double
0 else GoldenStats -> Double
statsMAD GoldenStats
actual Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ GoldenStats -> Double
statsMedian GoldenStats
actual
           else if GoldenStats -> Double
statsMean GoldenStats
actual Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
0 then Double
0 else GoldenStats -> Double
statsStddev GoldenStats
actual Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ GoldenStats -> Double
statsMean GoldenStats
actual

  in [[Warning]] -> [Warning]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat
       [ [ Double -> Double -> Double -> Double -> Warning
VarianceIncreased Double
goldenVar Double
actualVar Double
varDiff Double
varTolerance
         | Double
varDiff Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
varTolerance ]
       , [ Double -> Double -> Double -> Double -> Warning
VarianceDecreased Double
goldenVar Double
actualVar (Double -> Double
forall a. Num a => a -> a
abs Double
varDiff) Double
varTolerance
         | Double
varDiff Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double -> Double
forall a. Num a => a -> a
negate Double
varTolerance ]
       , [ Double -> Warning
HighVariance Double
cv | Double
cv Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
0.5 ]  -- CV > 50% is considered high
       ]

-- | Get the path for a golden file.
getGoldenPath :: FilePath -> FilePath -> String -> FilePath
getGoldenPath :: [Char] -> [Char] -> [Char] -> [Char]
getGoldenPath [Char]
outDir [Char]
archDir [Char]
name =
  [Char]
outDir [Char] -> [Char] -> [Char]
</> [Char]
archDir [Char] -> [Char] -> [Char]
</> [Char] -> [Char]
sanitizeName [Char]
name [Char] -> [Char] -> [Char]
<.> [Char]
"golden"

-- | Get the path for an actual results file.
getActualPath :: FilePath -> FilePath -> String -> FilePath
getActualPath :: [Char] -> [Char] -> [Char] -> [Char]
getActualPath [Char]
outDir [Char]
archDir [Char]
name =
  [Char]
outDir [Char] -> [Char] -> [Char]
</> [Char]
archDir [Char] -> [Char] -> [Char]
</> [Char] -> [Char]
sanitizeName [Char]
name [Char] -> [Char] -> [Char]
<.> [Char]
"actual"

-- | Sanitize a benchmark name for use in filenames.
sanitizeName :: String -> FilePath
sanitizeName :: [Char] -> [Char]
sanitizeName = Text -> [Char]
T.unpack (Text -> [Char]) -> ([Char] -> Text) -> [Char] -> [Char]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> Text
sanitizeForFilename (Text -> Text) -> ([Char] -> Text) -> [Char] -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Char] -> Text
T.pack

-- | Read a golden file.
readGoldenFile :: FilePath -> IO (Either String GoldenStats)
readGoldenFile :: [Char] -> IO (Either [Char] GoldenStats)
readGoldenFile = [Char] -> IO (Either [Char] GoldenStats)
forall a. FromJSON a => [Char] -> IO (Either [Char] a)
eitherDecodeFileStrict

-- | Write a golden file.
writeGoldenFile :: FilePath -> FilePath -> String -> GoldenStats -> IO ()
writeGoldenFile :: [Char] -> [Char] -> [Char] -> GoldenStats -> IO ()
writeGoldenFile [Char]
outDir [Char]
archDir [Char]
name GoldenStats
stats = do
  let path :: [Char]
path = [Char] -> [Char] -> [Char] -> [Char]
getGoldenPath [Char]
outDir [Char]
archDir [Char]
name
  [Char] -> GoldenStats -> IO ()
forall a. ToJSON a => [Char] -> a -> IO ()
encodeFile [Char]
path GoldenStats
stats

-- | Write an actual results file.
writeActualFile :: FilePath -> FilePath -> String -> GoldenStats -> IO ()
writeActualFile :: [Char] -> [Char] -> [Char] -> GoldenStats -> IO ()
writeActualFile [Char]
outDir [Char]
archDir [Char]
name GoldenStats
stats = do
  let path :: [Char]
path = [Char] -> [Char] -> [Char] -> [Char]
getActualPath [Char]
outDir [Char]
archDir [Char]
name
  [Char] -> GoldenStats -> IO ()
forall a. ToJSON a => [Char] -> a -> IO ()
encodeFile [Char]
path GoldenStats
stats

-- | Global state for tracking command-line flags set by hspec options.
{-# NOINLINE acceptGoldensRef #-}
acceptGoldensRef :: IORef Bool
acceptGoldensRef :: IORef Bool
acceptGoldensRef = IO (IORef Bool) -> IORef Bool
forall a. IO a -> a
unsafePerformIO (IO (IORef Bool) -> IORef Bool) -> IO (IORef Bool) -> IORef Bool
forall a b. (a -> b) -> a -> b
$ Bool -> IO (IORef Bool)
forall a. a -> IO (IORef a)
newIORef Bool
False

{-# NOINLINE skipBenchmarksRef #-}
skipBenchmarksRef :: IORef Bool
skipBenchmarksRef :: IORef Bool
skipBenchmarksRef = IO (IORef Bool) -> IORef Bool
forall a. IO a -> a
unsafePerformIO (IO (IORef Bool) -> IORef Bool) -> IO (IORef Bool) -> IORef Bool
forall a b. (a -> b) -> a -> b
$ Bool -> IO (IORef Bool)
forall a. a -> IO (IORef a)
newIORef Bool
False

-- | Set the accept goldens flag (called from BenchGolden Example instance).
setAcceptGoldens :: Bool -> IO ()
setAcceptGoldens :: Bool -> IO ()
setAcceptGoldens = IORef Bool -> Bool -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef Bool
acceptGoldensRef

-- | Set the skip benchmarks flag (called from BenchGolden Example instance).
setSkipBenchmarks :: Bool -> IO ()
setSkipBenchmarks :: Bool -> IO ()
setSkipBenchmarks = IORef Bool -> Bool -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef Bool
skipBenchmarksRef

-- | Check if golden files should be updated.
--
-- Returns 'True' if @GOLDS_GYM_ACCEPT@ environment variable is set.
-- 
-- Usage:
--
-- @
-- GOLDS_GYM_ACCEPT=1 cabal test
-- GOLDS_GYM_ACCEPT=1 stack test
-- @
shouldUpdateGolden :: IO Bool
shouldUpdateGolden :: IO Bool
shouldUpdateGolden = IORef Bool -> IO Bool
forall a. IORef a -> IO a
readIORef IORef Bool
acceptGoldensRef

-- | Check if benchmarks should be skipped entirely.
--
-- Returns 'True' if @GOLDS_GYM_SKIP@ environment variable is set.
-- Useful for CI environments where benchmark hardware is inconsistent.
--
-- Usage:
--
-- @
-- GOLDS_GYM_SKIP=1 cabal test
-- GOLDS_GYM_SKIP=1 stack test
-- @
shouldSkipBenchmarks :: IO Bool
shouldSkipBenchmarks :: IO Bool
shouldSkipBenchmarks = IORef Bool -> IO Bool
forall a. IORef a -> IO a
readIORef IORef Bool
skipBenchmarksRef

-- -----------------------------------------------------------------------------
-- Parameter Sweeps
-- -----------------------------------------------------------------------------

-- | Run a single point of a parameter sweep.
--
-- This is similar to 'runBenchGolden' but returns the 'GoldenStats' along
-- with the 'BenchResult', allowing the caller to accumulate stats for CSV export.
--
-- Each point is saved to its own golden file with the parameter value
-- included in the filename (e.g., @sort-scaling_n=1000.golden@).
runSweepPoint ::
     Show a
  => String       -- ^ Base sweep name
  -> BenchConfig
  -> T.Text       -- ^ Parameter name
  -> a            -- ^ Parameter value
  -> BenchAction
  -> IO (BenchResult, GoldenStats)
runSweepPoint :: forall a.
Show a =>
[Char]
-> BenchConfig
-> Text
-> a
-> BenchAction
-> IO (BenchResult, GoldenStats)
runSweepPoint [Char]
sweepName BenchConfig
config Text
paramName a
paramValue BenchAction
action = do
  let pointName :: [Char]
pointName = [Char]
sweepName [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"_" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Text -> [Char]
T.unpack Text
paramName [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"=" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ a -> [Char]
forall a. Show a => a -> [Char]
show a
paramValue

  -- Run the benchmark (this writes golden/actual files via runBenchGolden logic)
  Bool
skip <- IO Bool
shouldSkipBenchmarks
  if Bool
skip
    then do
      UTCTime
now <- IO UTCTime
getCurrentTime
      ArchConfig
arch <- IO ArchConfig
detectArchitecture
      let dummyStats :: GoldenStats
dummyStats = Double
-> Double
-> Double
-> Double
-> Double
-> [(Int, Double)]
-> Text
-> UTCTime
-> Double
-> Double
-> Double
-> [Double]
-> GoldenStats
GoldenStats Double
0 Double
0 Double
0 Double
0 Double
0 [] (ArchConfig -> Text
archId ArchConfig
arch) UTCTime
now Double
0 Double
0 Double
0 []
      (BenchResult, GoldenStats) -> IO (BenchResult, GoldenStats)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (GoldenStats -> GoldenStats -> [Warning] -> BenchResult
Pass GoldenStats
dummyStats GoldenStats
dummyStats [], GoldenStats
dummyStats)
    else do
      ArchConfig
arch <- IO ArchConfig
detectArchitecture
      let archDir :: [Char]
archDir = Text -> [Char]
T.unpack (Text -> [Char]) -> Text -> [Char]
forall a b. (a -> b) -> a -> b
$ ArchConfig -> Text
archId ArchConfig
arch

      -- Create output directory
      let dir :: [Char]
dir = BenchConfig -> [Char]
outputDir BenchConfig
config [Char] -> [Char] -> [Char]
</> [Char]
archDir
      Bool -> [Char] -> IO ()
createDirectoryIfMissing Bool
True [Char]
dir

      -- Run warm-up iterations
      Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (BenchConfig -> Int
warmupIterations BenchConfig
config Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
        BenchAction -> Word64 -> IO ()
runBenchAction BenchAction
action (Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Word64) -> Int -> Word64
forall a b. (a -> b) -> a -> b
$ BenchConfig -> Int
warmupIterations BenchConfig
config)

      -- Run the actual benchmark
      GoldenStats
actualStats <- [Char]
-> BenchAction -> BenchConfig -> ArchConfig -> IO GoldenStats
runBenchmark [Char]
pointName BenchAction
action BenchConfig
config ArchConfig
arch

      -- Write actual results
      [Char] -> [Char] -> [Char] -> GoldenStats -> IO ()
writeActualFile (BenchConfig -> [Char]
outputDir BenchConfig
config) [Char]
archDir [Char]
pointName GoldenStats
actualStats

      -- Check if we should force update
      Bool
update <- IO Bool
shouldUpdateGolden

      -- Read or create golden file
      let goldenPath :: [Char]
goldenPath = [Char] -> [Char] -> [Char] -> [Char]
getGoldenPath (BenchConfig -> [Char]
outputDir BenchConfig
config) [Char]
archDir [Char]
pointName
      Bool
goldenExists <- [Char] -> IO Bool
doesFileExist [Char]
goldenPath

      if Bool
update Bool -> Bool -> Bool
|| Bool -> Bool
not Bool
goldenExists
        then do
          -- First run or forced update: create/update golden
          [Char] -> [Char] -> [Char] -> GoldenStats -> IO ()
writeGoldenFile (BenchConfig -> [Char]
outputDir BenchConfig
config) [Char]
archDir [Char]
pointName GoldenStats
actualStats
          (BenchResult, GoldenStats) -> IO (BenchResult, GoldenStats)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (GoldenStats -> BenchResult
FirstRun GoldenStats
actualStats, GoldenStats
actualStats)
        else do
          -- Compare against existing golden
          Either [Char] GoldenStats
goldenResult <- [Char] -> IO (Either [Char] GoldenStats)
readGoldenFile [Char]
goldenPath
          case Either [Char] GoldenStats
goldenResult of
            Left [Char]
err -> [Char] -> IO (BenchResult, GoldenStats)
forall a. HasCallStack => [Char] -> a
error ([Char] -> IO (BenchResult, GoldenStats))
-> [Char] -> IO (BenchResult, GoldenStats)
forall a b. (a -> b) -> a -> b
$ [Char]
"Failed to read golden file: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
err
            Right GoldenStats
goldenStats -> do
              let result :: BenchResult
result = BenchConfig -> GoldenStats -> GoldenStats -> BenchResult
compareStats BenchConfig
config GoldenStats
goldenStats GoldenStats
actualStats
              (BenchResult, GoldenStats) -> IO (BenchResult, GoldenStats)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (BenchResult
result, GoldenStats
actualStats)

-- | Run a full parameter sweep and write CSV output.
--
-- This runs benchmarks for all parameter values, saves individual golden
-- files, and writes a single CSV file with all results for analysis.
--
-- The CSV file is placed at:
--
-- @
-- \<outputDir\>/\<sweep-name\>-\<arch-id\>.csv
-- @
runSweep ::
     Show a
  => String           -- ^ Sweep name
  -> BenchConfig
  -> T.Text           -- ^ Parameter name (for CSV column header)
  -> [a]              -- ^ Parameter values to sweep over
  -> (a -> BenchAction)  -- ^ Action generator
  -> IO [(a, BenchResult, GoldenStats)]
runSweep :: forall a.
Show a =>
[Char]
-> BenchConfig
-> Text
-> [a]
-> (a -> BenchAction)
-> IO [(a, BenchResult, GoldenStats)]
runSweep [Char]
sweepName BenchConfig
config Text
paramName [a]
paramValues a -> BenchAction
mkAction = do
  ArchConfig
arch <- IO ArchConfig
detectArchitecture
  
  -- Run each parameter value
  [(a, BenchResult, GoldenStats)]
results <- [a]
-> (a -> IO (a, BenchResult, GoldenStats))
-> IO [(a, BenchResult, GoldenStats)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [a]
paramValues ((a -> IO (a, BenchResult, GoldenStats))
 -> IO [(a, BenchResult, GoldenStats)])
-> (a -> IO (a, BenchResult, GoldenStats))
-> IO [(a, BenchResult, GoldenStats)]
forall a b. (a -> b) -> a -> b
$ \a
paramVal -> do
    let action :: BenchAction
action = a -> BenchAction
mkAction a
paramVal
    (BenchResult
result, GoldenStats
stats) <- [Char]
-> BenchConfig
-> Text
-> a
-> BenchAction
-> IO (BenchResult, GoldenStats)
forall a.
Show a =>
[Char]
-> BenchConfig
-> Text
-> a
-> BenchAction
-> IO (BenchResult, GoldenStats)
runSweepPoint [Char]
sweepName BenchConfig
config Text
paramName a
paramVal BenchAction
action
    (a, BenchResult, GoldenStats) -> IO (a, BenchResult, GoldenStats)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (a
paramVal, BenchResult
result, GoldenStats
stats)
  
  -- Write CSV with all results
  let csvRows :: [(Text, GoldenStats)]
csvRows = [([Char] -> Text
T.pack (a -> [Char]
forall a. Show a => a -> [Char]
show a
pv), GoldenStats
stats) | (a
pv, BenchResult
_, GoldenStats
stats) <- [(a, BenchResult, GoldenStats)]
results]
  [Char] -> Text -> [Char] -> Text -> [(Text, GoldenStats)] -> IO ()
writeSweepCSV (BenchConfig -> [Char]
outputDir BenchConfig
config) (ArchConfig -> Text
archId ArchConfig
arch) [Char]
sweepName Text
paramName [(Text, GoldenStats)]
csvRows
  
  [(a, BenchResult, GoldenStats)]
-> IO [(a, BenchResult, GoldenStats)]
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return [(a, BenchResult, GoldenStats)]
results