zklicensing 0.1.0-beta.2 → 0.1.0-beta.21

package/src/permitTree.ts

@@ -0,0 +1,232 @@
+// Copyright (c) 2025-2026 zkLicensing project developers
+// SPDX-License-Identifier: BUSL-1.1
+
+/**
+ * permitTree.ts — `zklicensing/permits/tree`
+ *
+ * A compact copy of a registry's permits tree, for holders to build their
+ * Merkle witness without the keeper learning which permit is theirs, and
+ * without rebuilding the 256-level tree (one 256-hash path per permit,
+ * about 0.1 s each).
+ *
+ * The tree is sparse: permits sit at random-looking positions, so below the
+ * top few levels nearly every permit is alone in its subtree and everything
+ * beside it is empty — known constants. The only non-constant siblings any
+ * witness needs are the two children of each node where two permits' paths
+ * split: with the permits sorted by position, one such node per adjacent
+ * pair, N − 1 in all. The compact tree is every permit's record plus those
+ * 2(N − 1) node values. Everyone downloads the same thing; a wrong copy only
+ * makes the holder's proof fail against the chain's root, never pass.
+ *
+ * Updates: `applyTreeDiff` takes the records that changed since a root the
+ * holder has and recomputes only the nodes above them.
+ */
+
+import { Field, MerkleMap, Poseidon } from 'o1js';
+import { permitLeafHash } from './permitMessages.js';
+import { recordToLeaf, type PermitRecord } from './permitRegistryRebuild.js';
+import { CompactMerkleMap, encodeFields, decodeFields } from './compactMerkleMap.js';
+
+const HEIGHT = 256;
+
+export type CompactTree = {
+  registry: string;
+  root: string;
+  slot: number;
+  /** Every permit, sorted by tree position. */
+  permits: PermitRecord[];
+  /** The 2(N − 1) split-children values, 32 bytes each, big-endian, base64. */
+  nodes: string;
+};
+
+export type TreeDiff = { registry: string; root: string; slot: number; since: string; changed: PermitRecord[] };
+
+// ── field and position helpers ───────────────────────────────────────────────
+
+let zeroCache: bigint[] | null = null;
+/** zeros[l]: the value of an empty subtree whose top is at level l. */
+function zeros(): bigint[] {
+  if (!zeroCache) {
+    const z = [0n];
+    for (let i = 1; i < HEIGHT; i++) z.push(Poseidon.hash([Field(z[i - 1]), Field(z[i - 1])]).toBigInt());
+    zeroCache = z;
+  }
+  return zeroCache;
+}
+const H = (a: bigint, b: bigint) => Poseidon.hash([Field(a), Field(b)]).toBigInt();
+
+const keyMap = new MerkleMap();
+/** A subject's leaf position (as MerkleMap places it). */
+export function subjectIndex(subject: string): bigint {
+  return (keyMap as any)._keyToIndex(Field(subject)) as bigint;
+}
+
+function bitLength(x: bigint): number { let n = 0; while (x > 0n) { x >>= 1n; n++; } return n; }
+
+/** The split-children positions, [level, index] pairs, in canonical order, for sorted leaf positions. */
+function splitChildren(sorted: bigint[]): Array<[number, bigint]> {
+  const out: Array<[number, bigint]> = [];
+  for (let i = 0; i + 1 < sorted.length; i++) {
+    const a = sorted[i], b = sorted[i + 1];
+    const l = bitLength(a ^ b) - 1; // the children's level; their parent at l + 1 is where the paths split
+    out.push([l, a >> BigInt(l)], [l, b >> BigInt(l)]);
+  }
+  return out;
+}
+
+function sortPermits(permits: PermitRecord[]): { permits: PermitRecord[]; idx: bigint[] } {
+  const withIdx = permits.map(p => ({ p, i: subjectIndex(p.subject) }));
+  withIdx.sort((x, y) => (x.i < y.i ? -1 : x.i > y.i ? 1 : 0));
+  return { permits: withIdx.map(x => x.p), idx: withIdx.map(x => x.i) };
+}
+
+const nodeKey = (l: number, j: bigint) => `${l}:${j}`;
+
+// ── keeper side ──────────────────────────────────────────────────────────────
+
+/**
+ * The compact tree from a full permits tree (the keeper's MerkleMap). Reads
+ * node values; no hashing.
+ */
+export function compactTreeFrom(args: { registry: string; slot: number; tree: MerkleMap | CompactMerkleMap; permits: Iterable<PermitRecord> }): CompactTree {
+  const { permits, idx } = sortPermits([...args.permits]);
+  let values: bigint[];
+  if (args.tree instanceof CompactMerkleMap) {
+    values = args.tree.splitValues();
+  } else {
+    const full = (args.tree as any).tree as { getNode(level: number, index: bigint): Field };
+    values = splitChildren(idx).map(([l, j]) => full.getNode(l, j).toBigInt());
+  }
+  if (values.length !== Math.max(0, permits.length - 1) * 2) throw new Error('compact tree: the tree and the permits list disagree');
+  return { registry: args.registry, root: args.tree.getRoot().toString(), slot: args.slot, permits, nodes: encodeFields(values) };
+}
+
+// ── holder side ──────────────────────────────────────────────────────────────
+
+type Indexed = { ct: CompactTree; idx: bigint[]; values: Map<string, bigint> };
+const indexed = new WeakMap<CompactTree, Indexed>();
+
+function index(ct: CompactTree): Indexed {
+  let x = indexed.get(ct);
+  if (x) return x;
+  const idx = ct.permits.map(p => subjectIndex(p.subject));
+  for (let i = 1; i < idx.length; i++) if (idx[i] <= idx[i - 1]) throw new Error('compact tree: permits are not sorted by position');
+  const positions = splitChildren(idx);
+  const vals = ct.nodes ? decodeFields(ct.nodes) : [];
+  if (vals.length !== positions.length) throw new Error(`compact tree: expected ${positions.length} node values, got ${vals.length}`);
+  const values = new Map<string, bigint>();
+  positions.forEach(([l, j], i) => values.set(nodeKey(l, j), vals[i]));
+  x = { ct, idx, values };
+  indexed.set(ct, x);
+  return x;
+}
+
+/**
+ * The Merkle witness for `subject` from the compact tree, checked by folding
+ * it to the tree's root (one path of hashes).
+ */
+export function witnessFromCompactTree(ct: CompactTree, subject: string): { root: string; witness: { isLefts: boolean[]; siblings: string[] }; permit: PermitRecord } {
+  const { values } = index(ct);
+  const permit = ct.permits.find(p => p.subject === subject);
+  if (!permit) throw new Error('The registry has no such permit.');
+  const z = zeros();
+  const s = subjectIndex(subject);
+  const isLefts: boolean[] = [];
+  const siblings: string[] = [];
+  let h = permitLeafHash(recordToLeaf(permit)).toBigInt();
+  for (let l = 0; l < HEIGHT - 1; l++) {
+    const at = s >> BigInt(l);
+    const left = (at & 1n) === 0n;
+    const sib = values.get(nodeKey(l, at ^ 1n)) ?? z[l];
+    isLefts.push(left);
+    siblings.push(sib.toString());
+    h = left ? H(h, sib) : H(sib, h);
+  }
+  if (h.toString() !== ct.root) throw new Error('The registry copy does not fold to its root; download it again.');
+  return { root: ct.root, witness: { isLefts, siblings }, permit };
+}
+
+/**
+ * Applies the records that changed since `old.root`, recomputing only the
+ * nodes above them, and checks the result against `diff.root`.
+ */
+export function applyTreeDiff(old: CompactTree, diff: TreeDiff): CompactTree {
+  if (diff.since !== old.root) throw new Error('This update is for another version of the registry.');
+  if (diff.registry !== old.registry) throw new Error('This update is for another registry.');
+  const prev = index(old);
+  const bySubject = new Map(old.permits.map(p => [p.subject, p] as const));
+  for (const p of diff.changed) bySubject.set(p.subject, p);
+  const { permits, idx } = sortPermits([...bySubject.values()]);
+  const changedIdx = diff.changed.map(p => subjectIndex(p.subject)).sort((a, b) => (a < b ? -1 : a > b ? 1 : 0));
+  const leafHash = new Map<bigint, bigint>(permits.map((p, i) => [idx[i], permitLeafHash(recordToLeaf(p)).toBigInt()]));
+  const z = zeros();
+
+  // First position >= x in a sorted array.
+  const lower = (arr: bigint[], x: bigint) => { let lo = 0, hi = arr.length; while (lo < hi) { const m = (lo + hi) >> 1; if (arr[m] < x) lo = m + 1; else hi = m; } return lo; };
+  const inRange = (arr: bigint[], lo: bigint, hi: bigint) => { const a = lower(arr, lo); return { a, b: lower(arr, hi + 1n) }; };
+
+  const memo = new Map<string, bigint>();
+  const node = (l: number, j: bigint): bigint => {
+    const k = nodeKey(l, j);
+    const m = memo.get(k);
+    if (m !== undefined) return m;
+    const lo = j << BigInt(l), hi = ((j + 1n) << BigInt(l)) - 1n;
+    const leaves = inRange(idx, lo, hi);
+    let v: bigint;
+    if (leaves.b === leaves.a) v = z[l];
+    else {
+      const touched = inRange(changedIdx, lo, hi);
+      const kept = touched.a === touched.b ? prev.values.get(k) : undefined;
+      if (kept !== undefined) v = kept;
+      else if (l === 0) v = leafHash.get(j)!;
+      else if (leaves.b - leaves.a === 1) {
+        // One permit below: hash its path up to here, beside empty subtrees.
+        const pos = idx[leaves.a];
+        let h = leafHash.get(pos)!;
+        for (let t = 0; t < l; t++) h = ((pos >> BigInt(t)) & 1n) === 0n ? H(h, z[t]) : H(z[t], h);
+        v = h;
+      } else v = H(node(l - 1, 2n * j), node(l - 1, 2n * j + 1n));
+    }
+    memo.set(k, v);
+    return v;
+  };
+
+  const values = splitChildren(idx).map(([l, j]) => node(l, j));
+  const root = node(HEIGHT - 1, 0n);
+  if (root.toString() !== diff.root) throw new Error('The updated registry copy does not match the new root; download it again.');
+  return { registry: old.registry, root: diff.root, slot: diff.slot, permits, nodes: encodeFields(values) };
+}
+
+/** Where a holder keeps its registry copies (IndexedDB in a browser, memory or a file elsewhere). */
+export type TreeCache = {
+  get(registry: string): Promise<CompactTree | null>;
+  set(registry: string, tree: CompactTree): Promise<void>;
+};
+
+/**
+ * The registry's current compact tree: the cached copy brought up to date
+ * with a diff when the keeper has one, else a full download.
+ */
+export async function syncRegistryTree(args: {
+  registry: string;
+  fetchTree: (since?: string) => Promise<CompactTree | TreeDiff | { root: string; slot: number; unchanged: true }>;
+  cache?: TreeCache;
+  onProgress?: (what: 'cached' | 'diff' | 'full') => void;
+}): Promise<CompactTree> {
+  const cached = args.cache ? await args.cache.get(args.registry).catch(() => null) : null;
+  if (cached) {
+    try {
+      const res: any = await args.fetchTree(cached.root);
+      let next: CompactTree | null = null;
+      if (res.unchanged) { next = { ...cached, slot: res.slot }; args.onProgress?.('cached'); }
+      else if (Array.isArray(res.changed)) { next = applyTreeDiff(cached, res as TreeDiff); args.onProgress?.('diff'); }
+      else if (Array.isArray(res.permits)) { next = res as CompactTree; args.onProgress?.('full'); }
+      if (next) { await args.cache!.set(args.registry, next).catch(() => {}); return next; }
+    } catch { /* fall through to a full download */ }
+  }
+  const full = await args.fetchTree() as CompactTree;
+  if (!Array.isArray(full.permits)) throw new Error('The keeper did not send the registry.');
+  args.onProgress?.('full');
+  await args.cache?.set(args.registry, full).catch(() => {});
+  return full;
+}