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 PackedLeaf { - 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 Tree { 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 Tree { 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( + left: &Arc>, + right: &Arc>, + ) -> 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 } }