From 5ce1f42831d7509da1a13d20d18c3ce12367fdfd Mon Sep 17 00:00:00 2001
From: Mac L <mjladson@pm.me>
Date: Tue, 5 Sep 2023 01:32:14 +1000
Subject: [PATCH] Use write locks to prevent work duplication

---
 src/packed_leaf.rs | 26 ++++++++++++++------------
 src/tree.rs        | 31 +++++++++++++++++--------------
 2 files changed, 31 insertions(+), 26 deletions(-)

diff --git a/src/packed_leaf.rs b/src/packed_leaf.rs
index b0f3770..68940ad 100644
--- a/src/packed_leaf.rs
+++ b/src/packed_leaf.rs
@@ -27,27 +27,29 @@ where
 }
 
 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;
-        }
-
+    fn compute_hash(&self, mut hash: Hash256) -> Hash256 {
         let hash_bytes = hash.as_bytes_mut();
-
         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;
         hash
     }
 
+    pub fn tree_hash(&self) -> Hash256 {
+        let mut write_lock = self.hash.write();
+        let hash = *write_lock;
+
+        if !hash.is_zero() {
+            hash
+        } else {
+            let tree_hash = self.compute_hash(hash);
+            *write_lock = tree_hash;
+            tree_hash
+        }
+    }
+
     pub fn empty() -> Self {
         PackedLeaf {
             hash: RwLock::new(Hash256::zero()),
diff --git a/src/tree.rs b/src/tree.rs
index d816e6e..2c88faa 100644
--- a/src/tree.rs
+++ b/src/tree.rs
@@ -531,10 +531,8 @@ impl<T: Value + Send + Sync> Tree<T> {
     pub fn tree_hash(&self) -> Hash256 {
         match self {
             Self::Leaf(Leaf { hash, value }) => {
-                // FIXME(sproul): upgradeable RwLock?
-                let read_lock = hash.read();
-                let existing_hash = *read_lock;
-                drop(read_lock);
+                let mut write_lock = hash.write();
+                let existing_hash = *write_lock;
 
                 // NOTE: We re-compute the hash whenever it is non-zero. Computed hashes may
                 // legitimately be zero, but this only occurs at the leaf level when the value is
@@ -547,26 +545,31 @@ impl<T: Value + Send + Sync> Tree<T> {
                     existing_hash
                 } else {
                     let tree_hash = value.tree_hash_root();
-                    *hash.write() = tree_hash;
+                    *write_lock = tree_hash;
                     tree_hash
                 }
             }
             Self::PackedLeaf(leaf) => leaf.tree_hash(),
             Self::Zero(depth) => Hash256::from_slice(&ZERO_HASHES[*depth]),
             Self::Node { hash, left, right } => {
-                let read_lock = hash.read();
-                let existing_hash = *read_lock;
-                drop(read_lock);
+                fn node_tree_hash<T: Value + Send + Sync>(
+                    left: &Arc<Tree<T>>,
+                    right: &Arc<Tree<T>>,
+                ) -> Hash256 {
+                    let (left_hash, right_hash) =
+                        rayon::join(|| left.tree_hash(), || right.tree_hash());
+                    Hash256::from(hash32_concat(left_hash.as_bytes(), right_hash.as_bytes()))
+                }
+
+                let mut write_lock = hash.write();
+                let existing_hash = *write_lock;
 
                 if !existing_hash.is_zero() {
                     existing_hash
                 } else {
-                    // Parallelism goes brrrr.
-                    let (left_hash, right_hash) =
-                        rayon::join(|| left.tree_hash(), || right.tree_hash());
-                    let tree_hash =
-                        Hash256::from(hash32_concat(left_hash.as_bytes(), right_hash.as_bytes()));
-                    *hash.write() = tree_hash;
+                    let tree_hash = node_tree_hash(left, right);
+
+                    *write_lock = tree_hash;
                     tree_hash
                 }
             }