Use OnceLock in place of RwLock

Cargo.toml

@@ -21,7 +21,6 @@ ethereum_hashing = "0.8"
 ethereum_ssz = "0.10"
 ethereum_ssz_derive = "0.10"
 itertools = "0.13.0"
-parking_lot = "0.12.1"
 rayon = "1.5.1"
 serde = { version = "1.0.0", features = ["derive"] }
 tree_hash = "0.12"

src/leaf.rs

@@ -1,15 +1,15 @@
 use crate::Arc;
 use educe::Educe;
-use parking_lot::RwLock;
+use std::sync::OnceLock;
 use tree_hash::Hash256;
 
 #[derive(Debug, Educe)]
 #[cfg_attr(feature = "arbitrary", derive(arbitrary::Arbitrary))]
 #[educe(PartialEq, Hash)]
 pub struct Leaf<T> {
     #[educe(PartialEq(ignore), Hash(ignore))]
-    #[cfg_attr(feature = "arbitrary", arbitrary(with = crate::utils::arb_rwlock))]
-    pub hash: RwLock<Hash256>,
+    #[cfg_attr(feature = "arbitrary", arbitrary(with = crate::utils::arb_oncelock))]
+    pub hash: OnceLock<Hash256>,
     #[cfg_attr(feature = "arbitrary", arbitrary(with = crate::utils::arb_arc))]
     pub value: Arc<T>,
 }
