# /// script
# requires-python = ">=3.10"
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Persistent tree: measure the cost of immutability.

1 million random keys into a persistent (path-copying) binary search tree
and into a mutable one: time and memory. A snapshot held after 500,000
inserts still sees only the first half; then the nodes the final tree shares
with that snapshot.
"""
import argparse, json, random, sys, time, tracemalloc

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--keys", type=int, default=1_000_000, help="Keys inserted.")
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)
sys.setrecursionlimit(10_000)
keys = random.sample(range(10 * args.keys), args.keys)

class Node:
    __slots__ = ("k", "l", "r")
    def __init__(self, k, l=None, r=None): self.k, self.l, self.r = k, l, r

def p_insert(n, k):                                   # persistent: copy the path, share the rest
    if n is None: return Node(k)
    if k < n.k: return Node(n.k, p_insert(n.l, k), n.r)
    if k > n.k: return Node(n.k, n.l, p_insert(n.r, k))
    return n

def m_insert(root, k):                                # mutable: change one pointer in place
    if root is None: return Node(k)
    n = root
    while True:
        if k < n.k:
            if n.l is None: n.l = Node(k); return root
            n = n.l
        elif k > n.k:
            if n.r is None: n.r = Node(k); return root
            n = n.r
        else: return root

def contains(n, k):
    while n:
        if k == n.k: return True
        n = n.l if k < n.k else n.r
    return False

rows = []
for name, ins in (("mutable", m_insert), ("persistent", p_insert)):
    tracemalloc.start(); t = time.perf_counter(); root = None; snapshot = None
    for i, k in enumerate(keys):
        root = ins(root, k)
        if i == args.keys // 2 - 1: snapshot = root
    dt = time.perf_counter() - t; peak = tracemalloc.get_traced_memory()[1]; tracemalloc.stop()
    rows.append({"tree": name, "seconds": dt, "peak_mb": peak / 1e6})
    if name == "persistent": p_root, p_snap = root, snapshot
    print(f"{name:<11} {dt:6.2f} s   peak memory {peak / 1e6:7.1f} MB")

first, second = keys[: args.keys // 2], keys[args.keys // 2:]
sees_first = all(contains(p_snap, k) for k in first[:10_000]); sees_second = any(contains(p_snap, k) for k in second[:10_000])
def ids(n, acc):
    stack = [n]
    while stack:
        x = stack.pop()
        if x: acc.add(id(x)); stack += [x.l, x.r]
    return acc
final_nodes, snap_nodes = ids(p_root, set()), ids(p_snap, set())
shared = len(final_nodes & snap_nodes)
print(f"snapshot after {args.keys // 2:,} inserts sees the first half: {sees_first}; sees any of the second: {sees_second}")
print(f"the final tree shares {shared:,} of the snapshot's {len(snap_nodes):,} nodes ({shared / len(snap_nodes):.0%})")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"trees": rows, "summary": [{"slowdown": rows[1]["seconds"] / rows[0]["seconds"],
                   "memory_x": rows[1]["peak_mb"] / rows[0]["peak_mb"], "sees_first": "yes" if sees_first else "no",
                   "sees_second": "yes" if sees_second else "no", "shared": shared, "snapshot_nodes": len(snap_nodes),
                   "shared_pct": shared / len(snap_nodes) * 100}]}}, f)