mod.rs
1 // Copyright (c) 2019-2025 Alpha-Delta Network Inc. 2 // This file is part of the alphavm library. 3 4 // Licensed under the Apache License, Version 2.0 (the "License"); 5 // you may not use this file except in compliance with the License. 6 // You may obtain a copy of the License at: 7 8 // http://www.apache.org/licenses/LICENSE-2.0 9 10 // Unless required by applicable law or agreed to in writing, software 11 // distributed under the License is distributed on an "AS IS" BASIS, 12 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 // See the License for the specific language governing permissions and 14 // limitations under the License. 15 16 mod helpers; 17 pub use helpers::*; 18 19 mod path; 20 pub use path::*; 21 22 #[cfg(test)] 23 mod tests; 24 25 use alphavm_console_types::prelude::*; 26 27 use alphastd::prelude::*; 28 29 use serde::{Deserialize, Serialize}; 30 use std::collections::BTreeMap; 31 32 #[cfg(not(feature = "serial"))] 33 use rayon::prelude::*; 34 35 #[derive(Clone, Deserialize, Serialize)] 36 #[serde(bound = "E: Serialize + DeserializeOwned, LH: Serialize + DeserializeOwned, PH: Serialize + DeserializeOwned")] 37 pub struct MerkleTree<E: Environment, LH: LeafHash<Hash = PH::Hash>, PH: PathHash<Hash = Field<E>>, const DEPTH: u8> { 38 /// The leaf hasher for the Merkle tree. 39 leaf_hasher: LH, 40 /// The path hasher for the Merkle tree. 41 path_hasher: PH, 42 /// The computed root of the full Merkle tree. 43 root: PH::Hash, 44 /// The internal hashes, from root to hashed leaves, of the full Merkle tree. 45 tree: Vec<PH::Hash>, 46 /// The canonical empty hash. 47 empty_hash: Field<E>, 48 /// The number of hashed leaves in the tree. 49 number_of_leaves: usize, 50 } 51 52 impl<E: Environment, LH: LeafHash<Hash = PH::Hash>, PH: PathHash<Hash = Field<E>>, const DEPTH: u8> 53 MerkleTree<E, LH, PH, DEPTH> 54 { 55 #[inline] 56 /// Initializes a new Merkle tree with the given leaves. 57 pub fn new(leaf_hasher: &LH, path_hasher: &PH, leaves: &[LH::Leaf]) -> Result<Self> { 58 let timer = timer!("MerkleTree::new"); 59 60 // Ensure the Merkle tree depth is greater than 0. 61 ensure!(DEPTH > 0, "Merkle tree depth must be greater than 0"); 62 // Ensure the Merkle tree depth is less than or equal to 64. 63 ensure!(DEPTH <= 64u8, "Merkle tree depth must be less than or equal to 64"); 64 65 // Compute the maximum number of leaves. 66 let max_leaves = match leaves.len().checked_next_power_of_two() { 67 Some(num_leaves) => num_leaves, 68 None => bail!("Integer overflow when computing the maximum number of leaves in the Merkle tree"), 69 }; 70 71 // Compute the number of nodes. 72 let num_nodes = max_leaves - 1; 73 // Compute the tree size as the maximum number of leaves plus the number of nodes. 74 let tree_size = max_leaves + num_nodes; 75 // Compute the number of levels in the Merkle tree (i.e. log2(tree_size)). 76 let tree_depth = tree_depth::<DEPTH>(tree_size)?; 77 // Compute the number of padded levels. 78 let padding_depth = DEPTH - tree_depth; 79 80 // Compute the empty hash. 81 let empty_hash = path_hasher.hash_empty()?; 82 83 // Calculate the size of the tree which excludes leafless nodes. 84 // The minimum tree size is either a single root node or the calculated number of nodes plus 85 // the supplied leaves; if the number of leaves is odd, an empty hash is added for padding. 86 let minimum_tree_size = 87 std::cmp::max(1, num_nodes + leaves.len() + if leaves.len() > 1 { leaves.len() % 2 } else { 0 }); 88 89 // Initialize the Merkle tree. 90 let mut tree = vec![empty_hash; minimum_tree_size]; 91 92 // Compute and store each leaf hash. 93 tree[num_nodes..num_nodes + leaves.len()].copy_from_slice(&leaf_hasher.hash_leaves(leaves)?); 94 lap!(timer, "Hashed {} leaves", leaves.len()); 95 96 // Compute and store the hashes for each level, iterating from the penultimate level to the root level. 97 let mut start_index = num_nodes; 98 // Compute the start index of the current level. 99 while let Some(start) = parent(start_index) { 100 // Compute the end index of the current level. 101 let end = left_child(start); 102 // Construct the children for each node in the current level; the leaves are padded, which means 103 // that there either are 2 children, or there are none, at which point we may stop iterating. 104 let tuples = (start..end) 105 .take_while(|&i| tree.get(left_child(i)).is_some()) 106 .map(|i| (tree[left_child(i)], tree[right_child(i)])) 107 .collect::<Vec<_>>(); 108 // Compute and store the hashes for each node in the current level. 109 let num_full_nodes = tuples.len(); 110 tree[start..][..num_full_nodes].copy_from_slice(&path_hasher.hash_all_children(&tuples)?); 111 // Use the precomputed empty node hash for every empty node, if there are any. 112 if start + num_full_nodes < end { 113 let empty_node_hash = path_hasher.hash_children(&empty_hash, &empty_hash)?; 114 for node in tree.iter_mut().take(end).skip(start + num_full_nodes) { 115 *node = empty_node_hash; 116 } 117 } 118 // Update the start index for the next level. 119 start_index = start; 120 } 121 lap!(timer, "Hashed {} levels", tree_depth); 122 123 // Compute the root hash, by iterating from the root level up to `DEPTH`. 124 let mut root_hash = tree[0]; 125 for _ in 0..padding_depth { 126 // Update the root hash, by hashing the current root hash with the empty hash. 127 root_hash = path_hasher.hash_children(&root_hash, &empty_hash)?; 128 } 129 lap!(timer, "Hashed {} padding levels", padding_depth); 130 131 finish!(timer); 132 133 Ok(Self { 134 leaf_hasher: leaf_hasher.clone(), 135 path_hasher: path_hasher.clone(), 136 root: root_hash, 137 tree, 138 empty_hash, 139 number_of_leaves: leaves.len(), 140 }) 141 } 142 143 #[inline] 144 /// Returns a new Merkle tree with the given new leaves appended to it. 145 pub fn prepare_append(&self, new_leaves: &[LH::Leaf]) -> Result<Self> { 146 let timer = timer!("MerkleTree::prepare_append"); 147 148 // Compute the maximum number of leaves. 149 let max_leaves = match (self.number_of_leaves + new_leaves.len()).checked_next_power_of_two() { 150 Some(num_leaves) => num_leaves, 151 None => bail!("Integer overflow when computing the maximum number of leaves in the Merkle tree"), 152 }; 153 // Compute the number of nodes. 154 let num_nodes = max_leaves - 1; 155 // Compute the tree size as the maximum number of leaves plus the number of nodes. 156 let tree_size = num_nodes + max_leaves; 157 // Compute the number of levels in the Merkle tree (i.e. log2(tree_size)). 158 let tree_depth = tree_depth::<DEPTH>(tree_size)?; 159 // Compute the number of padded levels. 160 let padding_depth = DEPTH - tree_depth; 161 162 // Initialize the Merkle tree. 163 let mut tree = vec![self.empty_hash; num_nodes]; 164 // Extend the new Merkle tree with the existing leaf hashes. 165 tree.extend(self.leaf_hashes()?); 166 // Extend the new Merkle tree with the new leaf hashes. 167 tree.extend(&self.leaf_hasher.hash_leaves(new_leaves)?); 168 169 // Calculate the size of the tree which excludes leafless nodes. 170 let new_number_of_leaves = self.number_of_leaves + new_leaves.len(); 171 let minimum_tree_size = std::cmp::max( 172 1, 173 num_nodes + new_number_of_leaves + if new_number_of_leaves > 1 { new_number_of_leaves % 2 } else { 0 }, 174 ); 175 176 // Resize the new Merkle tree with empty hashes to pad up to `tree_size`. 177 tree.resize(minimum_tree_size, self.empty_hash); 178 lap!(timer, "Hashed {} new leaves", new_leaves.len()); 179 180 // Initialize a start index to track the starting index of the current level. 181 let start_index = num_nodes; 182 // Initialize a middle index to separate the precomputed indices from the new indices that need to be computed. 183 let middle_index = num_nodes + self.number_of_leaves; 184 // Initialize a precompute index to track the starting index of each precomputed level. 185 let start_precompute_index = match self.number_of_leaves.checked_next_power_of_two() { 186 Some(num_leaves) => num_leaves - 1, 187 None => bail!("Integer overflow when computing the Merkle tree precompute index"), 188 }; 189 // Initialize a precompute index to track the middle index of each precomputed level. 190 let middle_precompute_index = match num_nodes == start_precompute_index { 191 // If the old tree and new tree are of the same size, then we can copy over the right half of the old tree. 192 true => Some(start_precompute_index + self.number_of_leaves + new_leaves.len() + 1), 193 // Otherwise, we need to compute the right half of the new tree. 194 false => None, 195 }; 196 197 // Compute and store the hashes for each level, iterating from the penultimate level to the root level. 198 self.compute_updated_tree( 199 &mut tree, 200 start_index, 201 middle_index, 202 start_precompute_index, 203 middle_precompute_index, 204 )?; 205 206 // Compute the root hash, by iterating from the root level up to `DEPTH`. 207 let mut root_hash = tree[0]; 208 for _ in 0..padding_depth { 209 // Update the root hash, by hashing the current root hash with the empty hash. 210 root_hash = self.path_hasher.hash_children(&root_hash, &self.empty_hash)?; 211 } 212 lap!(timer, "Hashed {} padding levels", padding_depth); 213 214 finish!(timer); 215 216 Ok(Self { 217 leaf_hasher: self.leaf_hasher.clone(), 218 path_hasher: self.path_hasher.clone(), 219 root: root_hash, 220 tree, 221 empty_hash: self.empty_hash, 222 number_of_leaves: self.number_of_leaves + new_leaves.len(), 223 }) 224 } 225 226 #[inline] 227 /// Updates the Merkle tree with the given new leaves appended to it. 228 pub fn append(&mut self, new_leaves: &[LH::Leaf]) -> Result<()> { 229 let timer = timer!("MerkleTree::append"); 230 231 // Compute the updated Merkle tree with the new leaves. 232 let updated_tree = self.prepare_append(new_leaves)?; 233 // Update the tree at the very end, so the original tree is not altered in case of failure. 234 *self = updated_tree; 235 236 finish!(timer); 237 Ok(()) 238 } 239 240 #[inline] 241 /// Updates the Merkle tree at the location of the given leaf index with the new leaf. 242 pub fn update(&mut self, leaf_index: usize, new_leaf: &LH::Leaf) -> Result<()> { 243 let timer = timer!("MerkleTree::update"); 244 245 // Compute the updated Merkle tree with the new leaves. 246 let updated_tree = self.prepare_update(leaf_index, new_leaf)?; 247 // Update the tree at the very end, so the original tree is not altered in case of failure. 248 *self = updated_tree; 249 250 finish!(timer); 251 Ok(()) 252 } 253 254 #[inline] 255 /// Returns a new Merkle tree with updates at the location of the given leaf index with the new leaf. 256 pub fn prepare_update(&self, leaf_index: usize, new_leaf: &LH::Leaf) -> Result<Self> { 257 let timer = timer!("MerkleTree::prepare_update"); 258 259 // Check that the leaf index is within the bounds of the Merkle tree. 260 ensure!( 261 leaf_index < self.number_of_leaves, 262 "Leaf index must be less than the number of leaves in the Merkle tree {leaf_index} , {}", 263 self.number_of_leaves 264 ); 265 266 // Allocate a vector to store the path hashes. 267 let mut path_hashes = Vec::with_capacity(DEPTH as usize); 268 269 // Compute and add the new leaf hash to the path hashes. 270 path_hashes.push(self.leaf_hasher.hash_leaf(new_leaf)?); 271 lap!(timer, "Hashed 1 new leaf"); 272 273 // Compute the start index (on the left) for the leaf hashes level in the Merkle tree. 274 let start = match self.number_of_leaves.checked_next_power_of_two() { 275 Some(num_leaves) => num_leaves - 1, 276 None => bail!("Integer overflow when computing the Merkle tree start index"), 277 }; 278 279 // Compute the new hashes for the path from the leaf to the root. 280 let mut index = start + leaf_index; 281 while let Some(parent) = parent(index) { 282 // Get the left and right child hashes of the parent. 283 let (left, right) = match is_left_child(index) { 284 true => (path_hashes.last().unwrap(), &self.tree[right_child(parent)]), 285 false => (&self.tree[left_child(parent)], path_hashes.last().unwrap()), 286 }; 287 // Compute and add the new parent hash to the path hashes. 288 path_hashes.push(self.path_hasher.hash_children(left, right)?); 289 // Update the index to the parent. 290 index = parent; 291 } 292 293 // Compute the number of levels in the Merkle tree (i.e. log2(tree_size)). 294 let tree_depth = tree_depth::<DEPTH>(self.tree.len())?; 295 // Compute the padding depth. 296 let padding_depth = DEPTH - tree_depth; 297 298 // Update the root hash. 299 // This unwrap is safe, as the path hashes vector is guaranteed to have at least one element. 300 let mut root_hash = *path_hashes.last().unwrap(); 301 for _ in 0..padding_depth { 302 // Update the root hash, by hashing the current root hash with the empty hash. 303 root_hash = self.path_hasher.hash_children(&root_hash, &self.empty_hash)?; 304 } 305 lap!(timer, "Hashed {} padding levels", padding_depth); 306 307 // Initialize the Merkle tree. 308 let mut tree = Vec::with_capacity(self.tree.len()); 309 // Extend the new Merkle tree with the existing leaf hashes. 310 tree.extend(&self.tree); 311 312 // Update the rest of the tree with the new path hashes. 313 let mut index = Some(start + leaf_index); 314 for path_hash in path_hashes { 315 tree[index.unwrap()] = path_hash; 316 index = parent(index.unwrap()); 317 } 318 319 finish!(timer); 320 321 Ok(Self { 322 leaf_hasher: self.leaf_hasher.clone(), 323 path_hasher: self.path_hasher.clone(), 324 root: root_hash, 325 tree, 326 empty_hash: self.empty_hash, 327 number_of_leaves: self.number_of_leaves, 328 }) 329 } 330 331 #[inline] 332 /// Updates the Merkle tree at the location of the given leaf indices with the new leaves. 333 pub fn update_many(&mut self, updates: &BTreeMap<usize, LH::Leaf>) -> Result<()> { 334 let timer = timer!("MerkleTree::update_many"); 335 336 // Check that there are updates to perform. 337 ensure!(!updates.is_empty(), "There must be at least one leaf to update in the Merkle tree"); 338 339 // Check that the latest leaf index is less than number of leaves in the Merkle tree. 340 // Note: This unwrap is safe since updates is guaranteed to be non-empty. 341 ensure!( 342 *updates.last_key_value().unwrap().0 < self.number_of_leaves, 343 "Leaf index must be less than the number of leaves in the Merkle tree" 344 ); 345 346 // Compute the start index (on the left) for the leaf hashes level in the Merkle tree. 347 let start = match self.number_of_leaves.checked_next_power_of_two() { 348 Some(num_leaves) => num_leaves - 1, 349 None => bail!("Integer overflow when computing the Merkle tree start index"), 350 }; 351 352 // A helper to compute the leaf hash. 353 let hash_update = |(leaf_index, leaf): &(&usize, &LH::Leaf)| { 354 self.leaf_hasher.hash_leaf(leaf).map(|hash| (start + **leaf_index, hash)) 355 }; 356 357 // Hash the leaves and add them to the updated hashes. 358 let leaf_hashes: Vec<(usize, LH::Hash)> = match updates.len() { 359 0..=100 => updates.iter().map(|update| hash_update(&update)).collect::<Result<Vec<_>>>()?, 360 _ => cfg_iter!(updates).map(|update| hash_update(&update)).collect::<Result<Vec<_>>>()?, 361 }; 362 lap!(timer, "Hashed {} new leaves", leaf_hashes.len()); 363 364 // Store the updated hashes by level. 365 let mut updated_hashes = Vec::new(); 366 updated_hashes.push(leaf_hashes); 367 368 // A helper function to compute the path hashes for a given level. 369 type Update<PH> = (usize, (<PH as PathHash>::Hash, <PH as PathHash>::Hash)); 370 let compute_path_hashes = |inputs: &[Update<PH>]| match inputs.len() { 371 0..=100 => inputs 372 .iter() 373 .map(|(index, (left, right))| self.path_hasher.hash_children(left, right).map(|hash| (*index, hash))) 374 .collect::<Result<Vec<_>>>(), 375 _ => cfg_iter!(inputs) 376 .map(|(index, (left, right))| self.path_hasher.hash_children(left, right).map(|hash| (*index, hash))) 377 .collect::<Result<Vec<_>>>(), 378 }; 379 380 // Compute the depth of the tree. This corresponds to the number of levels of hashes in the tree. 381 let tree_depth = tree_depth::<DEPTH>(self.tree.len())?; 382 // Allocate a vector to store the inputs to the path hasher. 383 let mut inputs = Vec::with_capacity(updated_hashes[0].len()); 384 // For each level in the tree, compute the path hashes. 385 // In the first iteration, we compute the path hashes for the updated leaf hashes. 386 // In the subsequent iterations, we compute the path hashes for the updated path hashes, until we reach the root. 387 for level in 0..tree_depth as usize { 388 let mut current = 0; 389 while current < updated_hashes[level].len() { 390 let (current_leaf_index, current_leaf_hash) = updated_hashes[level][current]; 391 // Get the sibling of the current leaf. 392 let sibling_leaf_index = match sibling(current_leaf_index) { 393 Some(sibling_index) => sibling_index, 394 // If there is no sibling, then we have reached the root. 395 None => break, 396 }; 397 // Check if the sibling hash is the next hash in the vector. 398 let sibling_is_next_hash = match current + 1 < updated_hashes[level].len() { 399 true => updated_hashes[level][current + 1].0 == sibling_leaf_index, 400 false => false, 401 }; 402 // Get the sibling hash. 403 // Note: This algorithm assumes that the sibling hash is either the next hash in the vector, 404 // or in the original Merkle tree. Consequently, updates need to be provided in sequential order. 405 // This is enforced by the type of `updates: `BTreeMap<usize, LH::Leaf>`. 406 // If this assumption is violated, then the algorithm will compute incorrect path hashes in the Merkle tree. 407 let sibling_leaf_hash = match sibling_is_next_hash { 408 true => updated_hashes[level][current + 1].1, 409 false => self.tree[sibling_leaf_index], 410 }; 411 // Order the current and sibling hashes. 412 let (left, right) = match is_left_child(current_leaf_index) { 413 true => (current_leaf_hash, sibling_leaf_hash), 414 false => (sibling_leaf_hash, current_leaf_hash), 415 }; 416 // Compute the parent index. 417 // Note that this unwrap is safe, since we check that the `current_leaf_index` is not the root. 418 let parent_index = parent(current_leaf_index).unwrap(); 419 // Add the parent hash to the updated hashes. 420 inputs.push((parent_index, (left, right))); 421 // Update the current index. 422 match sibling_is_next_hash { 423 true => current += 2, 424 false => current += 1, 425 } 426 } 427 // Compute the path hashes for the current level. 428 let path_hashes = compute_path_hashes(&inputs)?; 429 // Add the path hashes to the updated hashes. 430 updated_hashes.push(path_hashes); 431 // Clear the inputs. 432 inputs.clear(); 433 } 434 435 // Compute the padding depth. 436 let padding_depth = DEPTH - tree_depth; 437 438 // Update the root hash. 439 // This unwrap is safe, as the updated hashes is guaranteed to have at least one element. 440 let mut root_hash = updated_hashes.last().unwrap()[0].1; 441 for _ in 0..padding_depth { 442 // Update the root hash, by hashing the current root hash with the empty hash. 443 root_hash = self.path_hasher.hash_children(&root_hash, &self.empty_hash)?; 444 } 445 lap!(timer, "Hashed {} padding levels", padding_depth); 446 447 // Update the root hash. 448 self.root = root_hash; 449 450 // Update the rest of the tree with the updated hashes. 451 for (index, hash) in updated_hashes.into_iter().flatten() { 452 self.tree[index] = hash; 453 } 454 455 finish!(timer); 456 Ok(()) 457 } 458 459 #[inline] 460 /// Returns a new Merkle tree with the last 'n' leaves removed from it. 461 pub fn prepare_remove_last_n(&self, n: usize) -> Result<Self> { 462 let timer = timer!("MerkleTree::prepare_remove_last_n"); 463 464 ensure!(n > 0, "Cannot remove zero leaves from the Merkle tree"); 465 466 // Determine the updated number of leaves, after removing the last 'n' leaves. 467 let updated_number_of_leaves = self.number_of_leaves.checked_sub(n).ok_or_else(|| { 468 anyhow!("Failed to remove '{n}' leaves from the Merkle tree, as it only contains {}", self.number_of_leaves) 469 })?; 470 471 // Compute the maximum number of leaves. 472 let max_leaves = match (updated_number_of_leaves).checked_next_power_of_two() { 473 Some(num_leaves) => num_leaves, 474 None => bail!("Integer overflow when computing the maximum number of leaves in the Merkle tree"), 475 }; 476 // Compute the number of nodes. 477 let num_nodes = max_leaves - 1; 478 // Compute the tree size as the maximum number of leaves plus the number of nodes. 479 let tree_size = num_nodes + max_leaves; 480 // Compute the number of levels in the Merkle tree (i.e. log2(tree_size)). 481 let tree_depth = tree_depth::<DEPTH>(tree_size)?; 482 // Compute the number of padded levels. 483 let padding_depth = DEPTH - tree_depth; 484 485 // Calculate the size of the tree which excludes leafless nodes. 486 let minimum_tree_size = std::cmp::max( 487 1, 488 num_nodes 489 + updated_number_of_leaves 490 + if updated_number_of_leaves > 1 { updated_number_of_leaves % 2 } else { 0 }, 491 ); 492 493 // Initialize the Merkle tree. 494 let mut tree = vec![self.empty_hash; num_nodes]; 495 // Extend the new Merkle tree with the existing leaf hashes, excluding the last 'n' leaves. 496 tree.extend(&self.leaf_hashes()?[..updated_number_of_leaves]); 497 // Resize the new Merkle tree with empty hashes to pad up to `tree_size`. 498 tree.resize(minimum_tree_size, self.empty_hash); 499 lap!(timer, "Resizing to {} leaves", updated_number_of_leaves); 500 501 // Initialize a start index to track the starting index of the current level. 502 let start_index = num_nodes; 503 // Initialize a middle index to separate the precomputed indices from the new indices that need to be computed. 504 let middle_index = num_nodes + updated_number_of_leaves; 505 // Initialize a precompute index to track the starting index of each precomputed level. 506 let start_precompute_index = match self.number_of_leaves.checked_next_power_of_two() { 507 Some(num_leaves) => num_leaves - 1, 508 None => bail!("Integer overflow when computing the Merkle tree precompute index"), 509 }; 510 // Initialize a precompute index to track the middle index of each precomputed level. 511 let middle_precompute_index = match num_nodes == start_precompute_index { 512 // If the old tree and new tree are of the same size, then we can copy over the right half of the old tree. 513 true => Some(start_precompute_index + self.number_of_leaves + 1), 514 // true => None, 515 // Otherwise, do nothing, since shrinking the tree is already free. 516 false => None, 517 }; 518 519 // Compute and store the hashes for each level, iterating from the penultimate level to the root level. 520 self.compute_updated_tree( 521 &mut tree, 522 start_index, 523 middle_index, 524 start_precompute_index, 525 middle_precompute_index, 526 )?; 527 528 // Compute the root hash, by iterating from the root level up to `DEPTH`. 529 let mut root_hash = tree[0]; 530 for _ in 0..padding_depth { 531 // Update the root hash, by hashing the current root hash with the empty hash. 532 root_hash = self.path_hasher.hash_children(&root_hash, &self.empty_hash)?; 533 } 534 lap!(timer, "Hashed {} padding levels", padding_depth); 535 536 finish!(timer); 537 538 Ok(Self { 539 leaf_hasher: self.leaf_hasher.clone(), 540 path_hasher: self.path_hasher.clone(), 541 root: root_hash, 542 tree, 543 empty_hash: self.empty_hash, 544 number_of_leaves: updated_number_of_leaves, 545 }) 546 } 547 548 #[inline] 549 /// Updates the Merkle tree with the last 'n' leaves removed from it. 550 pub fn remove_last_n(&mut self, n: usize) -> Result<()> { 551 let timer = timer!("MerkleTree::remove_last_n"); 552 553 // Compute the updated Merkle tree with the last 'n' leaves removed. 554 let updated_tree = self.prepare_remove_last_n(n)?; 555 // Update the tree at the very end, so the original tree is not altered in case of failure. 556 *self = updated_tree; 557 558 finish!(timer); 559 Ok(()) 560 } 561 562 #[inline] 563 /// Returns the Merkle path for the given leaf index and leaf. 564 pub fn prove(&self, leaf_index: usize, leaf: &LH::Leaf) -> Result<MerklePath<E, DEPTH>> { 565 // Ensure the leaf index is valid. 566 ensure!(leaf_index < self.number_of_leaves, "The given Merkle leaf index is out of bounds"); 567 568 // Compute the leaf hash. 569 let leaf_hash = self.leaf_hasher.hash_leaf(leaf)?; 570 571 // Compute the start index (on the left) for the leaf hashes level in the Merkle tree. 572 let start = match self.number_of_leaves.checked_next_power_of_two() { 573 Some(num_leaves) => num_leaves - 1, 574 None => bail!("Integer overflow when computing the Merkle tree start index"), 575 }; 576 577 // Compute the absolute index of the leaf in the Merkle tree. 578 let mut index = start + leaf_index; 579 // Ensure the leaf index is valid. 580 ensure!(index < self.tree.len(), "The given Merkle leaf index is out of bounds"); 581 // Ensure the leaf hash matches the one in the tree. 582 ensure!(self.tree[index] == leaf_hash, "The given Merkle leaf does not match the one in the Merkle tree"); 583 584 // Initialize a vector for the Merkle path. 585 let mut path = Vec::with_capacity(DEPTH as usize); 586 587 // Iterate from the leaf hash to the root level, storing the sibling hashes along the path. 588 for _ in 0..DEPTH { 589 // Compute the index of the sibling hash, if it exists. 590 if let Some(sibling) = sibling(index) { 591 // Append the sibling hash to the path. 592 path.push(self.tree[sibling]); 593 // Compute the index of the parent hash, if it exists. 594 match parent(index) { 595 // Update the index to the parent index. 596 Some(parent) => index = parent, 597 // If the parent does not exist, the path is complete. 598 None => break, 599 } 600 } 601 } 602 603 // If the Merkle path length is not equal to `DEPTH`, pad the path with the empty hash. 604 path.resize(DEPTH as usize, self.empty_hash); 605 606 // Return the Merkle path. 607 MerklePath::try_from((U64::new(leaf_index as u64), path)) 608 } 609 610 /// Returns `true` if the given Merkle path is valid for the given root and leaf. 611 pub fn verify(&self, path: &MerklePath<E, DEPTH>, root: &PH::Hash, leaf: &LH::Leaf) -> bool { 612 path.verify(&self.leaf_hasher, &self.path_hasher, root, leaf) 613 } 614 615 /// Returns the Merkle root of the tree. 616 pub const fn root(&self) -> &PH::Hash { 617 &self.root 618 } 619 620 /// Returns the Merkle tree (excluding the hashes of the leaves). 621 pub fn tree(&self) -> &[PH::Hash] { 622 &self.tree 623 } 624 625 /// Returns the empty hash. 626 pub const fn empty_hash(&self) -> &PH::Hash { 627 &self.empty_hash 628 } 629 630 /// Returns the leaf hashes from the Merkle tree. 631 pub fn leaf_hashes(&self) -> Result<&[LH::Hash]> { 632 // Compute the start index (on the left) for the leaf hashes level in the Merkle tree. 633 let start = match self.number_of_leaves.checked_next_power_of_two() { 634 Some(num_leaves) => num_leaves - 1, 635 None => bail!("Integer overflow when computing the Merkle tree start index"), 636 }; 637 // Compute the end index (on the right) for the leaf hashes level in the Merkle tree. 638 let end = start + self.number_of_leaves; 639 // Return the leaf hashes. 640 Ok(&self.tree[start..end]) 641 } 642 643 /// Returns the number of leaves in the Merkle tree. 644 pub const fn number_of_leaves(&self) -> usize { 645 self.number_of_leaves 646 } 647 648 /// Compute and store the hashes for each level, iterating from the penultimate level to the root level. 649 /// 650 /// ```ignore 651 /// start_index middle_index end_index 652 /// start_precompute_index middle_precompute_index end_index 653 /// ``` 654 #[inline] 655 fn compute_updated_tree( 656 &self, 657 tree: &mut [Field<E>], 658 mut start_index: usize, 659 mut middle_index: usize, 660 mut start_precompute_index: usize, 661 mut middle_precompute_index: Option<usize>, 662 ) -> Result<()> { 663 // Initialize a timer for the while loop. 664 let timer = timer!("MerkleTree::compute_updated_tree"); 665 666 // Compute and store the hashes for each level, iterating from the penultimate level to the root level. 667 let empty_hash = self.path_hasher.hash_empty()?; 668 while let (Some(start), Some(middle)) = (parent(start_index), parent(middle_index)) { 669 // Compute the end index of the current level. 670 let end = left_child(start); 671 672 // If the current level has precomputed indices, copy them instead of recomputing them. 673 if let Some(start_precompute) = parent(start_precompute_index) { 674 // Compute the end index of the precomputed level. 675 let end_precompute = start_precompute + (middle - start); 676 // Copy the hashes for each node in the current level. 677 tree[start..middle].copy_from_slice(&self.tree[start_precompute..end_precompute]); 678 // Update the precompute index for the next level. 679 start_precompute_index = start_precompute; 680 } else { 681 // Ensure the start index is equal to the middle index, as all precomputed indices have been processed. 682 ensure!(start == middle, "Failed to process all left precomputed indices in the Merkle tree"); 683 } 684 lap!(timer, "Precompute (Left): {start} -> {middle}"); 685 686 // If the current level has precomputed indices, copy them instead of recomputing them. 687 // Note: This logic works because the old tree and new tree are the same power of two. 688 if let Some(middle_precompute) = middle_precompute_index { 689 if let Some(middle_precompute) = parent(middle_precompute) { 690 // Construct the children for the new indices in the current level. 691 let tuples = (middle..middle_precompute) 692 .map(|i| { 693 ( 694 tree.get(left_child(i)).copied().unwrap_or(empty_hash), 695 tree.get(right_child(i)).copied().unwrap_or(empty_hash), 696 ) 697 }) 698 .collect::<Vec<_>>(); 699 // Process the indices that need to be computed for the current level. 700 // If any level requires computing more than 100 nodes, borrow the tree for performance. 701 match tuples.len() >= 100 { 702 // Option 1: Borrow the tree to compute and store the hashes for the new indices in the current level. 703 true => cfg_iter_mut!(tree[middle..middle_precompute]).zip_eq(cfg_iter!(tuples)).try_for_each( 704 |(node, (left, right))| { 705 *node = self.path_hasher.hash_children(left, right)?; 706 Ok::<_, Error>(()) 707 }, 708 )?, 709 // Option 2: Compute and store the hashes for the new indices in the current level. 710 false => tree[middle..middle_precompute].iter_mut().zip_eq(&tuples).try_for_each( 711 |(node, (left, right))| { 712 *node = self.path_hasher.hash_children(left, right)?; 713 Ok::<_, Error>(()) 714 }, 715 )?, 716 } 717 lap!(timer, "Compute: {middle} -> {middle_precompute}"); 718 719 // Copy the hashes for each node in the current level. 720 tree[middle_precompute..end].copy_from_slice(&self.tree[middle_precompute..end]); 721 // Update the precompute index for the next level. 722 middle_precompute_index = Some(middle_precompute + 1); 723 lap!(timer, "Precompute (Right): {middle_precompute} -> {end}"); 724 } else { 725 // Ensure the middle precompute index is equal to the end index, as all precomputed indices have been processed. 726 ensure!( 727 middle_precompute == end, 728 "Failed to process all right precomputed indices in the Merkle tree" 729 ); 730 } 731 } else { 732 // Construct the children for the new indices in the current level. 733 let tuples = (middle..end) 734 .map(|i| { 735 ( 736 tree.get(left_child(i)).copied().unwrap_or(empty_hash), 737 tree.get(right_child(i)).copied().unwrap_or(empty_hash), 738 ) 739 }) 740 .collect::<Vec<_>>(); 741 // Process the indices that need to be computed for the current level. 742 // If any level requires computing more than 100 nodes, borrow the tree for performance. 743 match tuples.len() >= 100 { 744 // Option 1: Borrow the tree to compute and store the hashes for the new indices in the current level. 745 true => cfg_iter_mut!(tree[middle..end]).zip_eq(cfg_iter!(tuples)).try_for_each( 746 |(node, (left, right))| { 747 *node = self.path_hasher.hash_children(left, right)?; 748 Ok::<_, Error>(()) 749 }, 750 )?, 751 // Option 2: Compute and store the hashes for the new indices in the current level. 752 false => tree[middle..end].iter_mut().zip_eq(&tuples).try_for_each(|(node, (left, right))| { 753 *node = self.path_hasher.hash_children(left, right)?; 754 Ok::<_, Error>(()) 755 })?, 756 } 757 lap!(timer, "Compute: {middle} -> {end}"); 758 } 759 760 // Update the start index for the next level. 761 start_index = start; 762 // Update the middle index for the next level. 763 middle_index = middle; 764 } 765 766 // End the timer for the while loop. 767 finish!(timer); 768 769 Ok(()) 770 } 771 } 772 773 /// Returns the depth of the tree, given the size of the tree. 774 #[inline] 775 fn tree_depth<const DEPTH: u8>(tree_size: usize) -> Result<u8> { 776 let tree_size = u64::try_from(tree_size)?; 777 // Since we only allow tree sizes up to u64::MAX, the maximum possible depth is 63. 778 let tree_depth = u8::try_from(tree_size.checked_ilog2().unwrap_or(0))?; 779 // Ensure the tree depth is within the depth bound. 780 match tree_depth <= DEPTH { 781 // Return the tree depth. 782 true => Ok(tree_depth), 783 false => bail!("Merkle tree cannot exceed depth {DEPTH}: attempted to reach depth {tree_depth}"), 784 } 785 } 786 787 /// Returns the index of the left child, given an index. 788 #[inline] 789 const fn left_child(index: usize) -> usize { 790 2 * index + 1 791 } 792 793 /// Returns the index of the right child, given an index. 794 #[inline] 795 const fn right_child(index: usize) -> usize { 796 2 * index + 2 797 } 798 799 /// Returns the index of the sibling, given an index. 800 #[inline] 801 const fn sibling(index: usize) -> Option<usize> { 802 if is_root(index) { 803 None 804 } else if is_left_child(index) { 805 Some(index + 1) 806 } else { 807 Some(index - 1) 808 } 809 } 810 811 /// Returns true iff the index represents the root. 812 #[inline] 813 const fn is_root(index: usize) -> bool { 814 index == 0 815 } 816 817 /// Returns true iff the given index represents a left child. 818 #[inline] 819 const fn is_left_child(index: usize) -> bool { 820 index % 2 == 1 821 } 822 823 /// Returns the index of the parent, given the index of a child. 824 #[inline] 825 const fn parent(index: usize) -> Option<usize> { 826 if index > 0 { Some((index - 1) >> 1) } else { None } 827 }