Skip to content

Commit 1e17c26

Browse files
committed
refactor: batched level-merge for append_leaves_to_snapshot
Replaces the per-leaf frontier walk (MaxLeaves x TreeHeight poseidon hashes) with a level-by-level batched merge: the batch is prepended with the pending left sibling when it starts as a right child, dangling odd nodes become the new frontier entries, and remaining nodes are paired into the next level with lane bounds halving per level. Total cost drops to ~MaxLeaves + 3 x TreeHeight hashes (1,166 for 1024 leaves in the height-36 tree vs ~36,900 before). Same signature, semantics, and error messages. Adds an exhaustive small-tree sweep over every (start, num) combination and a fail-closed test for appending to a completely full tree.
1 parent 8186172 commit 1e17c26

1 file changed

Lines changed: 130 additions & 27 deletions

File tree

noir-projects/noir-protocol-circuits/crates/types/src/merkle_tree/append_only_tree.nr

Lines changed: 130 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -92,15 +92,19 @@ pub fn insert_subtree_root_to_snapshot_with_hasher<let TreeHeight: u32, let Subt
9292
/// Appends up to `MaxLeaves` leaves to a poseidon append-only tree at its current (arbitrary) next-available index.
9393
///
9494
/// Unlike `insert_subtree_root_to_snapshot`, the insertion index does not need to be aligned to a power-of-two
95-
/// boundary: leaves are absorbed one at a time into the tree's frontier (the incremental-merkle-tree "filled
96-
/// subtrees"), so a compact bundle can start wherever the previous append left off. Only `leaves[0..num_leaves]` are
97-
/// inserted; the remaining lanes are no-ops and must be zero.
95+
/// boundary: the whole batch is merged into the tree's frontier (the incremental-merkle-tree "filled subtrees") level
96+
/// by level, so a compact bundle can start wherever the previous append left off. Only `leaves[0..num_leaves]` are
97+
/// inserted; the remaining lanes are no-ops and must be zero. The pair-hash lanes halve at each level as the batch
98+
/// rises, so the total cost is ~`MaxLeaves + 3 * TreeHeight` hashes rather than `MaxLeaves * TreeHeight`.
9899
///
99100
/// `frontier_hint` is the frontier at `snapshot.next_available_leaf_index`, i.e. `frontier_hint[level]` is the pending
100101
/// left-sibling node at each level. It is validated against `snapshot.root` before use, pinning it to the snapshot the
101-
/// same way sibling-path hints are pinned elsewhere. Only the hint entries at levels where the start index is a right
102-
/// child are read during the append, and those are exactly the entries the root recomputation pins, so an incorrect
103-
/// hint cannot produce a wrong result: it either fails validation or is overwritten before it is read.
102+
/// same way sibling-path hints are pinned elsewhere. The merge only ever reads hint entries at levels where the start
103+
/// index is a right child, which are exactly the entries the validation pins, so an incorrect hint cannot produce a
104+
/// wrong result: it either fails validation or is never read.
105+
///
106+
/// A completely full tree (`next_available_leaf_index == 2^TreeHeight`) has no representable frontier, so even a
107+
/// `num_leaves = 0` call against one fails hint validation. This fails closed and is accepted behavior.
104108
pub fn append_leaves_to_snapshot<let TreeHeight: u32, let MaxLeaves: u32>(
105109
snapshot: AppendOnlyTreeSnapshot,
106110
leaves: [Field; MaxLeaves],
@@ -127,11 +131,11 @@ pub fn append_leaves_to_snapshot<let TreeHeight: u32, let MaxLeaves: u32>(
127131
let start_index = snapshot.next_available_leaf_index as u64;
128132
let end_index = start_index + num_leaves as u64;
129133
// Guard against wrapping past the tree capacity, which would otherwise silently produce a wrong root.
134+
let mut tree_is_filled_exactly = false;
130135
if TreeHeight < 64 {
131-
assert(
132-
end_index <= (1 as u64 << TreeHeight as u64),
133-
"append_leaves_to_snapshot exceeds the tree capacity",
134-
);
136+
let capacity = 1 as u64 << TreeHeight as u64;
137+
assert(end_index <= capacity, "append_leaves_to_snapshot exceeds the tree capacity");
138+
tree_is_filled_exactly = end_index == capacity;
135139
}
136140

137141
// Validate the hint pins to the snapshot: recompute the root of a tree holding `start_index` leaves from the hint's
@@ -142,28 +146,59 @@ pub fn append_leaves_to_snapshot<let TreeHeight: u32, let MaxLeaves: u32>(
142146
"Frontier hint does not match the snapshot root",
143147
);
144148

145-
// Absorb each real leaf into the frontier, tracking the running root; padding lanes are skipped.
149+
// Merge the batch into the frontier one level at a time. At each level, `working` holds the batch's nodes,
150+
// starting at absolute position `s` with `c` real entries (lanes past `c` hold garbage that is never read):
151+
// - If `s` is odd the batch starts as a right child, so its pending left sibling — the validated hint entry — is
152+
// prepended via a mux shift (no hashes) to make the batch start even.
153+
// - If the count is then odd, the dangling last node sits at an even position: it is the level's new pending left
154+
// child, so it becomes the frontier entry and drops out of pairing. On an even count the old entry is left in
155+
// place; at the new tree size it is unused garbage, just like unread hint lanes.
156+
// - The remaining nodes are hashed pairwise into the next level.
146157
let mut frontier = frontier_hint;
147-
let mut index = start_index;
148-
let mut root = snapshot.root;
158+
let mut working = [0; MaxLeaves + 2];
149159
for i in 0..MaxLeaves {
150-
if i < num_leaves {
151-
let mut node = leaves[i];
152-
let mut idx = index;
153-
for level in 0..TreeHeight {
154-
if idx % 2 == 0 {
155-
frontier[level] = node;
156-
node = merkle_hash(node, zero_subtree_roots[level]);
157-
} else {
158-
node = merkle_hash(frontier[level], node);
159-
}
160-
idx = idx / 2;
161-
}
162-
root = node;
163-
index += 1;
160+
working[i] = leaves[i];
161+
}
162+
let mut s = start_index;
163+
let mut c = num_leaves;
164+
for level in 0..TreeHeight {
165+
let prepend = (s % 2 == 1) & (c != 0);
166+
let mut shifted = [0; MaxLeaves + 2];
167+
shifted[0] = if prepend { frontier[level] } else { working[0] };
168+
// The loop bounds below only need to upper-bound the batch size at this level; the shift amounts are clamped
169+
// below 32 because a u32 shift by 32+ does not const-fold, and `MaxLeaves >> 31` is already 0.
170+
for j in 1..(MaxLeaves >> std::cmp::min(level, 31)) + 2 {
171+
shifted[j] = if prepend { working[j - 1] } else { working[j] };
164172
}
173+
let s2 = s - prepend as u64;
174+
let c2 = c + prepend as u32;
175+
176+
// `c2 - 1` is only meaningful when `c2` is odd (hence >= 1); the subtraction is kept unconditionally safe.
177+
let dangling_index = c2 - (c2 != 0) as u32;
178+
frontier[level] = if c2 % 2 == 1 {
179+
shifted[dangling_index]
180+
} else {
181+
frontier[level]
182+
};
183+
184+
for i in 0..(MaxLeaves >> std::cmp::min(level + 1, 31)) + 1 {
185+
working[i] = merkle_hash(shifted[2 * i], shifted[2 * i + 1]);
186+
}
187+
s = s2 / 2;
188+
c = c2 / 2;
165189
}
166190

191+
// The frontier is now the end-state frontier, so the final root falls out of the same recomputation used to
192+
// validate the hint. The exception is an append that fills the tree exactly: a full tree has no representable
193+
// frontier, but in that case the batch was merged all the way up and `working[0]` is the root itself.
194+
let root = if num_leaves == 0 {
195+
snapshot.root
196+
} else if tree_is_filled_exactly {
197+
working[0]
198+
} else {
199+
recompute_root_from_frontier(frontier, zero_subtree_roots, end_index)
200+
};
201+
167202
AppendOnlyTreeSnapshot {
168203
root,
169204
next_available_leaf_index: snapshot.next_available_leaf_index + num_leaves as Field,
@@ -524,4 +559,72 @@ mod append_leaves_tests {
524559
let leaves = [42, 0, 99, 0, 0, 0, 0, 0];
525560
let _ = append_leaves_to_snapshot::<5, 8>(snapshot, leaves, 1, [0; 5]);
526561
}
562+
563+
/// Extracts the frontier at `index` from a height-4 reference tree: the node at each level's position
564+
/// `(index >> level) - 1` when that position is odd, i.e. the pending left sibling.
565+
unconstrained fn frontier_of_height_4_tree(tree: MerkleTree<16>, index: u32) -> [Field; 4] {
566+
// Level offsets into `MerkleTree.nodes` for N = 16: level 1 starts at 0, level 2 at 8, level 3 at 12.
567+
let offsets = [0, 8, 12];
568+
let mut frontier = [0; 4];
569+
if index % 2 == 1 {
570+
frontier[0] = tree.leaves[index - 1];
571+
}
572+
for level in 1..4 {
573+
let p = index >> level;
574+
if p % 2 == 1 {
575+
frontier[level] = tree.nodes[offsets[level - 1] + p - 1];
576+
}
577+
}
578+
frontier
579+
}
580+
581+
#[test]
582+
unconstrained fn append_exhaustive_small_tree_sweep() {
583+
// Height-4 tree (capacity 16) with MaxLeaves = 6: every (start, num) with start + num <= capacity, compared
584+
// against a reference tree built over the concatenated leaves. Exercises every parity combination of the
585+
// prepend/dangling logic, including appends that fill the tree exactly.
586+
for start in 0..16 {
587+
for num in 0..7 {
588+
if start + num <= 16 {
589+
let mut pre_leaves = [0; 16];
590+
for i in 0..start {
591+
pre_leaves[i] = (100 + i) as Field;
592+
}
593+
let pre_tree = MerkleTree::new(pre_leaves);
594+
let snapshot = AppendOnlyTreeSnapshot {
595+
root: pre_tree.get_root(),
596+
next_available_leaf_index: start as Field,
597+
};
598+
let frontier = frontier_of_height_4_tree(pre_tree, start);
599+
600+
let mut batch = [0; 6];
601+
let mut all_leaves = pre_leaves;
602+
for i in 0..num {
603+
let value = (200 + start + i) as Field;
604+
batch[i] = value;
605+
all_leaves[start + i] = value;
606+
}
607+
608+
let result = append_leaves_to_snapshot::<4, 6>(snapshot, batch, num, frontier);
609+
610+
let expected_tree = MerkleTree::new(all_leaves);
611+
assert_eq(result.root, expected_tree.get_root());
612+
assert_eq(result.next_available_leaf_index, (start + num) as Field);
613+
}
614+
}
615+
}
616+
}
617+
618+
#[test(should_fail_with = "Frontier hint does not match the snapshot root")]
619+
fn append_to_completely_full_tree_fails() {
620+
// A full tree has no representable frontier; even a num_leaves = 0 call fails hint validation (fail-closed).
621+
let mut leaves = [0; 16];
622+
for i in 0..16 {
623+
leaves[i] = (100 + i) as Field;
624+
}
625+
let tree = MerkleTree::new(leaves);
626+
let snapshot =
627+
AppendOnlyTreeSnapshot { root: tree.get_root(), next_available_leaf_index: 16 };
628+
let _ = append_leaves_to_snapshot::<4, 6>(snapshot, [0; 6], 0, [0; 4]);
629+
}
527630
}

0 commit comments

Comments
 (0)