# /// script
# requires-python = ">=3.10"
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Sorted against shuffled: the branch benchmark.

Sums the elements >= 128 of 16 million random bytes-as-ints, 20 times, first
shuffled and then sorted. The work is identical; only the branch's
predictability changes. Compiled at -O1, -O2 and -O3 (where the compiler may
make the loop branchless or vectorize it), and at -O2 with a branch the
compiler cannot remove. Needs cc.
"""
import argparse, json, os, subprocess, tempfile

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--opts", default="-O1,-O2,-O3", help="Optimization levels to compile with.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

C = r"""
#include <stdio.h>
#include <stdlib.h>
#include <time.h>
#define N (1 << 24)
static int data[N];
static int cmp(const void* a, const void* b) { return *(int*)a - *(int*)b; }
static double run(void) {
    struct timespec a, b; clock_gettime(CLOCK_MONOTONIC, &a);
    long sum = 0;
    for (int rep = 0; rep < 20; rep++)
        for (int i = 0; i < N; i++)
#ifdef FORCE_BRANCH
            if (data[i] >= 128) { sum += data[i]; __asm__ volatile(""); }   /* a side effect the compiler cannot if-convert */
#else
            if (data[i] >= 128) sum += data[i];
#endif
    clock_gettime(CLOCK_MONOTONIC, &b);
    fprintf(stderr, "  sum=%ld\n", sum);                       /* stops the compiler removing the loop */
    return (b.tv_sec - a.tv_sec) + (b.tv_nsec - a.tv_nsec) / 1e9;
}
int main(void) {
    for (int i = 0; i < N; i++) data[i] = rand() % 256;
    double s = run();
    qsort(data, N, sizeof(int), cmp);
    printf("%f %f\n", s, run());
}
"""
d = tempfile.mkdtemp(); src = os.path.join(d, "branch.c"); open(src, "w").write(C); rows = []
builds = [(o, []) for o in args.opts.split(",")] + [("-O2", ["-DFORCE_BRANCH"])]
for opt, extra in builds:
    exe = os.path.join(d, "b" + opt + "".join(extra)); subprocess.run(["cc", opt, *extra, src, "-o", exe], check=True)
    shuffled, sorted_ = map(float, subprocess.run([exe], capture_output=True, text=True, check=True).stdout.split())
    asm = subprocess.run(["cc", opt, *extra, "-S", "-o", "-", src], capture_output=True, text=True).stdout
    branchless = any(x in asm for x in ("csel", "cmov", "cmhs", "cmgt", "vpcmp"))     # select or vector compare: no branch
    rows.append({"opt": opt + (" + forced branch" if extra else ""), "shuffled_s": shuffled, "sorted_s": sorted_, "ratio": shuffled / sorted_,
                 "branchless": "yes" if branchless else "no"})
for r in rows: print(f"{r['opt']}: shuffled {r['shuffled_s']:.2f} s, sorted {r['sorted_s']:.2f} s -> {r['ratio']:.1f}x   (conditional select in the code: {r['branchless']})")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"builds": rows}}, f)