# /// script
# requires-python = ">=3.10"
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""SPSC ring buffer: head and tail on one cache line, or two.

One producer and one consumer move 100 million longs through a ring:
head and tail in the same cache line, on separate lines (128 bytes apart,
the line size on Apple silicon), and with batched publishing (the producer
updates tail once per k items). The C program is compiled and run by this
script (needs cc).
"""
import argparse, json, os, subprocess, tempfile

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--items", type=int, default=100_000_000, help="Items moved.")
p.add_argument("--batches", type=lambda s: [int(v) for v in s.split(",")], default=[1, 2, 4, 8, 16, 64], help="Batch sizes k.")
p.add_argument("--repeats", type=int, default=3, help="Runs per variant; the median is reported.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

C = r"""
#include <pthread.h>
#include <stdatomic.h>
#include <stdio.h>
#include <time.h>
#define CAP 4096
#ifdef SAME_LINE
struct { _Atomic long head; _Atomic long tail; } q;
#else
struct { _Alignas(128) _Atomic long head; _Alignas(128) _Atomic long tail; } q;
#endif
long buf[CAP];
static void *producer(void *a) {
    long t = 0, pub = 0;
    for (long i = 0; i < N; i++) {
        while (t - atomic_load_explicit(&q.head, memory_order_acquire) >= CAP) ;
        buf[t % CAP] = i; t++;
        if (t - pub >= K || i == N - 1) { atomic_store_explicit(&q.tail, t, memory_order_release); pub = t; }
    }
    return NULL;
}
static void *consumer(void *a) {
    long h = 0, sum = 0;
    while (h < N) {
        long t = atomic_load_explicit(&q.tail, memory_order_acquire);
        while (h < t) { sum += buf[h % CAP]; h++; }
        atomic_store_explicit(&q.head, h, memory_order_release);
    }
    return (void *)sum;
}
int main(void) {
    pthread_t p, c; struct timespec a, b;
    clock_gettime(CLOCK_MONOTONIC, &a);
    pthread_create(&c, NULL, consumer, NULL); pthread_create(&p, NULL, producer, NULL);
    pthread_join(p, NULL); pthread_join(c, NULL);
    clock_gettime(CLOCK_MONOTONIC, &b);
    printf("%.4f\n", (b.tv_sec - a.tv_sec) + (b.tv_nsec - a.tv_nsec) / 1e9);
}
"""
d = tempfile.mkdtemp(); src = os.path.join(d, "ring.c"); open(src, "w").write(C)
def run(label, flags, k):
    exe = os.path.join(d, "ring" + str(abs(hash(label))))
    subprocess.run(["cc", "-O2", "-pthread", src, "-o", exe, f"-DN={args.items}L", f"-DK={k}", *flags], check=True)
    times = sorted(float(subprocess.run([exe], capture_output=True, text=True, check=True).stdout) for _ in range(args.repeats))
    s = times[len(times) // 2]
    return {"variant": label, "k": k, "million_per_s": args.items / s / 1e6}
rows = [run("head and tail on one line", ["-DSAME_LINE"], 1)]
rows += [run(f"separate lines, batch {k}", [], k) for k in args.batches]
for r in rows: print(f"{r['variant']:<28} {r['million_per_s']:8.1f} million items/s")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"variants": rows, "batches": [r for r in rows if r["variant"].startswith("separate")]}}, f)