binary-merkle-tree: preallocate proof vector to avoid reallocations

prdoc/pr_12288.prdoc

@@ -0,0 +1,16 @@
+# Schema: Polkadot SDK PRDoc Schema (prdoc) v1.0.0
+# See doc at https://raw.githubusercontent.com/paritytech/polkadot-sdk/master/prdoc/schema_user.json
+
+title: 'binary-merkle-tree: preallocate proof vector to avoid reallocations'
+
+doc:
+  - audience: Runtime Dev
+    description: |
+      `ProofCollection` now preallocates the proof `Vec` with a capacity of
+      `ceil(log2(number_of_leaves))`, the maximum number of nodes a single-leaf
+      Merkle proof can contain. This avoids intermediate reallocations while
+      collecting the proof in `merkle_proof` and `merkle_proof_raw`.
+
+crates:
+  - name: binary-merkle-tree
+    bump: patch

substrate/utils/binary-merkle-tree/src/lib.rs

@@ -148,15 +148,31 @@ impl<T> Visitor<T> for () {
 	fn visit(&mut self, _index: u32, _left: &Option<T>, _right: &Option<T>) {}
 }
 
+/// Maximum number of nodes a single-leaf proof can contain for a tree of `number_of_leaves`
+/// leaves, i.e. `ceil(log2(number_of_leaves))` (the height of the tree above the leaves).
+///
+/// `number_of_leaves <= 1` is guarded explicitly to avoid `0u32.ilog2()` panicking, and the
+/// `(n - 1).ilog2() + 1` form computes the ceiling without the overflow risk of
+/// `n.next_power_of_two().ilog2()`.
+fn proof_capacity(number_of_leaves: u32) -> usize {
+	if number_of_leaves <= 1 {
+		0
+	} else {
+		(number_of_leaves - 1).ilog2() as usize + 1
+	}
+}
+
 /// The struct collects a proof for single leaf.
 struct ProofCollection<T> {
 	proof: Vec<T>,
 	position: u32,
 }
 
 impl<T> ProofCollection<T> {
-	fn new(position: u32) -> Self {
-		ProofCollection { proof: Default::default(), position }
+	fn new(position: u32, number_of_leaves: u32) -> Self {
+		// Preallocate the proof to its maximum possible size to avoid intermediate
+		// reallocations while collecting it.
+		ProofCollection { proof: Vec::with_capacity(proof_capacity(number_of_leaves)), position }
 	}
 }
 
@@ -209,7 +225,7 @@ where
 	});
 
 	let number_of_leaves = iter.len() as u32;
-	let mut collect_proof = ProofCollection::new(leaf_index);
+	let mut collect_proof = ProofCollection::new(leaf_index, number_of_leaves);
 
 	let root = merkelize::<H, _, _>(iter, &mut collect_proof);
 	let leaf = leaf.expect("Requested `leaf_index` is greater than number of leaves.");
@@ -257,7 +273,7 @@ where
 	});
 
 	let number_of_leaves = iter.len() as u32;
-	let mut collect_proof = ProofCollection::new(leaf_index);
+	let mut collect_proof = ProofCollection::new(leaf_index, number_of_leaves);
 
 	let root = merkelize::<H, _, _>(iter, &mut collect_proof);
 	let leaf = leaf.expect("Requested `leaf_index` is greater than number of leaves.");
@@ -863,4 +879,44 @@ mod tests {
 			}
 		);
 	}
+
+	#[test]
+	fn proof_capacity_matches_ceil_log2() {
+		// (number_of_leaves, expected `ceil(log2(n))`).
+		let cases = [
+			(0, 0),
+			(1, 0), // both guarded to avoid `0u32.ilog2()` panicking
+			(2, 1),
+			(3, 2),
+			(4, 2),
+			(5, 3),
+			(7, 3),
+			(8, 3),
+			(9, 4),
+			(1 << 20, 20),
+			(u32::MAX, 32), // no overflow at the extreme
+		];
+		for (n, expected) in cases {
+			assert_eq!(proof_capacity(n), expected, "n={n}");
+		}
+	}
+
+	#[test]
+	fn proof_length_never_exceeds_tree_height() {
+		// The preallocation relies on a single-leaf proof never holding more than
+		// `ceil(log2(number_of_leaves))` nodes. The range straddles the power-of-two
+		// boundary (128) where last-odd-node promotion is most likely to surprise.
+		for n in 1u32..=130 {
+			let data: Vec<H256> = (0..n).map(|i| H256::repeat_byte(i as u8)).collect();
+			let max = proof_capacity(n);
+			for leaf_index in 0..n {
+				let proof = merkle_proof::<Keccak256, _, _>(data.clone(), leaf_index);
+				assert!(
+					proof.proof.len() <= max,
+					"n={n}, leaf={leaf_index}: proof len {} exceeds height {max}",
+					proof.proof.len(),
+				);
+			}
+		}
+	}
 }