#!/bin/bash
# verify-preview.sh — check, without source code, whether the hand-written code paths cited in
# ClickHouse/ClickBench#2165 (irprobe/src/q*.rs, tpch_q01.rs) execute while the submitted pgrust v0.4-preview binary runs ClickBench.
#
# Run on a fresh Ubuntu 24.04 arm64 machine (c8g.4xlarge, 500 GB disk), as a sudo-capable user:
#     bash verify-preview.sh
# Takes ~30 min (most of it downloading hits.parquet and loading 100M rows).
#
# Method (no symbols, no source):
#   1. the binary embeds Rust panic-location records {file path, line, col}; for each cited source
#      file we find its path string, the records pointing at it, the code that loads those records,
#      and (from the unwind tables, present even when stripped) the function containing that code;
#   2. we put a kernel uprobe on the entry of every such function and count hits;
#   3. we run the 43 ClickBench queries x 3 as submitted, then a CONTROL that forces the
#      hand-written path (pgrust.gravity_admit=shape) — the probes must fire there, or they prove nothing.
# Limits: a file with no panic sites left in the binary cannot be located; a function containing
# panic sites from several files is reported as "shared" (ambiguous), not as belonging to one file.
set -euo pipefail

BIN_URL=https://pgrust.com/downloads/archive/2026-09-clickbench-snapshot/pgrust-v0.4-preview-aarch64-neoverse-v2.tar.gz
BIN_SHA=a8b08f029255a2f27364c51dae119fba6aff4beb20fa58125345527e78d3b2e3
CLICKBENCH_COMMIT=964e362a891f2ac19d0a896a642758555341ab74   # ClickBench main just before the pgrust entry was removed
# the files listed in ClickBench#2165 as "one hand-written kernel per benchmark query"
CITED_RE='crates/gravity/backend/irprobe/src/(q[0-9]+|tpch_q01)\.rs'
W=$HOME/verify-preview; mkdir -p "$W"; cd "$W"

echo "== 1. tools + the submission's own ClickBench scripts"
sudo apt-get update -qq && sudo apt-get install -y -qq git binutils python3 curl >/dev/null
[ -d ClickBench ] || git clone -q https://github.com/ClickHouse/ClickBench.git
git -C ClickBench checkout -q "$CLICKBENCH_COMMIT"
cd ClickBench/pgrust

echo "== 2. the submitted binary (pinned by sha256)"
curl -fsSL -o "$W/preview.tar.gz" "$BIN_URL"
echo "$BIN_SHA  $W/preview.tar.gz" | sha256sum -c -
export PGRUST_TARBALL="$W/preview.tar.gz"
sudo chmod a+r "$W/preview.tar.gz"; chmod a+x "$HOME" "$W"

echo "== 3. install, start, load (the submission's scripts, unmodified)"
./install && ./start && for i in $(seq 1 300); do ./check >/dev/null 2>&1 && break; sleep 1; done && ./load
. ./env.sh
BIN="$PGRUST_PREFIX/bin/postgres"
sha256sum "$BIN"

echo "== 4. locate code from the cited files via panic-location records"
cat > "$W/locate.py" <<'PY'
import bisect, re, struct, sys
from collections import defaultdict
data = open(sys.argv[1], "rb").read(); cited = re.compile(sys.argv[2])
shoff, = struct.unpack_from("<Q", data, 0x28); shentsize, shnum, shstrndx = struct.unpack_from("<HHH", data, 0x3A)
secs = [struct.unpack_from("<IIQQQQ", data, shoff + i * shentsize) for i in range(shnum)]
st = secs[shstrndx][4]
S = {data[st + s[0]: data.index(b"\0", st + s[0])].decode(): (s[3], s[4], s[5]) for s in secs}
def v2o(va):
    for a, o, sz in S.values():
        if a and a <= va < a + sz: return va - a + o
ra, ro, rs = S[".rodata"]; rod = data[ro:ro + rs]
strs = {}                                   # string vaddr -> path, for EVERY source path (so mixed functions are detectable)
for m in re.finditer(rb"crates/[A-Za-z0-9_./-]*?\.rs", rod): strs[ra + m.start()] = m.group(0).decode()
xa, xo, xs = S[".rela.dyn"]; recs = {}      # Location record vaddr -> path
for i in range(0, xs, 24):
    off, info, add = struct.unpack_from("<QQq", data, xo + i)
    if (info & 0xffffffff) == 1027 and add in strs:          # R_AARCH64_RELATIVE -> path string
        fo = v2o(off)
        if fo is not None and struct.unpack_from("<Q", data, fo + 8)[0] == len(strs[add]): recs[off] = strs[add]
