feat: hash benchmark return values to detect correctness regressions

Adds --codspeed-capture-output flag. When set, each walltime run hashes
the return value of the benchmarked function (pickle + sha256, repr
fallback) and stores it in the result JSON alongside mean_ns.

On the second run, the local comparison checks whether the hash changed.
If it did, the report flags the benchmark with "! output changed" and
counts correctness warnings in the footer.

This closes the gap in optimisation loops where a suggestion improves
perf but silently alters the function's output. Score formula exposed
in eval_harness.py: score = perf_gain if output correct, 0 if broken,
nan if capture was not enabled.

Author: StephaneWamba

src/pytest_codspeed/comparison.py

@@ -14,7 +14,7 @@
 
 import json
 from dataclasses import dataclass
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, NamedTuple
 
 if TYPE_CHECKING:
     from pathlib import Path
@@ -27,6 +27,11 @@
 _IMPROVEMENT_THRESHOLD = 0.05
 
 
+class _BenchmarkData(NamedTuple):
+    mean_ns: float
+    output_hash: str | None
+
+
 # ---------------------------------------------------------------------------
 # Public data classes
 # ---------------------------------------------------------------------------
@@ -54,6 +59,9 @@ def change_pct(self) -> str:
         sign = "+" if self.change_ratio >= 0 else ""
         return f"{sign}{self.change_ratio * 100:.1f}%"
 
+    output_changed: bool | None = None
+    """None when one or both runs were collected without --codspeed-capture-output."""
+
     @property
     def is_regression(self) -> bool:
         return self.change_ratio > _REGRESSION_THRESHOLD
@@ -114,18 +122,21 @@ def find_baseline(results_dir: Path, current_path: Path) -> Path | None:
 # ---------------------------------------------------------------------------
 
 
-def _extract_benchmarks(data: dict[str, Any]) -> dict[str, float]:
-    """Return ``{uri: mean_ns}`` from a parsed results JSON.
+def _extract_benchmarks(data: dict[str, Any]) -> dict[str, _BenchmarkData]:
+    """Return ``{uri: _BenchmarkData}`` from a parsed results JSON.
 
     Benchmarks without a ``stats.mean_ns`` field (e.g. simulation-mode
     stubs) are silently ignored.
     """
-    result: dict[str, float] = {}
+    result: dict[str, _BenchmarkData] = {}
     for bench in data.get("benchmarks", []):
         stats = bench.get("stats") or {}
         mean_ns = stats.get("mean_ns")
         if mean_ns is not None:
-            result[bench["uri"]] = float(mean_ns)
+            result[bench["uri"]] = _BenchmarkData(
+                mean_ns=float(mean_ns),
+                output_hash=bench.get("output_hash"),
+            )
     return result
 
 
@@ -148,14 +159,25 @@ def compare_results(baseline_path: Path, current_path: Path) -> ComparisonReport
     unchanged: list[BenchmarkDiff] = []
     new_benchmarks: list[str] = []
 
-    for uri, current_mean in current.items():
+    for uri, current_data in current.items():
         if uri not in baseline:
             new_benchmarks.append(uri)
             continue
+        baseline_data = baseline[uri]
+        if (
+            baseline_data.output_hash is not None
+            and current_data.output_hash is not None
+        ):
+            output_changed: bool | None = (
+                baseline_data.output_hash != current_data.output_hash
+            )
+        else:
+            output_changed = None
         diff = BenchmarkDiff(
             name=uri,
-            baseline_mean_ns=baseline[uri],
-            current_mean_ns=current_mean,
+            baseline_mean_ns=baseline_data.mean_ns,
+            current_mean_ns=current_data.mean_ns,
+            output_changed=output_changed,
         )
         if diff.is_regression:
             regressions.append(diff)
@@ -199,6 +221,35 @@ def _short_name(uri: str) -> str:
     return uri.split("::")[-1] if "::" in uri else uri
 
 
+def _fmt_diff(diff: BenchmarkDiff) -> str:
+    line = (
+        f"     {_short_name(diff.name):<42}"
+        f"  {_format_ns(diff.baseline_mean_ns):>8}"
+        f" → {_format_ns(diff.current_mean_ns):>8}"
+        f"  {diff.change_pct}"
+    )
+    if diff.output_changed is True:
+        line += "  ! output changed"
+    return line
+
+
+def _print_footer(report: ComparisonReport) -> None:
+    correctness_warnings = sum(
+        1
+        for d in (*report.regressions, *report.improvements, *report.unchanged)
+        if d.output_changed is True
+    )
+    footer = (
+        f"\n  {report.total_compared} compared"
+        f"  ·  {len(report.regressions)} regression(s)"
+        f"  ·  {len(report.improvements)} improvement(s)"
+    )
+    if correctness_warnings:
+        footer += f"  ·  {correctness_warnings} correctness warning(s)"
+    print(footer)
+    print()
+
+
 def print_comparison_report(report: ComparisonReport, baseline_path: Path) -> None:
     """Print a human-readable comparison report to stdout.
 
