# /// script
# requires-python = ">=3.10"
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Arena versus malloc: the 50x claim, and fragmentation.

Experiment 1: allocate and free 100 million 24-byte objects with one
malloc/free each, then with a bump-pointer arena reset after every batch of
1,000,000. Experiment 2: allocate objects of random sizes, free a random half
in random order, allocate again, and print the process's resident memory next
to the live bytes; then the same with an arena. The C program is compiled
(-O2) and run by this script (needs cc).
"""
import argparse, json, os, subprocess, tempfile

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--batches", type=int, default=100, help="Experiment 1: batches of 1,000,000 objects.")
p.add_argument("--objects", type=int, default=500_000, help="Experiment 2: objects of random size.")
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 <string.h>
#include <time.h>
#include <sys/mman.h>
#ifdef __APPLE__
#include <mach/mach.h>
static double rss_mb(void){ task_vm_info_data_t i; mach_msg_type_number_t n = TASK_VM_INFO_COUNT;   /* footprint: counts compressed pages too */
  task_info(mach_task_self(), TASK_VM_INFO, (task_info_t)&i, &n); return i.phys_footprint / 1e6; }
#else
#include <unistd.h>
static double rss_mb(void){ long pages = 0, r = 0; FILE *f = fopen("/proc/self/statm", "r"); fscanf(f, "%ld %ld", &pages, &r); fclose(f);
  return r * sysconf(_SC_PAGESIZE) / 1e6; }
#endif
typedef struct { long v; long a; long b; } Item;                     /* 24 bytes */
typedef struct { char *base; size_t used, cap; } Arena;
static void *arena_alloc(Arena *a, size_t n){ n = (n + 15) & ~(size_t)15;
  if (a->used + n > a->cap) return NULL; void *p = a->base + a->used; a->used += n; return p; }
static void arena_reset(Arena *a){ a->used = 0; }
static double now(void){ struct timespec t; clock_gettime(CLOCK_MONOTONIC, &t); return t.tv_sec + t.tv_nsec / 1e9; }
static unsigned long long s = 88172645463325252ULL;
static unsigned long long rnd(void){ s ^= s << 13; s ^= s >> 7; s ^= s << 17; return s; }
static void fill(char *p, size_t n){ for (size_t i = 0; i + 8 <= n; i += 8) { unsigned long long r = rnd(); memcpy(p + i, &r, 8); } }  /* random bytes: nothing to compress */
int main(int argc, char **argv){
  int batches = atoi(argv[1]); long objs = atol(argv[2]); long sink = 0;
  /* experiment 1 */
  static Item *keep[1000000];
  double t = now();
  for (int b = 0; b < batches; b++) {
    for (int i = 0; i < 1000000; i++) { Item *x = malloc(sizeof(Item)); x->v = i; keep[i] = x; }
    for (int i = 0; i < 1000000; i++) { sink += keep[i]->v; free(keep[i]); }
  }
  double t_malloc = now() - t;
  Arena a = { malloc(1000000 * 32), 0, 1000000 * 32 };
  t = now();
  for (int b = 0; b < batches; b++) {
    for (int i = 0; i < 1000000; i++) { Item *x = arena_alloc(&a, sizeof(Item)); x->v = i; keep[i] = x; }
    for (int i = 0; i < 1000000; i++) sink += keep[i]->v;
    arena_reset(&a);                                                /* frees all 1,000,000 objects at once */
  }
  double t_arena = now() - t;
  printf("alloc %.6f %.6f %ld\n", t_malloc / (batches * 1e6) * 1e9, t_arena / (batches * 1e6) * 1e9, sink & 1);
  /* experiment 2: random sizes, free a random half, allocate again */
  char **ptr = malloc(objs * sizeof(char *)); size_t *sz = malloc(objs * sizeof(size_t)); double live = 0;
  double base = rss_mb();
  for (long i = 0; i < objs; i++) { sz[i] = 16 + rnd() % 1024; ptr[i] = malloc(sz[i]); fill(ptr[i], sz[i]); live += sz[i]; }
  double rss_full = rss_mb() - base, live_full = live / 1e6;
  for (long i = objs - 1; i > 0; i--) { long j = rnd() % (i + 1); char *tp = ptr[i]; ptr[i] = ptr[j]; ptr[j] = tp; size_t ts = sz[i]; sz[i] = sz[j]; sz[j] = ts; }
  for (long i = 0; i < objs / 2; i++) { free(ptr[i]); live -= sz[i]; ptr[i] = NULL; }
  double rss_half = rss_mb() - base, live_half = live / 1e6;
  for (long i = 0; i < objs / 4; i++) { sz[i] = 2048 + rnd() % 2048; ptr[i] = malloc(sz[i]); fill(ptr[i], sz[i]); live += sz[i]; }  /* larger than the holes */
  double rss_again = rss_mb() - base, live_again = live / 1e6;
  double after_free = 0; (void)after_free;
  printf("frag malloc %.1f %.1f %.1f %.1f %.1f %.1f\n", live_full, rss_full, live_half, rss_half, live_again, rss_again);
  for (long i = 0; i < objs; i++) free(ptr[i]);
  /* the same lifetimes with arenas: one arena per generation, dropped whole when its generation dies */
  base = rss_mb(); size_t cap = (size_t)objs * 1100;
  Arena g1 = { mmap(NULL, cap, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANON, -1, 0), 0, cap }; live = 0;   /* arenas take their chunks from the OS */
  for (long i = 0; i < objs; i++) { size_t n = 16 + rnd() % 1024; char *q = arena_alloc(&g1, n); fill(q, n); sink += q[n - 1]; live += n; }
  rss_full = rss_mb() - base; live_full = live / 1e6;
  munmap(g1.base, cap); live_half = live_full / 2; live = 0;               /* a generation dies together: the arena is released */
  Arena g2 = { mmap(NULL, cap / 2, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANON, -1, 0), 0, cap / 2 };
  for (long i = 0; i < objs / 2; i++) { size_t n = 16 + rnd() % 1024; char *q = arena_alloc(&g2, n); fill(q, n); sink += q[n - 1]; live += n; }
  rss_half = rss_mb() - base; live_half = live / 1e6;
  printf("frag arena %.1f %.1f %.1f %.1f\n", live_full, rss_full, live_half, rss_half);
  fprintf(stderr, "%ld\n", sink & 1);                                /* read the bytes, so the writes cannot be dropped */
  return 0;
}
"""
d = tempfile.mkdtemp(); src = os.path.join(d, "arena.c"); exe = os.path.join(d, "arena"); open(src, "w").write(C)
subprocess.run(["cc", "-O2", src, "-o", exe], check=True)
out = subprocess.run([exe, str(args.batches), str(args.objects)], capture_output=True, text=True, check=True).stdout.splitlines()
a = out[0].split(); fm = list(map(float, out[1].split()[2:])); fa = list(map(float, out[2].split()[2:]))
alloc = [{"allocator": "malloc/free per object", "ns_per_alloc": float(a[1])}, {"allocator": "bump arena, reset per batch", "ns_per_alloc": float(a[2])}]
frag = [{"phase": "malloc: all allocated", "live_mb": fm[0], "rss_mb": fm[1]},
        {"phase": "malloc: random half freed", "live_mb": fm[2], "rss_mb": fm[3]},
        {"phase": "malloc: larger objects allocated again", "live_mb": fm[4], "rss_mb": fm[5]},
        {"phase": "arena: all allocated", "live_mb": fa[0], "rss_mb": fa[1]},
        {"phase": "arena: generation dropped, half reallocated", "live_mb": fa[2], "rss_mb": fa[3]}]
for r in frag: r["gap_mb"] = r["rss_mb"] - r["live_mb"]
summary = {"ratio": alloc[0]["ns_per_alloc"] / alloc[1]["ns_per_alloc"]}
for r in alloc: print(f"{r['allocator']:<28} {r['ns_per_alloc']:.2f} ns per allocation")
print(f"arena / malloc speed-up: {summary['ratio']:.0f}x")
for r in frag: print(f"{r['phase']:<44} live {r['live_mb']:7.1f} MB   resident {r['rss_mb']:7.1f} MB   gap {r['gap_mb']:7.1f} MB")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"summary": summary, "tables": {"alloc": alloc, "fragmentation": frag}}, f)