diff --git a/src/interface.rs b/src/interface.rs index 6f27952..7cb6153 100644 --- a/src/interface.rs +++ b/src/interface.rs @@ -2,7 +2,7 @@ use crate::level_iter::LevelIter; use crate::update_map::UpdateMap; use crate::utils::{Length, updated_length}; use crate::{ - Cow, Error, Value, + Arc, Cow, Error, Tree, Value, interface_iter::{InterfaceIter, InterfaceIterCow}, iter::Iter, }; @@ -22,6 +22,12 @@ pub trait ImmList { fn iter_from(&self, index: usize) -> Iter<'_, T>; fn level_iter_from(&self, index: usize) -> LevelIter<'_, T>; + + fn tree(&self) -> &Arc>; + + fn depth(&self) -> usize; + + fn packing_depth(&self) -> usize; } pub trait MutList: ImmList { @@ -144,6 +150,248 @@ where self.updates = updates; Ok(()) } + + /// Compute the tree hash root, accounting for pending updates. + /// + /// This method computes the tree hash without applying pending updates to the backing tree. + /// It recursively traverses the tree, using values from pending updates where they exist, + /// and cached hashes from the tree where they don't. + pub fn tree_hash_root(&self, full_length: usize) -> Hash256 + where + T: Send + Sync, + B: ImmList, + { + let tree = self.backing.tree(); + let depth = self.backing.depth(); + let packing_depth = self.backing.packing_depth(); + + if self.updates.is_empty() { + // No pending updates, just use the cached tree hash + tree.tree_hash() + } else { + // Compute tree hash while considering pending updates + self.tree_hash_recursive(tree, &self.updates, 0, depth, packing_depth, full_length) + } + } + + /// Recursively compute tree hash for a subtree, considering pending updates. + fn tree_hash_recursive( + &self, + node: &Tree, + updates: &U, + prefix: usize, + depth: usize, + packing_depth: usize, + full_length: usize, + ) -> Hash256 + where + T: Send + Sync, + { + use crate::tree::Tree; + use ethereum_hashing::{ZERO_HASHES, hash32_concat}; + use std::ops::ControlFlow; + + match node { + Tree::Leaf(leaf) if depth == 0 => { + // Check if this leaf has a pending update + if let Some(updated_value) = updates.get(prefix) { + updated_value.tree_hash_root() + } else { + // Use cached hash if available, otherwise compute + let read_lock = leaf.hash.read(); + let existing_hash = *read_lock; + drop(read_lock); + + if !existing_hash.is_zero() { + existing_hash + } else { + leaf.value.tree_hash_root() + } + } + } + Tree::PackedLeaf(packed_leaf) if depth == 0 => { + // Check if any values in this packed leaf have pending updates + let mut has_updates = false; + let packing_factor = T::tree_hash_packing_factor(); + for i in 0..packing_factor { + let index = prefix + i; + if index >= full_length { + break; + } + if updates.get(index).is_some() { + has_updates = true; + break; + } + } + + if has_updates { + // Need to recompute hash with updated values + let mut values = packed_leaf.values.clone(); + for i in 0..packing_factor { + let index = prefix + i; + if index >= full_length { + break; + } + if let Some(updated_value) = updates.get(index) { + if i < values.len() { + values[i] = updated_value.clone(); + } else { + values.push(updated_value.clone()); + } + } + } + // Compute tree hash for packed values + use crate::packed_leaf::PackedLeaf; + PackedLeaf { + values, + hash: parking_lot::RwLock::new(Hash256::ZERO), + } + .tree_hash() + } else { + // No updates, use cached hash + packed_leaf.tree_hash() + } + } + Tree::Zero(zero_depth) if *zero_depth == depth => { + // Check if there are any updates for this zero tree's range + let subtree_start = prefix; + let subtree_end = prefix + (1 << (depth + packing_depth)); + let mut has_updates = false; + let _ = updates.for_each_range::<_, Error>(subtree_start, subtree_end, |_, _| { + has_updates = true; + ControlFlow::Break(()) + }); + + if !has_updates { + // No updates in this zero tree, return zero hash + Hash256::from(ZERO_HASHES[depth]) + } else if depth == 0 { + // Leaf level with updates + if let Some(packing_factor) = crate::utils::opt_packing_factor::() { + // Packed leaf case + let mut values = Vec::with_capacity(packing_factor); + for i in 0..packing_factor { + let index = prefix + i; + if index >= full_length { + break; + } + if let Some(value) = updates.get(index) { + values.push(value.clone()); + } + } + if !values.is_empty() { + use crate::packed_leaf::PackedLeaf; + PackedLeaf { + values, + hash: parking_lot::RwLock::new(Hash256::ZERO), + } + .tree_hash() + } else { + Hash256::from(ZERO_HASHES[depth]) + } + } else { + // Regular leaf case + if let Some(value) = updates.get(prefix) { + value.tree_hash_root() + } else { + Hash256::from(ZERO_HASHES[depth]) + } + } + } else { + // Internal node with updates - need to recursively build from zero trees + let new_depth = depth - 1; + let left_prefix = prefix; + let right_prefix = prefix | (1 << (new_depth + packing_depth)); + + let left_zero = Tree::zero(new_depth); + let right_zero = Tree::zero(new_depth); + + let left_hash = self.tree_hash_recursive( + &left_zero, + updates, + left_prefix, + new_depth, + packing_depth, + full_length, + ); + let right_hash = self.tree_hash_recursive( + &right_zero, + updates, + right_prefix, + new_depth, + packing_depth, + full_length, + ); + + Hash256::from(hash32_concat(left_hash.as_slice(), right_hash.as_slice())) + } + } + Tree::Node { hash, left, right } if depth > 0 => { + let new_depth = depth - 1; + let left_prefix = prefix; + let right_prefix = prefix | (1 << (new_depth + packing_depth)); + let right_subtree_end = prefix + (1 << (depth + packing_depth)); + + // Check if there are updates in left or right subtrees + let mut has_left_updates = false; + let _ = updates.for_each_range::<_, Error>(left_prefix, right_prefix, |_, _| { + has_left_updates = true; + ControlFlow::Break(()) + }); + + let mut has_right_updates = false; + let _ = + updates.for_each_range::<_, Error>(right_prefix, right_subtree_end, |_, _| { + has_right_updates = true; + ControlFlow::Break(()) + }); + + if !has_left_updates && !has_right_updates { + // No updates in this subtree, use cached hash if available + let read_lock = hash.read(); + let existing_hash = *read_lock; + drop(read_lock); + + if !existing_hash.is_zero() { + return existing_hash; + } + } + + // Compute hashes for left and right subtrees + let left_hash = if has_left_updates { + self.tree_hash_recursive( + left, + updates, + left_prefix, + new_depth, + packing_depth, + full_length, + ) + } else { + left.tree_hash() + }; + + let right_hash = if has_right_updates { + self.tree_hash_recursive( + right, + updates, + right_prefix, + new_depth, + packing_depth, + full_length, + ) + } else { + right.tree_hash() + }; + + Hash256::from(hash32_concat(left_hash.as_slice(), right_hash.as_slice())) + } + _ => { + // Fallback to regular tree hash for unexpected cases + node.tree_hash() + } + } + } } #[cfg(test)] diff --git a/src/list.rs b/src/list.rs index 069dcfb..82494fd 100644 --- a/src/list.rs +++ b/src/list.rs @@ -277,6 +277,18 @@ impl ImmList for ListInner { fn level_iter_from(&self, index: usize) -> LevelIter<'_, T> { LevelIter::from_index(index, &self.tree, self.depth, self.length) } + + fn tree(&self) -> &Arc> { + &self.tree + } + + fn depth(&self) -> usize { + self.depth + } + + fn packing_depth(&self) -> usize { + self.packing_depth + } } impl MutList for ListInner @@ -393,10 +405,7 @@ impl> TreeHash for List Hash256 { - // FIXME(sproul): remove assert - assert!(!self.interface.has_pending_updates()); - - let root = self.interface.backing.tree.tree_hash(); + let root = self.interface.tree_hash_root(self.len()); tree_hash::mix_in_length(&root, self.len()) } } diff --git a/src/tests/mod.rs b/src/tests/mod.rs index b6010f3..7286b09 100644 --- a/src/tests/mod.rs +++ b/src/tests/mod.rs @@ -9,3 +9,4 @@ mod proptest; mod repeat; mod serde; mod size_of; +mod tree_hash_pending; diff --git a/src/tests/proptest/operations.rs b/src/tests/proptest/operations.rs index 699a2fa..790924d 100644 --- a/src/tests/proptest/operations.rs +++ b/src/tests/proptest/operations.rs @@ -2,6 +2,7 @@ use super::{Large, arb_hash256, arb_index, arb_large, arb_list, arb_vect}; use crate::{Error, List, Value, Vector}; use proptest::prelude::*; use ssz::{Decode, Encode}; +use ssz_types::{FixedVector, VariableList}; use std::fmt::Debug; use std::marker::PhantomData; use std::ops::Deref; @@ -90,6 +91,19 @@ impl Spec { Err(Error::PushNotSupported) } } + + pub fn tree_hash_root(&self) -> Hash256 + where + T: Send + Sync, + { + if self.allow_push { + let var_list = VariableList::::new(self.values.clone()).unwrap(); + var_list.tree_hash_root() + } else { + let fixed_vect = FixedVector::::new(self.values.clone()).unwrap(); + fixed_vect.tree_hash_root() + } + } } #[derive(Debug, Clone)] @@ -251,8 +265,7 @@ where checkpoint = list.clone(); } Op::TreeHash => { - list.apply_updates().unwrap(); - list.tree_hash_root(); + assert_eq!(list.tree_hash_root(), spec.tree_hash_root()); } Op::Rebase => { list.apply_updates().unwrap(); @@ -354,8 +367,7 @@ where vect.apply_updates().unwrap(); } Op::TreeHash => { - vect.apply_updates().unwrap(); - vect.tree_hash_root(); + assert_eq!(vect.tree_hash_root(), spec.tree_hash_root()); } Op::Checkpoint => { vect.apply_updates().unwrap(); diff --git a/src/tests/tree_hash_pending.rs b/src/tests/tree_hash_pending.rs new file mode 100644 index 0000000..77a8635 --- /dev/null +++ b/src/tests/tree_hash_pending.rs @@ -0,0 +1,285 @@ +use crate::{List, Vector}; +use tree_hash::TreeHash; +use typenum::{U8, U32, U1024}; + +#[test] +fn tree_hash_with_pending_updates() { + let mut list = List::::new(vec![1, 2, 3, 4]).unwrap(); + + // Compute initial hash (no pending updates) + let hash1 = list.tree_hash_root(); + + // Make some updates without applying them + *list.get_mut(0).unwrap() = 10; + *list.get_mut(2).unwrap() = 30; + + // Verify we have pending updates + assert!(list.has_pending_updates()); + + // Compute hash WITH pending updates (without calling apply_updates) + let hash2 = list.tree_hash_root(); + + // Hashes should be different + assert_ne!(hash1, hash2); + + // Apply updates and compute hash again + list.apply_updates().unwrap(); + let hash3 = list.tree_hash_root(); + + // Hash after applying updates should match hash computed with pending updates + assert_eq!(hash2, hash3); +} + +#[test] +fn tree_hash_with_pending_push() { + let mut list = List::::new(vec![1, 2, 3]).unwrap(); + + // Compute initial hash + let hash1 = list.tree_hash_root(); + + // Push without applying + list.push(4).unwrap(); + assert!(list.has_pending_updates()); + + // Compute hash with pending push + let hash2 = list.tree_hash_root(); + assert_ne!(hash1, hash2); + + // Apply and verify + list.apply_updates().unwrap(); + let hash3 = list.tree_hash_root(); + assert_eq!(hash2, hash3); + + // Compare with a list built with the final values + let list_direct = List::::new(vec![1, 2, 3, 4]).unwrap(); + assert_eq!(hash3, list_direct.tree_hash_root()); +} + +#[test] +fn tree_hash_packed_with_pending_updates() { + // Test with packed types (u8) + let mut list = List::::new(vec![1, 2, 3, 4, 5, 6, 7, 8]).unwrap(); + + let hash1 = list.tree_hash_root(); + + // Update some packed values + *list.get_mut(0).unwrap() = 10; + *list.get_mut(7).unwrap() = 80; + + assert!(list.has_pending_updates()); + + // Hash with pending updates + let hash2 = list.tree_hash_root(); + assert_ne!(hash1, hash2); + + // Apply and verify + list.apply_updates().unwrap(); + let hash3 = list.tree_hash_root(); + assert_eq!(hash2, hash3); +} + +#[test] +fn tree_hash_mixed_pending_and_applied() { + let mut list = List::::new(vec![1, 2, 3, 4]).unwrap(); + + // Apply some updates + *list.get_mut(0).unwrap() = 10; + list.apply_updates().unwrap(); + let hash1 = list.tree_hash_root(); + + // Add more pending updates + *list.get_mut(1).unwrap() = 20; + *list.get_mut(2).unwrap() = 30; + assert!(list.has_pending_updates()); + + // Hash with new pending updates + let hash2 = list.tree_hash_root(); + assert_ne!(hash1, hash2); + + // Apply and verify + list.apply_updates().unwrap(); + let hash3 = list.tree_hash_root(); + assert_eq!(hash2, hash3); + + // Compare with directly constructed list + let list_direct = List::::new(vec![10, 20, 30, 4]).unwrap(); + assert_eq!(hash3, list_direct.tree_hash_root()); +} + +#[test] +fn tree_hash_empty_list_with_push() { + // Test empty list with pending push + let mut list = List::::empty(); + + let hash1 = list.tree_hash_root(); + + // Push to empty list + list.push(1).unwrap(); + assert!(list.has_pending_updates()); + + let hash2 = list.tree_hash_root(); + assert_ne!(hash1, hash2); + + // Apply and verify + list.apply_updates().unwrap(); + let hash3 = list.tree_hash_root(); + assert_eq!(hash2, hash3); + + // Compare with directly constructed list + let list_direct = List::::new(vec![1]).unwrap(); + assert_eq!(hash3, list_direct.tree_hash_root()); +} + +#[test] +fn tree_hash_large_list_with_pending_updates() { + // Test with a larger list to exercise deeper tree structures + let vec: Vec = (0..100).collect(); + let mut list = List::::new(vec.clone()).unwrap(); + + let hash1 = list.tree_hash_root(); + + // Update multiple elements at different depths in the tree + *list.get_mut(0).unwrap() = 1000; + *list.get_mut(16).unwrap() = 1016; + *list.get_mut(50).unwrap() = 1050; + *list.get_mut(99).unwrap() = 1099; + + assert!(list.has_pending_updates()); + + let hash2 = list.tree_hash_root(); + assert_ne!(hash1, hash2); + + // Apply and verify + list.apply_updates().unwrap(); + let hash3 = list.tree_hash_root(); + assert_eq!(hash2, hash3); +} + +#[test] +fn tree_hash_multiple_pending_pushes() { + // Test multiple pushes without applying + let mut list = List::::new(vec![1, 2]).unwrap(); + + let hash1 = list.tree_hash_root(); + + // Push multiple elements + list.push(3).unwrap(); + list.push(4).unwrap(); + list.push(5).unwrap(); + + assert!(list.has_pending_updates()); + + let hash2 = list.tree_hash_root(); + assert_ne!(hash1, hash2); + + // Apply and verify + list.apply_updates().unwrap(); + let hash3 = list.tree_hash_root(); + assert_eq!(hash2, hash3); + + // Compare with directly constructed list + let list_direct = List::::new(vec![1, 2, 3, 4, 5]).unwrap(); + assert_eq!(hash3, list_direct.tree_hash_root()); +} + +#[test] +fn tree_hash_vector_with_pending_updates() { + // Test Vector with pending updates + let mut vector = Vector::::new(vec![1, 2, 3, 4, 5, 6, 7, 8]).unwrap(); + + let hash1 = vector.tree_hash_root(); + + // Update some elements + *vector.get_mut(0).unwrap() = 10; + *vector.get_mut(7).unwrap() = 80; + + assert!(vector.has_pending_updates()); + + let hash2 = vector.tree_hash_root(); + assert_ne!(hash1, hash2); + + // Apply and verify + vector.apply_updates().unwrap(); + let hash3 = vector.tree_hash_root(); + assert_eq!(hash2, hash3); +} + +#[test] +fn tree_hash_packed_multiple_in_same_leaf() { + // Test updating multiple values in the same packed leaf + let mut list = List::::new(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]).unwrap(); + + let hash1 = list.tree_hash_root(); + + // Update multiple values that would be in the same packed leaf + *list.get_mut(0).unwrap() = 10; + *list.get_mut(1).unwrap() = 20; + *list.get_mut(2).unwrap() = 30; + + assert!(list.has_pending_updates()); + + let hash2 = list.tree_hash_root(); + assert_ne!(hash1, hash2); + + // Apply and verify + list.apply_updates().unwrap(); + let hash3 = list.tree_hash_root(); + assert_eq!(hash2, hash3); + + // Compare with directly constructed list + let list_direct = List::::new(vec![10, 20, 30, 4, 5, 6, 7, 8, 9, 10]).unwrap(); + assert_eq!(hash3, list_direct.tree_hash_root()); +} + +#[test] +fn tree_hash_single_element_list() { + // Test with a single element list + let mut list = List::::new(vec![1]).unwrap(); + + let hash1 = list.tree_hash_root(); + + // Update the only element + *list.get_mut(0).unwrap() = 100; + + assert!(list.has_pending_updates()); + + let hash2 = list.tree_hash_root(); + assert_ne!(hash1, hash2); + + // Apply and verify + list.apply_updates().unwrap(); + let hash3 = list.tree_hash_root(); + assert_eq!(hash2, hash3); + + // Compare with directly constructed list + let list_direct = List::::new(vec![100]).unwrap(); + assert_eq!(hash3, list_direct.tree_hash_root()); +} + +#[test] +fn tree_hash_all_elements_updated() { + // Test updating all elements in a list + let mut list = List::::new(vec![1, 2, 3, 4]).unwrap(); + + let hash1 = list.tree_hash_root(); + + // Update all elements + *list.get_mut(0).unwrap() = 10; + *list.get_mut(1).unwrap() = 20; + *list.get_mut(2).unwrap() = 30; + *list.get_mut(3).unwrap() = 40; + + assert!(list.has_pending_updates()); + + let hash2 = list.tree_hash_root(); + assert_ne!(hash1, hash2); + + // Apply and verify + list.apply_updates().unwrap(); + let hash3 = list.tree_hash_root(); + assert_eq!(hash2, hash3); + + // Compare with directly constructed list + let list_direct = List::::new(vec![10, 20, 30, 40]).unwrap(); + assert_eq!(hash3, list_direct.tree_hash_root()); +} diff --git a/src/vector.rs b/src/vector.rs index 1e6cda7..67e75c8 100644 --- a/src/vector.rs +++ b/src/vector.rs @@ -238,6 +238,18 @@ impl ImmList for VectorInner { fn level_iter_from(&self, index: usize) -> LevelIter<'_, T> { LevelIter::from_index(index, &self.tree, self.depth, Length(N::to_usize())) } + + fn tree(&self) -> &Arc> { + &self.tree + } + + fn depth(&self) -> usize { + self.depth + } + + fn packing_depth(&self) -> usize { + self.packing_depth + } } impl MutList for VectorInner @@ -306,9 +318,8 @@ impl> tree_hash::TreeHash f } fn tree_hash_root(&self) -> Hash256 { - // FIXME(sproul): remove assert - assert!(!self.interface.has_pending_updates()); - self.interface.backing.tree.tree_hash() + // For vectors, the length is fixed at N, so we use that as the full_length + self.interface.tree_hash_root(N::to_usize()) } }