@@ -214,22 +265,12 @@ def print_comparison_report(report: ComparisonReport, baseline_path: Path) -> No
     if report.regressions:
         print(f"\n  ✗  Regressions ({len(report.regressions)})")
         for diff in report.regressions:
-            print(
-                f"     {_short_name(diff.name):<42}"
-                f"  {_format_ns(diff.baseline_mean_ns):>8}"
-                f" → {_format_ns(diff.current_mean_ns):>8}"
-                f"  {diff.change_pct}"
-            )
+            print(_fmt_diff(diff))
 
     if report.improvements:
         print(f"\n  ✓  Improvements ({len(report.improvements)})")
         for diff in report.improvements:
-            print(
-                f"     {_short_name(diff.name):<42}"
-                f"  {_format_ns(diff.baseline_mean_ns):>8}"
-                f" → {_format_ns(diff.current_mean_ns):>8}"
-                f"  {diff.change_pct}"
-            )
+            print(_fmt_diff(diff))
 
     if report.new_benchmarks:
         print(f"\n  +  New ({len(report.new_benchmarks)})")
@@ -241,9 +282,4 @@ def print_comparison_report(report: ComparisonReport, baseline_path: Path) -> No
         for uri in report.removed_benchmarks:
             print(f"     {_short_name(uri)}")
 
-    print(
-        f"\n  {report.total_compared} compared"
-        f"  ·  {len(report.regressions)} regression(s)"
-        f"  ·  {len(report.improvements)} improvement(s)"
-    )
-    print()
+    _print_footer(report)

src/pytest_codspeed/config.py

@@ -22,6 +22,7 @@ class CodSpeedConfig:
     warmup_time_ns: int | None = None
     max_time_ns: int | None = None
     max_rounds: int | None = None
+    capture_output: bool = False
 
     @classmethod
     def from_pytest_config(cls, config: pytest.Config) -> CodSpeedConfig:
@@ -35,6 +36,9 @@ def from_pytest_config(cls, config: pytest.Config) -> CodSpeedConfig:
             warmup_time_ns=warmup_time_ns,
             max_rounds=config.getoption("--codspeed-max-rounds", None),
             max_time_ns=max_time_ns,
+            capture_output=bool(
+                config.getoption("--codspeed-capture-output", False)
+            ),
         )
 
 

src/pytest_codspeed/eval_harness.py

@@ -0,0 +1,115 @@
+"""Scoring utilities for evaluating LLM-generated optimisation suggestions.
+
+After two consecutive walltime runs with ``--codspeed-capture-output``, the
+local comparison report includes both a performance delta and a correctness
+signal (``output_changed``).  This module turns those two signals into a
+single scalar score that can be used to rank or accept/reject suggestions
+automatically.
+
+Score semantics
+---------------
+* ``1.0``  - 100 % faster, output unchanged (theoretical maximum).
+* ``0.33`` - 33 % faster, output unchanged.
+* ``0.0``  - any speedup, but output changed (correctness broken).
+* ``nan``  - run was not collected with ``--codspeed-capture-output``;
+             correctness is unknown, score cannot be computed.
+
+Usage in an optimisation loop::
+
+    from pytest_codspeed.comparison import compare_results
+    from pytest_codspeed.eval_harness import compute_score, EvalReport
+
+    report = compare_results(baseline_path, current_path)
+    eval_report = EvalReport.from_comparison(report)
+    for entry in eval_report.entries:
+        print(entry.name, entry.score)
+"""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from pytest_codspeed.comparison import BenchmarkDiff, ComparisonReport
+
+
+def compute_score(perf_gain: float, output_changed: bool | None) -> float:
+    """Return a scalar quality score for a single optimisation suggestion.
+
+    Args:
+        perf_gain: Relative performance improvement, positive = faster.
+                   Computed as ``-diff.change_ratio`` (so +0.33 means 33 %
+                   faster).
+        output_changed: ``True`` if the function's return value changed,
+                        ``False`` if it stayed the same, ``None`` if the
+                        capture flag was not used.
+
+    Returns:
+        * ``0.0`` when ``output_changed is True`` -- correctness is broken.
+        * ``float('nan')`` when ``output_changed is None`` -- unknown.
+        * ``max(0.0, perf_gain)`` otherwise -- reward speed, ignore slowdowns.
+    """
+    if output_changed is True:
+        return 0.0
+    if output_changed is None:
+        return float("nan")
+    return max(0.0, perf_gain)
+
+
+@dataclass(frozen=True)
+class EvalEntry:
+    """Score for a single benchmark in an optimisation suggestion."""
+
+    name: str
+    perf_gain: float
+    output_changed: bool | None
+    score: float
+
+    @property
+    def is_acceptable(self) -> bool:
+        """True when the suggestion improves perf without breaking correctness."""
+        return not math.isnan(self.score) and self.score > 0.0
+
+    @property
+    def correctness_unknown(self) -> bool:
+        return self.output_changed is None
+
+
+@dataclass(frozen=True)
+class EvalReport:
+    """Aggregated evaluation of an optimisation suggestion across all benchmarks."""
+
+    entries: tuple[EvalEntry, ...]
+
+    @classmethod
+    def from_comparison(cls, report: ComparisonReport) -> EvalReport:
+        """Build an :class:`EvalReport` from a :class:`ComparisonReport`."""
+        all_diffs: tuple[BenchmarkDiff, ...] = (
+            *report.regressions,
+            *report.improvements,
+            *report.unchanged,
+        )
+        entries = tuple(
+            EvalEntry(
+                name=diff.name,
+                perf_gain=-diff.change_ratio,
+                output_changed=diff.output_changed,
+                score=compute_score(-diff.change_ratio, diff.output_changed),
+            )
+            for diff in all_diffs
+        )
+        return cls(entries=entries)
+
+    @property
+    def acceptable(self) -> tuple[EvalEntry, ...]:
+        return tuple(e for e in self.entries if e.is_acceptable)
+
+    @property
+    def correctness_broken(self) -> tuple[EvalEntry, ...]:
+        return tuple(e for e in self.entries if e.output_changed is True)
+
+    @property
+    def correctness_unknown(self) -> tuple[EvalEntry, ...]:
+        return tuple(e for e in self.entries if e.correctness_unknown)

