Make List::tree_hash_root infallible by computing hashes with pending updates

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<T: Value> {
     fn iter_from(&self, index: usize) -> Iter<'_, T>;
 
     fn level_iter_from(&self, index: usize) -> LevelIter<'_, T>;
+
+    fn tree(&self) -> &Arc<Tree<T>>;
+
+    fn depth(&self) -> usize;
+
+    fn packing_depth(&self) -> usize;
 }
 
 pub trait MutList<T: Value>: ImmList<T> {
@@ -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<T>,
+    {
+        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<T>,
+        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::<T>() {
+                        // 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)]

src/list.rs

@@ -277,6 +277,18 @@ impl<T: Value, N: Unsigned> ImmList<T> for ListInner<T, N> {
     fn level_iter_from(&self, index: usize) -> LevelIter<'_, T> {
         LevelIter::from_index(index, &self.tree, self.depth, self.length)
     }
+
+    fn tree(&self) -> &Arc<Tree<T>> {
+        &self.tree
+    }
+
+    fn depth(&self) -> usize {
+        self.depth
+    }
+
+    fn packing_depth(&self) -> usize {
+        self.packing_depth
+    }
 }
 
 impl<T, N> MutList<T> for ListInner<T, N>
@@ -393,10 +405,7 @@ impl<T: Value + Send + Sync, N: Unsigned, U: UpdateMap<T>> TreeHash for List<T,
     }
 
     fn tree_hash_root(&self) -> 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())
     }
 }

src/tests/mod.rs

@@ -9,3 +9,4 @@ mod proptest;
 mod repeat;
 mod serde;
 mod size_of;
+mod tree_hash_pending;

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<T: Value, N: Unsigned> Spec<T, N> {
             Err(Error::PushNotSupported)
         }
     }
+
+    pub fn tree_hash_root(&self) -> Hash256
+    where
+        T: Send + Sync,
+    {
+        if self.allow_push {
+            let var_list = VariableList::<T, N>::new(self.values.clone()).unwrap();
+            var_list.tree_hash_root()
+        } else {
+            let fixed_vect = FixedVector::<T, N>::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();