// atproto MST — WRITE path (slice 2).
//
// A faithful port of millipds's `atmst` (mst/node.py + node_wrangler.py) — the exact
// code that produced our golden CAR. Byte-exactness is verified two ways: rebuilding
// the golden repo's MST, and matching atmst's own root CIDs for a multi-height tree
// (build + delete oracles). See test/mst-write.test.js.
//
// Node shape (interleaved): subtrees[0], (keys[0],vals[0]), subtrees[1], … ,subtrees[n].
// key layer = leading-zero BITS of sha256(key), // 2 (fanout 4).
import * as dagCbor from '@ipld/dag-cbor'
import { sha256 as nobleSha256 } from '@noble/hashes/sha2'
import { cidForBytesSync } from './cid.js'
const cidEq = (a, b) => (a === null || b === null) ? a === b : a.equals(b)
const replaceAt = (arr, i, v) => { const c = arr.slice(); c[i] = v; return c }
const insertAt = (arr, i, v) => arr.slice(0, i).concat([v], arr.slice(i))
const removeAt = (arr, i) => arr.slice(0, i).concat(arr.slice(i + 1))
export function keyHeight (key) {
const h = nobleSha256(Buffer.from(key, 'utf8')) // 32 bytes
let zeros = 0
for (const byte of h) {
if (byte === 0) { zeros += 8; continue }
zeros += Math.clz32(byte) - 24 // leading zero bits within this byte
break
}
return zeros >> 1 // // 2
}
export class MSTNode {
constructor (keys, vals, subtrees) {
this.keys = keys // [string]
this.vals = vals // [CID]
this.subtrees = subtrees // [CID|null], length keys.length + 1
}
static empty () { return new MSTNode([], [], [null]) }
isEmpty () { return this.keys.length === 0 && this.subtrees.length === 1 && this.subtrees[0] === null }
serialise () {
const e = []
let prevKey = Buffer.alloc(0)
for (let i = 0; i < this.keys.length; i++) {
const keyBytes = Buffer.from(this.keys[i], 'utf8')
const max = Math.min(prevKey.length, keyBytes.length)
let p = 0
while (p < max && prevKey[p] === keyBytes[p]) p++
e.push({ k: keyBytes.subarray(p), p, t: this.subtrees[i + 1], v: this.vals[i] })
prevKey = keyBytes
}
return dagCbor.encode({ e, l: this.subtrees[0] }) // dag-cbor canonicalises map key order
}
cid () { return cidForBytesSync(this.serialise()) }
toOptional () { return this.isEmpty() ? null : this.cid() }
gteIndex (key) { let i = 0; while (i < this.keys.length && key > this.keys[i]) i++; return i }
height () {
if (this.keys.length) return keyHeight(this.keys[0])
if (this.subtrees[0] === null) return 0
return null // indeterminate (empty intermediate node)
}
static deserialise (bytes) {
const cbor = dagCbor.decode(bytes)
const subtrees = [cbor.l]
const keys = []
const vals = []
let prevKey = Buffer.alloc(0)
for (const e of cbor.e) {
const suffix = Buffer.from(e.k)
const thisKey = Buffer.concat([prevKey.subarray(0, e.p), suffix])
keys.push(thisKey.toString('utf8'))
vals.push(e.v)
subtrees.push(e.t)
prevKey = thisKey
}
return new MSTNode(keys, vals, subtrees)
}
}
// In-memory node store — maps CID string -> { cid, bytes }. get(null) returns the
// (stored) empty node, matching atmst NodeStore.get_node(None).
export class NodeStore {
constructor () { this.blocks = new Map() }
storedNode (node) {
const bytes = node.serialise()
const cid = cidForBytesSync(bytes)
this.blocks.set(cid.toString(), { cid, bytes })
return node
}
getNode (cid) {
if (cid === null || cid === undefined) return this.storedNode(MSTNode.empty())
const b = this.blocks.get(cid.toString())
if (!b) throw new Error(`missing MST node ${cid}`)
return MSTNode.deserialise(b.bytes)
}
}
export class NodeWrangler {
constructor (ns) { this.ns = ns }
putRecord (rootCid, key, val) {
const root = this.ns.getNode(rootCid)
if (root.isEmpty()) return this._putHere(root, key, val).cid()
return this._putRecursive(root, key, val, keyHeight(key), root.height()).cid()
}
delRecord (rootCid, key) {
const root = this.ns.getNode(rootCid)
const afterDel = this._deleteRecursive(root, key, keyHeight(key), root.height())
return this.ns.getNode(this._squashTop(afterDel)).cid()
}
_putHere (node, key, val) {
const i = node.gteIndex(key)
if (i < node.keys.length && node.keys[i] === key) {
if (cidEq(node.vals[i], val)) return node
return this.ns.storedNode(new MSTNode(node.keys.slice(), replaceAt(node.vals, i, val), node.subtrees.slice()))
}
const [lsub, rsub] = this._splitOnKey(node.subtrees[i], key)
return this.ns.storedNode(new MSTNode(
insertAt(node.keys, i, key),
insertAt(node.vals, i, val),
node.subtrees.slice(0, i).concat([lsub, rsub], node.subtrees.slice(i + 1))
))
}
_putRecursive (node, key, val, keyH, treeH) {
if (keyH > treeH) { // grow the tree upward
return this.ns.storedNode(this._putRecursive(
this.ns.storedNode(new MSTNode([], [], [node.cid()])),
key, val, keyH, treeH + 1
))
}
if (keyH < treeH) { // descend
const i = node.gteIndex(key)
return this.ns.storedNode(new MSTNode(
node.keys.slice(), node.vals.slice(),
replaceAt(node.subtrees, i, this._putRecursive(this.ns.getNode(node.subtrees[i]), key, val, keyH, treeH - 1).cid())
))
}
return this._putHere(node, key, val) // keyH === treeH
}
_splitOnKey (nodeCid, key) {
if (nodeCid === null || nodeCid === undefined) return [null, null]
const node = this.ns.getNode(nodeCid)
const i = node.gteIndex(key)
const [lsub, rsub] = this._splitOnKey(node.subtrees[i], key)
const left = this.ns.storedNode(new MSTNode(node.keys.slice(0, i), node.vals.slice(0, i), node.subtrees.slice(0, i).concat([lsub]))).toOptional()
const right = this.ns.storedNode(new MSTNode(node.keys.slice(i), node.vals.slice(i), [rsub].concat(node.subtrees.slice(i + 1)))).toOptional()
return [left, right]
}
_squashTop (nodeCid) {
const node = this.ns.getNode(nodeCid)
if (node.keys.length) return nodeCid
if (node.subtrees[0] === null) return nodeCid
return this._squashTop(node.subtrees[0])
}
_deleteRecursive (node, key, keyH, treeH) {
if (keyH > treeH) return node.toOptional()
let i = node.gteIndex(key)
if (keyH < treeH) {
if (node.subtrees[i] === null) return node.toOptional()
return this.ns.storedNode(new MSTNode(
node.keys.slice(), node.vals.slice(),
replaceAt(node.subtrees, i, this._deleteRecursive(this.ns.getNode(node.subtrees[i]), key, keyH, treeH - 1))
)).toOptional()
}
i = node.gteIndex(key)
if (i === node.keys.length || node.keys[i] !== key) return node.toOptional()
return this.ns.storedNode(new MSTNode(
removeAt(node.keys, i), removeAt(node.vals, i),
node.subtrees.slice(0, i).concat([this._merge(node.subtrees[i], node.subtrees[i + 1])], node.subtrees.slice(i + 2))
)).toOptional()
}
_merge (leftCid, rightCid) {
if (leftCid === null) return rightCid
if (rightCid === null) return leftCid
const left = this.ns.getNode(leftCid)
const right = this.ns.getNode(rightCid)
return this.ns.storedNode(new MSTNode(
left.keys.concat(right.keys), left.vals.concat(right.vals),
left.subtrees.slice(0, -1).concat([this._merge(left.subtrees[left.subtrees.length - 1], right.subtrees[0])], right.subtrees.slice(1))
)).toOptional()
}
}
// Build an MST from records [{ key, value: CID }] (any order). Returns
// { root: CID, ns: NodeStore }. ns.blocks holds every node created; use
// reachableBlocks(ns, root) for just the live set (CAR output).
export function buildMst (records) {
const ns = new NodeStore()
const w = new NodeWrangler(ns)
let root = ns.storedNode(MSTNode.empty()).cid()
for (const { key, value } of records) root = w.putRecord(root, key, value)
return { root, ns, wrangler: w }
}
// BFS the live MST nodes reachable from root (for CAR block sets).
export function reachableBlocks (ns, rootCid) {
const out = new Map()
const stack = [rootCid]
while (stack.length) {
const cid = stack.pop()
if (cid === null || cid === undefined) continue
const key = cid.toString()
if (out.has(key)) continue
const b = ns.blocks.get(key)
if (!b) throw new Error(`missing node ${key}`)
out.set(key, b)
const node = MSTNode.deserialise(b.bytes)
for (const t of node.subtrees) if (t) stack.push(t)
}
return out
}
// ---- Firehose MST-inversion proof (slice 3b) ----
//
// A #commit's `blocks` CAR must be a MINIMAL, self-contained proof of the
// old→new transition — not the whole tree (that would trip a relay's `tooBig`).
// This ports millipds `repo_ops.py:apply_writes` step 2b + `atmst` proof.py:
// blocks = created | deletion_proof_cids | inductive_proof_nodes (+ new records)
// where `inductive_proof_nodes` = every MST node READ while replaying the ops
// backward from the new root, which must reproduce the old (prevData) root.
// The prior record CIDs travel in each op's `prev` (like millipds), so a cold
// relay can invert with no old-tree state of its own.
// A NodeStore overlay: writes land in an in-memory upper layer (NOT logged);
// reads that fall through to `lower` are recorded in `.gets`. Mirrors millipds's
// OverlayBlockStore(MemoryBlockStore, LoggingBlockStoreWrapper(bs)).
export class OverlayNodeStore {
constructor (lower) { this.lower = lower; this.upper = new Map(); this.gets = new Set() }
storedNode (node) {
const bytes = node.serialise()
const cid = cidForBytesSync(bytes)
this.upper.set(cid.toString(), { cid, bytes })
return node
}
getNode (cid) {
if (cid === null || cid === undefined) return MSTNode.empty()
const k = cid.toString()
const u = this.upper.get(k)
if (u) return MSTNode.deserialise(u.bytes)
this.gets.add(k) // a read of a pre-existing node — part of the inductive proof
return this.lower.getNode(cid)
}
}
// Follow the search path for `key` from `rootCid`, collecting every node visited.
// Returns { value: CID|null, proof: Set<cidStr> }. value===null ⇒ exclusion proof;
// value!==null ⇒ inclusion proof. Port of atmst find_rpath_and_build_proof.
export function findRpathProof (ns, rootCid, key) {
const proof = new Set()
let cid = rootCid
while (cid !== null && cid !== undefined) {
const node = ns.getNode(cid)
proof.add(cid.toString())
const i = node.gteIndex(key)
if (i < node.keys.length && node.keys[i] === key) return { value: node.vals[i], proof }
cid = node.subtrees[i] // descend the subtree left of the first key >= target
}
return { value: null, proof }
}
// Build a NodeStore holding exactly the given blocks (cid-string -> { cid, bytes }).
export function nodeStoreFromBlocks (blocks) {
const ns = new NodeStore()
for (const [k, v] of blocks) ns.blocks.set(k, v)
return ns
}
// Compute the minimal firehose block set for a single commit's ops.
// oldRecords / newRecords : [{ key, value: CID }] — the record sets BEFORE/AFTER.
// ops : [{ action:'create'|'update'|'delete', path, cid: CID|null, prev: CID|null }]
// Returns { blocks: Map<cidStr,{cid,bytes}>, dataCid, oldRoot } — MST nodes only
// (record + commit blocks are added by the caller). Throws if the backward replay
// does not reproduce the old root (i.e. the proof would be invalid).
export function firehoseProofBlocks ({ oldRecords, newRecords, ops }) {
const oldB = buildMst(oldRecords)
const newB = buildMst(newRecords)
const oldRoot = oldB.root
const newRoot = newB.root
const oldSet = reachableBlocks(oldB.ns, oldRoot)
const newSet = reachableBlocks(newB.ns, newRoot)
// combined store holds every node of BOTH trees (proof reads span both)
const combined = new NodeStore()
for (const [k, v] of oldSet) combined.blocks.set(k, v)
for (const [k, v] of newSet) combined.blocks.set(k, v)
const proofCids = new Set()
// created: new-tree nodes not present in the old tree
for (const k of newSet.keys()) if (!oldSet.has(k)) proofCids.add(k)
// deletion proofs: for each delete, the exclusion-proof path in the NEW tree
for (const op of ops) {
if (op.action !== 'delete') continue
const { value, proof } = findRpathProof(combined, newRoot, op.path)
if (value !== null) throw new Error(`delete proof: ${op.path} still present in new tree`)
for (const c of proof) proofCids.add(c)
}
// inductive proof: replay ops backward from newRoot; every node read is required.
const overlay = new OverlayNodeStore(combined)
const w = new NodeWrangler(overlay)
let inv = newRoot
for (let i = ops.length - 1; i >= 0; i--) {
const op = ops[i]
if (op.action === 'create') {
inv = w.delRecord(inv, op.path)
} else { // update | delete → restore the prior value
if (op.prev === null || op.prev === undefined) throw new Error(`inversion: ${op.action} ${op.path} missing prev`)
inv = w.putRecord(inv, op.path, op.prev)
}
}
if (!inv.equals(oldRoot)) throw new Error('inversion proof failed: backward replay did not reproduce prevData root')
for (const k of overlay.gets) proofCids.add(k)
const blocks = new Map()
for (const k of proofCids) blocks.set(k, combined.blocks.get(k))
return { blocks, dataCid: newRoot, oldRoot }
}