# /// script
# requires-python = ">=3.10"
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Varint and zigzag: where the bet pays and where it loses.

Bytes per value for varint on four kinds of data, with and without zigzag
for signed values, against a fixed 8-byte int64.
"""
import argparse, json, random

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--values", type=int, default=100_000, help="Values per set.")
p.add_argument("--seed", type=int, default=0, help="Random seed.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()
random.seed(args.seed)
def varint_len(u):
    n = 1
    while u >= 0x80: u >>= 7; n += 1
    return n
zigzag = lambda n: ((n << 1) ^ (n >> 63)) & ((1 << 64) - 1)
twos = lambda n: n & ((1 << 64) - 1)
N = args.values
sets = {"small counts 0..100": [random.randrange(100) for _ in range(N)],
        "ids up to 10^9": [random.randrange(10**9) for _ in range(N)],
        "random 64-bit": [random.getrandbits(64) for _ in range(N)],
        "small signed -50..50": [random.randint(-50, 50) for _ in range(N)]}
rows = []
for name, v in sets.items():
    plain = sum(varint_len(twos(x)) for x in v) / N; zz = sum(varint_len(zigzag(x)) for x in v) / N
    rows.append({"data": name, "varint_bytes": plain, "zigzag_bytes": zz, "fixed_bytes": 8})
    print(f"{name:24s} varint {plain:5.2f} B   zigzag+varint {zz:5.2f} B   (fixed int64 = 8 B)")

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