pages = defaultdict(set)
for r in recs: pages[r & ~0xfff].add(r)
ta, to, ts = S[".text"]; n = ts // 4; ins = struct.unpack_from(f"<{n}I", data, to); refs = []
for i in range(n):                          # ADRP Xd, page ; ADD Xd', Xd, #lo12  ->  record
    w = ins[i]
    if (w & 0x9F000000) != 0x90000000: continue
    pc = ta + 4 * i; imm = ((w >> 5) & 0x7FFFF) << 2 | ((w >> 29) & 3)
    if imm & (1 << 20): imm -= 1 << 21
    page = (pc & ~0xfff) + (imm << 12)
    if page not in pages: continue
    for j in range(i + 1, min(n, i + 64)):
        x = ins[j]
        if (x & 0xFF800000) == 0x91000000 and ((x >> 5) & 0x1f) == (w & 0x1f):
            t = page + (((x >> 10) & 0xfff) << (12 if (x >> 22) & 1 else 0))
            if t in pages[page]: refs.append((ta + 4 * j, t))
ha, ho, hs = S[".eh_frame_hdr"]; cnt, = struct.unpack_from("<I", data, ho + 8)
tab = [struct.unpack_from("<ii", data, ho + 12 + 8 * k) for k in range(cnt)]
starts = [ha + a for a, _ in tab]; fdes = [ha + f for _, f in tab]
files = defaultdict(set)                    # function start -> {cited files}
for iva, rec in refs:
    k = bisect.bisect_right(starts, iva) - 1
    if k < 0: continue
    fo = v2o(fdes[k]); b = fdes[k] + 8 + struct.unpack_from("<i", data, fo + 8)[0]; rng, = struct.unpack_from("<I", data, fo + 12)
    if b == starts[k] and b <= iva < b + rng: files[b].add(recs[rec])
keep = {b: fs for b, fs in files.items() if any(cited.fullmatch(p) for p in fs)}
cited_paths = sorted({p for p in strs.values() if cited.fullmatch(p)})
sys.stderr.write(f"cited paths in binary {len(cited_paths)}, functions containing cited code {len(keep)}\n")
located = {p for fs in keep.values() for p in fs}
for p in cited_paths:
    if p not in located: sys.stderr.write(f"not locatable (no panic sites left): {p}\n")
# column 2: the cited file(s); column 3: "own" if every panic site in the function is from cited files, else "shared"
for b, fs in sorted(keep.items()):
    c = sorted(p for p in fs if cited.fullmatch(p))
    print(f"{b:#x}\t{','.join(c)}\t{'own' if len(c) == len(fs) and len(c) == 1 else 'shared'}")
PY
python3 -I "$W/locate.py" "$BIN" "$CITED_RE" > "$W/functions.tsv" 2> "$W/locate.txt"; cat "$W/locate.txt"

echo "== 5. one uprobe per located function"
cat > "$W/probes.py" <<'PY'
import os, subprocess, sys
from collections import defaultdict
T = "/sys/kernel/tracing"
def tw(p, s):
    fd = os.open(p, os.O_WRONLY | os.O_APPEND)
    try: os.write(fd, s.encode())
    finally: os.close(fd)
def clear():
    try: tw(f"{T}/events/vp/enable", "0")
    except OSError: pass
    for l in open(f"{T}/uprobe_events"):
        if l.startswith("p:vp/"): tw(f"{T}/uprobe_events", "-:" + l.split()[0].split(":", 1)[1] + "\n")
def register(binp, fns, mapf):
    segs = []
    for l in subprocess.run(["readelf", "-lW", binp], capture_output=True, text=True).stdout.splitlines():
        p = l.split()
        if p and p[0] == "LOAD" and "E" in " ".join(p[6:-1]): segs.append((int(p[2], 16), int(p[2], 16) + int(p[5], 16), int(p[1], 16)))
    tw(f"{T}/tracing_on", "0")   # counters in uprobe_profile stay exact without tracing
    with open(mapf, "w") as m:
        for i, l in enumerate(open(fns)):
            va, files, kind = l.split(); va = int(va, 16)
            seg = next(s for s in segs if s[0] <= va < s[1])
            tw(f"{T}/uprobe_events", f"p:vp/f{i} {binp}:{va - seg[0] + seg[2]:#x}\n")
            m.write(f"f{i}\t{va:#x}\t{files}\t{kind}\n")
    tw(f"{T}/events/vp/enable", "1")
