From bd49d472689554433b6b49f1d952fee50da29605 Mon Sep 17 00:00:00 2001
From: Pepper Lebeck-Jobe <eljobe@gmail.com>
Date: Tue, 21 May 2024 12:52:01 +0200
Subject: [PATCH] Remove deadlocks.

---
 arbitrator/prover/src/merkle.rs | 51 ++++++++++++++++-----------------
 1 file changed, 24 insertions(+), 27 deletions(-)

diff --git a/arbitrator/prover/src/merkle.rs b/arbitrator/prover/src/merkle.rs
index b37e0140ba..70884da096 100644
--- a/arbitrator/prover/src/merkle.rs
+++ b/arbitrator/prover/src/merkle.rs
@@ -6,6 +6,8 @@ use digest::Digest;
 
 use enum_iterator::Sequence;
 
+use parking_lot::Mutex;
+
 #[cfg(feature = "counters")]
 use enum_iterator::all;
 use itertools::Itertools;
@@ -30,7 +32,7 @@ use sha3::Keccak256;
 use std::{
     collections::HashSet,
     convert::{TryFrom, TryInto},
-    sync::{Arc, Mutex},
+    sync::Arc,
 };
 
 #[cfg(feature = "rayon")]
@@ -262,8 +264,7 @@ impl Merkle {
         }
     }
 
-    fn rehash(&self) {
-        let layers = &mut self.layers.lock().unwrap();
+    fn rehash(&self, layers: &mut Layers) {
         if layers.dirt.is_empty() || layers.dirt[0].is_empty() {
             return;
         }
@@ -294,8 +295,9 @@ impl Merkle {
     pub fn root(&self) -> Bytes32 {
         #[cfg(feature = "counters")]
         ROOT_COUNTERS[&self.ty].fetch_add(1, Ordering::Relaxed);
-        self.rehash();
-        if let Some(layer) = self.layers.lock().unwrap().data.last() {
+        let mut layers = self.layers.lock();
+        self.rehash(&mut layers);
+        if let Some(layer) = layers.data.last() {
             assert_eq!(layer.len(), 1);
             layer[0]
         } else {
@@ -306,7 +308,7 @@ impl Merkle {
     // Returns the total number of leaves the tree can hold.
     #[inline]
     fn capacity(&self) -> usize {
-        let layers = self.layers.lock().unwrap();
+        let layers = self.layers.lock();
         if layers.data.is_empty() {
             return 0;
         }
@@ -316,18 +318,17 @@ impl Merkle {
 
     // Returns the number of leaves in the tree.
     pub fn len(&self) -> usize {
-        self.layers.lock().unwrap().data[0].len()
+        self.layers.lock().data[0].len()
     }
 
     pub fn is_empty(&self) -> bool {
-        let layers = self.layers.lock().unwrap();
+        let layers = self.layers.lock();
         layers.data.is_empty() || layers.data[0].is_empty()
     }
 
     #[must_use]
     pub fn prove(&self, idx: usize) -> Option<Vec<u8>> {
-        let layers = self.layers.lock().unwrap();
-        if layers.data.is_empty() || idx >= layers.data[0].len() {
+        if self.layers.lock().data.is_empty() || idx >= self.layers.lock().data[0].len() {
             return None;
         }
         Some(self.prove_any(idx))
@@ -336,8 +337,8 @@ impl Merkle {
     /// creates a merkle proof regardless of if the leaf has content
     #[must_use]
     pub fn prove_any(&self, mut idx: usize) -> Vec<u8> {
-        self.rehash();
-        let layers = self.layers.lock().unwrap();
+        let mut layers = self.layers.lock();
+        self.rehash(&mut layers);
         let mut proof = vec![u8::try_from(layers.data.len() - 1).unwrap()];
         for (layer_i, layer) in layers.data.iter().enumerate() {
             if layer_i == layers.data.len() - 1 {
@@ -358,7 +359,7 @@ impl Merkle {
     /// Adds a new leaf to the merkle
     /// Currently O(n) in the number of leaves (could be log(n))
     pub fn push_leaf(&mut self, leaf: Bytes32) {
-        let mut leaves = self.layers.lock().unwrap().data.swap_remove(0);
+        let mut leaves = self.layers.lock().data.swap_remove(0);
         leaves.push(leaf);
         *self = Self::new_advanced(self.ty, leaves, self.min_depth);
     }
@@ -366,7 +367,7 @@ impl Merkle {
     /// Removes the rightmost leaf from the merkle
     /// Currently O(n) in the number of leaves (could be log(n))
     pub fn pop_leaf(&mut self) {
-        let mut leaves = self.layers.lock().unwrap().data.swap_remove(0);
+        let mut leaves = self.layers.lock().data.swap_remove(0);
         leaves.pop();
         *self = Self::new_advanced(self.ty, leaves, self.min_depth);
     }
@@ -376,7 +377,7 @@ impl Merkle {
     pub fn set(&self, idx: usize, hash: Bytes32) {
         #[cfg(feature = "counters")]
         SET_COUNTERS[&self.ty].fetch_add(1, Ordering::Relaxed);
-        let mut layers = self.layers.lock().unwrap();
+        let mut layers = self.layers.lock();
         if layers.data[0][idx] == hash {
             return;
         }
@@ -393,7 +394,7 @@ impl Merkle {
                 "Cannot resize to a length greater than the capacity of the tree.".to_owned(),
             );
         }
-        let mut layers = self.layers.lock().unwrap();
+        let mut layers = self.layers.lock();
         let mut layer_size = new_len;
         for (layer_i, layer) in layers.data.iter_mut().enumerate() {
             layer.resize(layer_size, *empty_hash_at(self.ty, layer_i));
@@ -416,27 +417,23 @@ impl PartialEq for Merkle {
 impl Eq for Merkle {}
 
 pub mod arc_mutex_sedre {
-    pub fn serialize<S, T>(
-        data: &std::sync::Arc<std::sync::Mutex<T>>,
-        serializer: S,
-    ) -> Result<S::Ok, S::Error>
+    use parking_lot::Mutex;
+    use std::sync::Arc;
+
+    pub fn serialize<S, T>(data: &Arc<Mutex<T>>, serializer: S) -> Result<S::Ok, S::Error>
     where
         S: serde::Serializer,
         T: serde::Serialize,
     {
-        data.lock().unwrap().serialize(serializer)
+        data.lock().serialize(serializer)
     }
 
-    pub fn deserialize<'de, D, T>(
-        deserializer: D,
-    ) -> Result<std::sync::Arc<std::sync::Mutex<T>>, D::Error>
+    pub fn deserialize<'de, D, T>(deserializer: D) -> Result<Arc<Mutex<T>>, D::Error>
     where
         D: serde::Deserializer<'de>,
         T: serde::Deserialize<'de>,
     {
-        Ok(std::sync::Arc::new(std::sync::Mutex::new(T::deserialize(
-            deserializer,
-        )?)))
+        Ok(Arc::new(Mutex::new(T::deserialize(deserializer)?)))
     }
 }