@@ -20,20 +20,26 @@ where
 {
     fn clone(&self) -> Self {
         Self {
-            hash: RwLock::new(*self.hash.read()),
+            hash: match self.hash.get() {
+                Some(&h) => OnceLock::from(h),
+                None => OnceLock::new(),
+            },
             value: self.value.clone(),
         }
     }
 }
 
 impl<T> Leaf<T> {
     pub fn new(value: T) -> Self {
-        Self::with_hash(value, Hash256::ZERO)
+        Self::with_hash(value, None)
     }
 
-    pub fn with_hash(value: T, hash: Hash256) -> Self {
+    pub fn with_hash(value: T, hash: Option<Hash256>) -> Self {
         Self {
-            hash: RwLock::new(hash),
+            hash: match hash {
+                Some(h) => OnceLock::from(h),
+                None => OnceLock::new(),
+            },
             value: Arc::new(value),
         }
     }

src/packed_leaf.rs

@@ -1,16 +1,16 @@
 use crate::{Error, UpdateMap};
 use educe::Educe;
-use parking_lot::RwLock;
 use std::ops::ControlFlow;
+use std::sync::OnceLock;
 use tree_hash::{BYTES_PER_CHUNK, Hash256, TreeHash};
 
 #[derive(Debug, Educe)]
 #[cfg_attr(feature = "arbitrary", derive(arbitrary::Arbitrary))]
 #[educe(PartialEq, Hash)]
 pub struct PackedLeaf<T: TreeHash + Clone> {
     #[educe(PartialEq(ignore), Hash(ignore))]
-    #[cfg_attr(feature = "arbitrary", arbitrary(with = crate::utils::arb_rwlock))]
-    pub hash: RwLock<Hash256>,
+    #[cfg_attr(feature = "arbitrary", arbitrary(with = crate::utils::arb_oncelock))]
+    pub hash: OnceLock<Hash256>,
     pub values: Vec<T>,
 }
 
@@ -20,37 +20,34 @@ where
 {
     fn clone(&self) -> Self {
         Self {
-            hash: RwLock::new(*self.hash.read()),
+            hash: match self.hash.get() {
+                Some(&h) => OnceLock::from(h),
+                None => OnceLock::new(),
+            },
             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;
+        if let Some(&cached) = self.hash.get() {
+            return cached;
         }
-
+        let mut hash = Hash256::ZERO;
         let hash_bytes = hash.as_mut_slice();
-
         let value_len = BYTES_PER_CHUNK / T::tree_hash_packing_factor();
         for (i, value) in self.values.iter().enumerate() {
             hash_bytes[i * value_len..(i + 1) * value_len]
                 .copy_from_slice(&value.tree_hash_packed_encoding());
         }
-
-        *self.hash.write() = hash;
+        let _ = self.hash.set(hash);
         hash
     }
 
     pub fn empty() -> Self {
         PackedLeaf {
-            hash: RwLock::new(Hash256::ZERO),
+            hash: OnceLock::new(),
             values: Vec::with_capacity(T::tree_hash_packing_factor()),
         }
     }
@@ -60,22 +57,22 @@ impl<T: TreeHash + Clone> PackedLeaf<T> {
         values.push(value);
 
         PackedLeaf {
-            hash: RwLock::new(Hash256::ZERO),
+            hash: OnceLock::new(),
             values,
         }
     }
 
     pub fn repeat(value: T, n: usize) -> Self {
         assert!(n <= T::tree_hash_packing_factor());
         PackedLeaf {
-            hash: RwLock::new(Hash256::ZERO),
+            hash: OnceLock::new(),
             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),
+            hash: OnceLock::new(),
             values: self.values.clone(),
         };
         let sub_index = index % T::tree_hash_packing_factor();
@@ -86,11 +83,14 @@ impl<T: TreeHash + Clone> PackedLeaf<T> {
     pub fn update<U: UpdateMap<T>>(
         &self,
         prefix: usize,
-        hash: Hash256,
+        hash: Option<Hash256>,
         updates: &U,
     ) -> Result<Self, Error> {
         let mut updated = PackedLeaf {
-            hash: RwLock::new(hash),
+            hash: match hash {
+                Some(h) => OnceLock::from(h),
+                None => OnceLock::new(),
+            },
             values: self.values.clone(),
         };
 
@@ -104,8 +104,8 @@ 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;
+        // Ensure hash is cleared.
+        self.hash = OnceLock::new();
 
         if sub_index == self.values.len() {
             self.values.push(value);

src/repeat.rs

@@ -1,7 +1,6 @@
 use crate::utils::{Length, opt_packing_factor};
 use crate::{Arc, Error, Leaf, List, PackedLeaf, Tree, UpdateMap, Value};
 use smallvec::{SmallVec, smallvec};
-use tree_hash::Hash256;
 use typenum::Unsigned;
 
 /// Efficiently construct a list from `n` copies of `elem`.
@@ -44,55 +43,46 @@ where
     for depth in 0..tree_depth {
         let new_layer = match &layer[..] {
             [(repeat_leaf, 1)] => {
-                smallvec![(
-                    Tree::node(repeat_leaf.clone(), Tree::zero(depth), Hash256::ZERO),
-                    1,
-                )]
+                smallvec![(Tree::node(repeat_leaf.clone(), Tree::zero(depth), None), 1,)]
             }
             [(repeat_leaf, repeat_count)] if repeat_count.is_multiple_of(2) => {
                 smallvec![(
-                    Tree::node(repeat_leaf.clone(), repeat_leaf.clone(), Hash256::ZERO),
+                    Tree::node(repeat_leaf.clone(), repeat_leaf.clone(), None),
                     repeat_count / 2,
                 )]
             }
             [(repeat_leaf, repeat_count)] => {
                 smallvec![
                     (
-                        Tree::node(repeat_leaf.clone(), repeat_leaf.clone(), Hash256::ZERO),
+                        Tree::node(repeat_leaf.clone(), repeat_leaf.clone(), None),
                         repeat_count / 2,
                     ),
-                    (
-                        Tree::node(repeat_leaf.clone(), Tree::zero(depth), Hash256::ZERO),
-                        1,
-                    ),
+                    (Tree::node(repeat_leaf.clone(), Tree::zero(depth), None), 1,),
                 ]
             }
             [(repeat_leaf, 1), (lonely_leaf, 1)] => {
                 smallvec![(
-                    Tree::node(repeat_leaf.clone(), lonely_leaf.clone(), Hash256::ZERO),
+                    Tree::node(repeat_leaf.clone(), lonely_leaf.clone(), None),
                     1,
                 )]
             }
             [(repeat_leaf, repeat_count), (lonely_leaf, 1)] => {
                 if repeat_count.is_multiple_of(2) {
                     smallvec![
                         (
-                            Tree::node(repeat_leaf.clone(), repeat_leaf.clone(), Hash256::ZERO),
+                            Tree::node(repeat_leaf.clone(), repeat_leaf.clone(), None),
                             repeat_count / 2,
                         ),
-                        (
-                            Tree::node(lonely_leaf.clone(), Tree::zero(depth), Hash256::ZERO),
-                            1,
-                        ),
+                        (Tree::node(lonely_leaf.clone(), Tree::zero(depth), None), 1,),
                     ]
                 } else {
                     smallvec![
                         (
-                            Tree::node(repeat_leaf.clone(), repeat_leaf.clone(), Hash256::ZERO),
+                            Tree::node(repeat_leaf.clone(), repeat_leaf.clone(), None),
                             repeat_count / 2,
                         ),
                         (
-                            Tree::node(repeat_leaf.clone(), lonely_leaf.clone(), Hash256::ZERO),
+                            Tree::node(repeat_leaf.clone(), lonely_leaf.clone(), None),
                             1,
                         ),
                     ]

src/tests/size_of.rs

@@ -1,6 +1,6 @@
 use crate::{Arc, Leaf, PackedLeaf, Tree};
-use parking_lot::RwLock;
 use std::mem::size_of;
+use std::sync::OnceLock;
 use tree_hash::Hash256;
 
 /// It's important that the Tree nodes have a predictable size.
@@ -10,8 +10,8 @@ fn size_of_hash256() {
     assert_eq!(size_of::<Leaf<Hash256>>(), 48);
     assert_eq!(size_of::<PackedLeaf<Hash256>>(), 64);
 
-    let rw_lock_size = size_of::<RwLock<Hash256>>();
-    assert_eq!(rw_lock_size, 40);
+    let once_lock_size = size_of::<OnceLock<Hash256>>();
+    assert_eq!(once_lock_size, 36);
 
     let arc_size = size_of::<Arc<Tree<Hash256>>>();
     assert_eq!(arc_size, 8);
@@ -27,11 +27,11 @@ fn size_of_u8() {
     assert_eq!(size_of::<PackedLeaf<u8>>(), 64);
     assert_eq!(
         size_of::<PackedLeaf<u8>>(),
-        size_of::<RwLock<Hash256>>() + size_of::<Vec<u8>>()
+        size_of::<OnceLock<Hash256>>() + size_of::<Vec<u8>>() + 4 // 4 bytes alignment padding
     );
 
-    let rw_lock_size = size_of::<RwLock<u8>>();
-    assert_eq!(rw_lock_size, 16);
+    let once_lock_size = size_of::<OnceLock<u8>>();
+    assert_eq!(once_lock_size, 8);
 
     let arc_size = size_of::<Arc<Tree<u8>>>();
     assert_eq!(arc_size, 8);