def snap(out): open(out, "w").write(open(f"{T}/uprobe_profile").read())
def load(p): return {x[1].split("/")[-1]: int(x[-1]) for x in (l.split() for l in open(p)) if len(x) >= 3 and x[-1].isdigit()}
def report(mapf, *phases):   # phases: name before after ...
    rows = [l.rstrip("\n").split("\t") for l in open(mapf)]
    names = phases[0::3]
    hits = {nm: (lambda a, b: {e: b.get(e, 0) - a.get(e, 0) for e in b})(load(a), load(b)) for nm, a, b in zip(phases[0::3], phases[1::3], phases[2::3])}
    per = defaultdict(lambda: defaultdict(lambda: [0, 0]))   # file -> phase -> [own-function hits, shared-function hits]
    nfn = defaultdict(lambda: [0, 0])
    for ev, va, files, kind in rows:
        fs = files.split(","); k = 0 if kind == "own" else 1
        for f in fs:
            nfn[f][k] += 1
            for nm in names: per[f][nm][k] += hits[nm].get(ev, 0)
    hdr = "file\tfunctions(own/shared)\t" + "\t".join(f"{nm}: own / shared" for nm in names)
    print(hdr)
    for f in sorted(per):
        print(f"{f}\t{nfn[f][0]}/{nfn[f][1]}\t" + "\t".join(f"{per[f][nm][0]} / {per[f][nm][1]}" for nm in names))
    print("\nown = every panic site in the function comes from this one cited file; shared = the function also")
    print("contains panic sites from other source files (inlined or generic code), so it cannot be attributed to the cited file alone.")
c = sys.argv[1]
if c == "clear": clear()
elif c == "register": register(*sys.argv[2:5])
elif c == "snap": snap(sys.argv[2])
elif c == "report": report(*sys.argv[2:])
PY
P="sudo python3 -I $W/probes.py"
sudo ./stop; $P clear; $P register "$BIN" "$W/functions.tsv" "$W/map.tsv"
echo "probes: $(wc -l < "$W/map.tsv")"

restart() { sudo ./stop; [ "${1:-}" = cold ] && { sync; echo 3 | sudo tee /proc/sys/vm/drop_caches >/dev/null; }
            ./start; for i in $(seq 1 300); do ./check >/dev/null 2>&1 && break; sleep 1; done; }
q() { sed -n "$(( $1 + 1 ))p" queries.sql; }

echo "== 6. the benchmark as submitted: 43 queries x 3, cold restart before each query's first run"
restart; $P snap "$W/s0"
for n in $(seq 0 42); do restart cold; for r in 1 2 3; do q $n | ./query >/dev/null 2>&1 || echo "Q$n run $r: ERROR"; done; done
$P snap "$W/s1"

echo "== 7. CONTROL: force the hand-written path (the probes must fire here)"
$PSQL -tAc "ALTER DATABASE test SET pgrust.gravity_admit = 'shape'" >/dev/null
# the irprobe kernels also need PGRUST_GRAVITY_IRPROBE=tpde in the server's environment; ./start passes no
# extra env (sudo -u resets it), so the control server is started by ./start with that one variable added
START_CTL=$(sed 's/nohup env /nohup env PGRUST_GRAVITY_IRPROBE=tpde /' start)
grep -q 'PGRUST_GRAVITY_IRPROBE=tpde' <<<"$START_CTL" || { echo "control: cannot add PGRUST_GRAVITY_IRPROBE to ./start"; exit 1; }
sudo ./stop; bash -c "$START_CTL" ./start; for i in $(seq 1 300); do ./check >/dev/null 2>&1 && break; sleep 1; done
$P snap "$W/s2"
for n in $(seq 0 42); do for r in 1 2; do q $n | ./query >/dev/null 2>&1 || true; done; done
$P snap "$W/s3"
$PSQL -tAc "ALTER DATABASE test RESET pgrust.gravity_admit" >/dev/null

echo "== 8. results"
sudo ./stop; $P report "$W/map.tsv" benchmark "$W/s0" "$W/s1" control "$W/s2" "$W/s3" | tee "$W/results.tsv"
$P clear
echo "results: $W/results.tsv   located functions: $W/functions.tsv"