src/pytest_codspeed/instruments/walltime.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 
+import hashlib
 import os
+import pickle
 import warnings
 from dataclasses import asdict, dataclass
 from math import ceil
@@ -146,13 +148,25 @@ def from_list(
         )
 
 
+def _hash_output(value: object) -> str | None:
+    """Return a short hex digest of *value*, or ``None`` if un-serialisable."""
+    try:
+        raw = pickle.dumps(value, protocol=pickle.HIGHEST_PROTOCOL)
+    except Exception:
+        try:
+            raw = repr(value).encode()
+        except Exception:
+            return None
+    return hashlib.sha256(raw).hexdigest()[:16]
+
+
 @dataclass
 class Benchmark:
     name: str
     uri: str
-
     config: BenchmarkConfig
     stats: BenchmarkStats
+    output_hash: str | None = None
 
 
 class WallTimeInstrument(Instrument):
@@ -200,6 +214,9 @@ def __codspeed_root_frame__() -> T:
 
         # Compute the actual result of the function
         out = __codspeed_root_frame__()
+        output_hash = (
+            _hash_output(out) if self.config.capture_output else None
+        )
 
         # Warmup
         times_per_round_ns: list[float] = []
@@ -258,7 +275,13 @@ def __codspeed_root_frame__() -> T:
         )
 
         self.benchmarks.append(
-            Benchmark(name=name, uri=uri, config=benchmark_config, stats=stats)
+            Benchmark(
+                name=name,
+                uri=uri,
+                config=benchmark_config,
+                stats=stats,
+                output_hash=output_hash,
+            )
         )
         return out
 
@@ -316,11 +339,20 @@ def __codspeed_root_frame__(*args, **kwargs) -> T:
         # Compute the actual result of the function
         args, kwargs = pedantic_options.setup_and_get_args_kwargs()
         out = __codspeed_root_frame__(*args, **kwargs)
+        output_hash = (
+            _hash_output(out) if self.config.capture_output else None
+        )
         if pedantic_options.teardown is not None:
             pedantic_options.teardown(*args, **kwargs)
 
         self.benchmarks.append(
-            Benchmark(name=name, uri=uri, config=benchmark_config, stats=stats)
+            Benchmark(
+                name=name,
+                uri=uri,
+                config=benchmark_config,
+                stats=stats,
+                output_hash=output_hash,
+            )
         )
         return out
 

src/pytest_codspeed/plugin.py

@@ -80,6 +80,15 @@ def pytest_addoption(parser: pytest.Parser):
             ", only for walltime mode"
         ),
     )
+    group.addoption(
+        "--codspeed-capture-output",
+        action="store_true",
+        default=False,
+        help=(
+            "Hash the return value of each benchmark and store it in the "
+            "result file, enabling correctness checks in local comparisons"
+        ),
+    )
 
 
 @dataclass(unsafe_hash=True)