Compute packed leaf hash on demand (Option 2)

src/packed_leaf.rs

@@ -1,16 +1,12 @@
-use crate::{utils::arb_rwlock, Error, UpdateMap};
+use crate::{Error, UpdateMap};
 use arbitrary::Arbitrary;
 use derivative::Derivative;
-use parking_lot::RwLock;
 use std::ops::ControlFlow;
 use tree_hash::{Hash256, TreeHash, BYTES_PER_CHUNK};
 
 #[derive(Debug, Derivative, Arbitrary)]
 #[derivative(PartialEq, Hash)]
 pub struct PackedLeaf<T: TreeHash + Clone> {
-    #[derivative(PartialEq = "ignore", Hash = "ignore")]
-    #[arbitrary(with = arb_rwlock)]
-    pub hash: RwLock<Hash256>,
     pub(crate) values: Vec<T>,
 }
 
@@ -20,21 +16,14 @@ where
 {
     fn clone(&self) -> Self {
         Self {
-            hash: RwLock::new(*self.hash.read()),
             values: self.values.clone(),
         }
     }
 }
 
 impl<T: TreeHash + Clone> PackedLeaf<T> {
     pub fn tree_hash(&self) -> Hash256 {
-        let read_lock = self.hash.read();
-        let mut hash = *read_lock;
-        drop(read_lock);
-
-        if !hash.is_zero() {
-            return hash;
-        }
+        let mut hash = Hash256::zero();
 
         let hash_bytes = hash.as_bytes_mut();
 
@@ -44,13 +33,11 @@ impl<T: TreeHash + Clone> PackedLeaf<T> {
                 .copy_from_slice(&value.tree_hash_packed_encoding());
         }
 
-        *self.hash.write() = hash;
         hash
     }
 
     pub fn empty() -> Self {
         PackedLeaf {
-            hash: RwLock::new(Hash256::zero()),
             values: Vec::with_capacity(T::tree_hash_packing_factor()),
         }
     }
@@ -59,38 +46,27 @@ impl<T: TreeHash + Clone> PackedLeaf<T> {
         let mut values = Vec::with_capacity(T::tree_hash_packing_factor());
         values.push(value);
 
-        PackedLeaf {
-            hash: RwLock::new(Hash256::zero()),
-            values,
-        }
+        PackedLeaf { values }
     }
 
     pub fn repeat(value: T, n: usize) -> Self {
         assert!(n <= T::tree_hash_packing_factor());
         PackedLeaf {
-            hash: RwLock::new(Hash256::zero()),
             values: vec![value; n],
         }
     }
 
     pub fn insert_at_index(&self, index: usize, value: T) -> Result<Self, Error> {
         let mut updated = PackedLeaf {
-            hash: RwLock::new(Hash256::zero()),
             values: self.values.clone(),
         };
         let sub_index = index % T::tree_hash_packing_factor();
         updated.insert_mut(sub_index, value)?;
         Ok(updated)
     }
 
-    pub fn update<U: UpdateMap<T>>(
-        &self,
-        prefix: usize,
-        hash: Hash256,
-        updates: &U,
-    ) -> Result<Self, Error> {
+    pub fn update<U: UpdateMap<T>>(&self, prefix: usize, updates: &U) -> Result<Self, Error> {
         let mut updated = PackedLeaf {
-            hash: RwLock::new(hash),
             values: self.values.clone(),
         };
 
@@ -104,9 +80,6 @@ impl<T: TreeHash + Clone> PackedLeaf<T> {
     }
 
     pub fn insert_mut(&mut self, sub_index: usize, value: T) -> Result<(), Error> {
-        // Ensure hash is 0.
-        *self.hash.get_mut() = Hash256::zero();
-
         if sub_index == self.values.len() {
             self.values.push(value);
         } else if sub_index < self.values.len() {

src/tests/size_of.rs

@@ -6,9 +6,9 @@ use tree_hash::Hash256;
 /// It's important that the Tree nodes have a predictable size.
 #[test]
 fn size_of_hash256() {
-    assert_eq!(size_of::<Tree<Hash256>>(), 72);
+    assert_eq!(size_of::<Tree<Hash256>>(), 64);
     assert_eq!(size_of::<Leaf<Hash256>>(), 48);
-    assert_eq!(size_of::<PackedLeaf<Hash256>>(), 64);
+    assert_eq!(size_of::<PackedLeaf<Hash256>>(), 24);
 
     let rw_lock_size = size_of::<RwLock<Hash256>>();
     assert_eq!(rw_lock_size, 40);
@@ -18,26 +18,23 @@ fn size_of_hash256() {
 
     assert_eq!(
         size_of::<Tree<Hash256>>(),
-        size_of::<PackedLeaf<Hash256>>() + 8
+        size_of::<PackedLeaf<Hash256>>() + 40
     );
 }
 
 /// It's important that the Tree nodes have a predictable size.
 #[test]
 fn size_of_u8() {
-    assert_eq!(size_of::<Tree<u8>>(), 72);
+    assert_eq!(size_of::<Tree<u8>>(), 64);
     assert_eq!(size_of::<Leaf<u8>>(), 48);
-    assert_eq!(size_of::<PackedLeaf<u8>>(), 64);
-    assert_eq!(
-        size_of::<PackedLeaf<u8>>(),
-        size_of::<RwLock<Hash256>>() + size_of::<Vec<u8>>()
-    );
+    assert_eq!(size_of::<PackedLeaf<u8>>(), 24);
+    assert_eq!(size_of::<PackedLeaf<u8>>(), size_of::<Vec<u8>>());
 
     let rw_lock_size = size_of::<RwLock<u8>>();
     assert_eq!(rw_lock_size, 16);
 
     let arc_size = size_of::<Arc<Tree<u8>>>();
     assert_eq!(arc_size, 8);
 
-    assert_eq!(size_of::<Tree<u8>>(), size_of::<PackedLeaf<u8>>() + 8);
+    assert_eq!(size_of::<Tree<u8>>(), size_of::<PackedLeaf<u8>>() + 40);
 }

src/tree.rs

@@ -176,7 +176,7 @@ impl<T: TreeHash + Clone> Tree<T> {
                 Ok(Self::leaf_with_hash(value, hash))
             }
             Self::PackedLeaf(packed_leaf) if depth == 0 => Ok(Arc::new(Self::PackedLeaf(
-                packed_leaf.update(prefix, hash, updates)?,
+                packed_leaf.update(prefix, updates)?,
             ))),
             Self::Node { left, right, .. } if depth > 0 => {
                 let packing_depth = opt_packing_depth::<T>().unwrap_or(0);
@@ -217,7 +217,7 @@ impl<T: TreeHash + Clone> Tree<T> {
             Self::Zero(zero_depth) if *zero_depth == depth => {
                 if depth == 0 {
                     if opt_packing_factor::<T>().is_some() {
-                        let packed_leaf = PackedLeaf::empty().update(prefix, hash, updates)?;
+                        let packed_leaf = PackedLeaf::empty().update(prefix, updates)?;
                         Ok(Arc::new(Self::PackedLeaf(packed_leaf)))
                     } else {
                         let index = prefix;
@@ -270,8 +270,7 @@ impl<T: PartialEq + TreeHash + Clone + Encode + Decode> Tree<T> {
                     }
                 }
                 if !equal {
-                    let hash = *l2.hash.read();
-                    diff.hashes.insert((depth, prefix), hash);
+                    diff.hashes.insert((depth, prefix), l2.tree_hash());
                 }
                 Ok(())
             }
@@ -342,8 +341,7 @@ impl<T: PartialEq + TreeHash + Clone + Encode + Decode> Tree<T> {
                 Ok(())
             }
             Self::PackedLeaf(packed_leaf) if depth == 0 => {
-                diff.hashes
-                    .insert((depth, prefix), *packed_leaf.hash.read());
+                diff.hashes.insert((depth, prefix), packed_leaf.tree_hash());
                 for (i, value) in packed_leaf.values.iter().enumerate() {
                     diff.leaves.insert(prefix | i, value.clone());
                 }