f - in-place merkle root computation

Author: jkczyz

lightning/src/offers/merkle.rs

@@ -26,27 +26,29 @@ pub(super) fn tagged_root_hash(tag: sha256::Hash, data: &[u8]) -> sha256::Hash {
 	let leaf_tag = sha256::Hash::hash("LnLeaf".as_bytes());
 	let branch_tag = sha256::Hash::hash("LnBranch".as_bytes());
 
-	let mut merkle_leaves = Vec::new();
+	let mut leaves = Vec::new();
 	for record in TlvStream::new(&data[..]) {
 		if !SIGNATURE_TYPES.contains(&record.r#type.0) {
-			merkle_leaves.push(tagged_hash(leaf_tag, &record));
-			merkle_leaves.push(tagged_hash(nonce_tag, &record));
+			leaves.push(tagged_hash(leaf_tag, &record));
+			leaves.push(tagged_hash(nonce_tag, &record));
 		}
 	}
 
-	while merkle_leaves.len() != 1 {
-		let mut parents = Vec::with_capacity(merkle_leaves.len() / 2 + merkle_leaves.len() % 2);
-		let mut iter = merkle_leaves.chunks_exact(2);
-		for leaves in iter.by_ref() {
-			parents.push(tagged_branch_hash(branch_tag, leaves[0], leaves[1]));
+	let num_leaves = leaves.len();
+	for level in 0.. {
+		let step = 2 << level;
+		let offset = step / 2;
+		if offset >= num_leaves {
+			break;
 		}
-		for leaf in iter.remainder() {
-			parents.push(*leaf);
+
+		for (i, j) in (0..num_leaves).step_by(step).zip((offset..num_leaves).step_by(step)) {
+			leaves[i] = tagged_branch_hash(branch_tag, leaves[i], leaves[j]);
 		}
-		core::mem::swap(&mut parents, &mut merkle_leaves);
 	}
 
-	tagged_hash(tag, merkle_leaves.first().unwrap())
+	// TODO: Can we ever have zero leaves?
+	tagged_hash(tag, leaves.first().unwrap())
 }
 
 fn tagged_branch_hash(tag: sha256::Hash, leaf1: sha256::Hash, leaf2: sha256::Hash) -> sha256::Hash {