#!/usr/bin/env bash
set -euo pipefail

# Allocator comparison matrix using the SAME benchmark binary via LD_PRELOAD.
#
# Why:
# - Different binaries introduce layout tax (text size/I-cache) and can make hakmem look much worse/better.
# - This script uses `bench_random_mixed_system` as the single fixed binary and swaps allocators via LD_PRELOAD.
#
# What it runs:
# - system (no LD_PRELOAD)
# - hakmem (LD_PRELOAD=./libhakmem.so)
# - mimalloc (LD_PRELOAD=$MIMALLOC_SO) if provided
# - jemalloc (LD_PRELOAD=$JEMALLOC_SO) if provided
# - tcmalloc (LD_PRELOAD=$TCMALLOC_SO) if provided
#
# SSOT alignment:
# - Applies the same "cleanenv defaults" as `scripts/run_mixed_10_cleanenv.sh`.
# - IMPORTANT: never LD_PRELOAD the shell/script itself; apply LD_PRELOAD only to the benchmark binary exec.
#
# Usage:
#   make bench_random_mixed_system shared
#   export MIMALLOC_SO=/path/to/libmimalloc.so.2      # optional
#   export JEMALLOC_SO=/path/to/libjemalloc.so.2      # optional
#   export TCMALLOC_SO=/path/to/libtcmalloc.so        # optional
#   RUNS=10 scripts/run_allocator_preload_matrix.sh
#
# Tunables:
#   HAKMEM_PROFILE=MIXED_TINYV3_C7_SAFE ITERS=20000000 WS=400 RUNS=10

root_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
cd "${root_dir}"

profile="${HAKMEM_PROFILE:-MIXED_TINYV3_C7_SAFE}"
iters="${ITERS:-20000000}"
ws="${WS:-400}"
runs="${RUNS:-10}"

if [[ ! -x ./bench_random_mixed_system ]]; then
  echo "[preload-matrix] Missing ./bench_random_mixed_system (build via: make bench_random_mixed_system)" >&2
  exit 1
fi
extract_throughput() {
  rg -o "Throughput = +[0-9]+ ops/s" | rg -o "[0-9]+"
}

stats_py='
import statistics,sys
xs=[int(x) for x in sys.stdin.read().strip().split() if x.strip()]
if not xs:
  sys.exit(1)
xs_sorted=sorted(xs)
mean=sum(xs)/len(xs)
median=statistics.median(xs_sorted)
stdev=statistics.pstdev(xs) if len(xs)>1 else 0.0
cv=(stdev/mean*100.0) if mean>0 else 0.0
print(f"runs={len(xs)} mean={mean/1e6:.2f}M median={median/1e6:.2f}M cv={cv:.2f}% min={min(xs)/1e6:.2f}M max={max(xs)/1e6:.2f}M")
'

apply_cleanenv_defaults() {
  # Keep reproducible even if user exported env vars.
  case "${profile}" in
    MIXED_TINYV3_C7_BALANCED)
      export HAKMEM_SS_MEM_LEAN=1
      export HAKMEM_SS_MEM_LEAN_DECOMMIT=OFF
      export HAKMEM_SS_MEM_LEAN_TARGET_MB=10
      ;;
    *)
      export HAKMEM_SS_MEM_LEAN=0
      export HAKMEM_SS_MEM_LEAN_DECOMMIT=OFF
      export HAKMEM_SS_MEM_LEAN_TARGET_MB=10
      ;;
  esac

  # Force known research knobs OFF to avoid accidental carry-over.
  export HAKMEM_TINY_HEADER_WRITE_ONCE=0
  export HAKMEM_TINY_C7_PRESERVE_HEADER=0
  export HAKMEM_TINY_TCACHE=0
  export HAKMEM_TINY_TCACHE_CAP=64
  export HAKMEM_MALLOC_TINY_DIRECT=0
  export HAKMEM_FRONT_FASTLANE_ALLOC_LEGACY_DIRECT=0
  export HAKMEM_FORCE_LIBC_ALLOC=0
  export HAKMEM_ENV_SNAPSHOT_SHAPE=0
  export HAKMEM_TINY_C7_ULTRA_HEADER_LIGHT=0
  export HAKMEM_TINY_C2_LOCAL_CACHE=0
  export HAKMEM_TINY_INLINE_SLOTS_SWITCHDISPATCH_FIXED=0

  # Keep cleanenv aligned with promoted knobs.
  export HAKMEM_FASTLANE_DIRECT=1
  export HAKMEM_FREE_TINY_FAST_MONO_DUALHOT=1
  export HAKMEM_FREE_TINY_FAST_MONO_LEGACY_DIRECT=1
  export HAKMEM_WARM_POOL_SIZE=16
  export HAKMEM_TINY_C4_INLINE_SLOTS=1
  export HAKMEM_TINY_C5_INLINE_SLOTS=1
  export HAKMEM_TINY_C6_INLINE_SLOTS=1
  export HAKMEM_TINY_INLINE_SLOTS_FIXED=1
  export HAKMEM_TINY_INLINE_SLOTS_SWITCHDISPATCH=1
}

run_preload_n() {
  local label="$1"
  local preload="$2"

  echo ""
  echo "== ${label} (profile=${profile}) =="

  apply_cleanenv_defaults

  for i in $(seq 1 "${runs}"); do
    if [[ -n "${preload}" ]]; then
      local preload_abs
      preload_abs="$(realpath "${preload}")"
      # Apply LD_PRELOAD ONLY to the benchmark binary exec (not to bash/rg/python).
      HAKMEM_PROFILE="${profile}" LD_PRELOAD="${preload_abs}" \
        ./bench_random_mixed_system "${iters}" "${ws}" 1 2>&1 | extract_throughput || true
    else
      HAKMEM_PROFILE="${profile}" \
        ./bench_random_mixed_system "${iters}" "${ws}" 1 2>&1 | extract_throughput || true
    fi
  done | python3 -c "${stats_py}"
}

run_preload_n "system (no preload)" ""

if [[ -x ./libhakmem.so ]]; then
  run_preload_n "hakmem (LD_PRELOAD libhakmem.so)" ./libhakmem.so
else
  echo ""
  echo "== hakmem (LD_PRELOAD libhakmem.so) =="
  echo "skipped (missing ./libhakmem.so; build via: make shared)"
fi

if [[ -n "${MIMALLOC_SO:-}" && -e "${MIMALLOC_SO}" ]]; then
  run_preload_n "mimalloc (LD_PRELOAD)" "${MIMALLOC_SO}"
fi
if [[ -n "${JEMALLOC_SO:-}" && -e "${JEMALLOC_SO}" ]]; then
  run_preload_n "jemalloc (LD_PRELOAD)" "${JEMALLOC_SO}"
fi
if [[ -n "${TCMALLOC_SO:-}" && -e "${TCMALLOC_SO}" ]]; then
  run_preload_n "tcmalloc (LD_PRELOAD)" "${TCMALLOC_SO}"
fi