WIP: interior mutability

starkware-libs · Apr 16, 2024 · 2e166d8 · 2e166d8
1 parent 5341978
commit 2e166d8
Showing 1 changed file with 43 additions and 35 deletions.
diff --git a/crates/committer/src/patricia_merkle_tree/updated_skeleton_tree.rs b/crates/committer/src/patricia_merkle_tree/updated_skeleton_tree.rs
@@ -31,7 +31,6 @@ struct UpdatedSkeletonTreeImpl<L: LeafDataTrait + std::clone::Clone> {
     skeleton_tree: HashMap<NodeIndex, UpdatedSkeletonNode<L>>,
 }
 
-#[allow(dead_code)]
 impl<L: LeafDataTrait + std::clone::Clone + std::marker::Sync + std::marker::Send>
     UpdatedSkeletonTreeImpl<L>
 {
@@ -45,6 +44,29 @@ impl<L: LeafDataTrait + std::clone::Clone + std::marker::Sync + std::marker::Sen
         }
     }
 
+    fn write_to_output_map_with_interior_mutuability(
+        output_map: Arc<RwLock<HashMap<NodeIndex, Mutex<FilledNode<L>>>>>,
+        index: NodeIndex,
+        hash: HashOutput,
+        data: NodeData<L>,
+    ) -> Result<(), UpdatedSkeletonTreeError> {
+        let read_map = output_map.read().map_err(|_| {
+            UpdatedSkeletonTreeError::PoisonedLock("Cannot read from output map.".to_owned())
+        })?;
+
+        match read_map.get(&index) {
+            Some(node) => {
+                let mut node = node.lock().map_err(|_| {
+                    UpdatedSkeletonTreeError::PoisonedLock("Cannot lock node.".to_owned())
+                })?;
+                node.hash = hash;
+                node.data = data;
+                Ok(())
+            }
+            None => Err(UpdatedSkeletonTreeError::MissingNode),
+        }
+    }
+
     fn compute_filled_tree_rec<H: HashFunction, TH: TreeHashFunction<L, H>>(
         &self,
         index: NodeIndex,
@@ -78,17 +100,12 @@ impl<L: LeafDataTrait + std::clone::Clone + std::marker::Sync + std::marker::Sen
                     right_hash,
                 });
                 let hash_value = TH::compute_node_hash(&data);
-                // TODO (Aner, 15/4/21): Change to use interior mutability.
-                let mut write_locked_map = output_map.write().map_err(|_| {
-                    UpdatedSkeletonTreeError::PoisonedLock("Cannot write to output map.".to_owned())
-                })?;
-                write_locked_map.insert(
-                    index,
-                    Mutex::new(FilledNode {
-                        hash: hash_value,
-                        data,
-                    }),
-                );
+                // TODO (Aner, 15/4/21): Change to use interior mutability. Doing so will require
+                // inserting all nodes as "unknown nodes" in advance, so they are only edited
+                // here.
+                Self::write_to_output_map_with_interior_mutuability(
+                    output_map, index, hash_value, data,
+                )?;
                 Ok(hash_value)
             }
             UpdatedSkeletonNode::Edge { path_to_bottom } => {
@@ -100,33 +117,24 @@ impl<L: LeafDataTrait + std::clone::Clone + std::marker::Sync + std::marker::Sen
                     bottom_hash,
                 });
                 let hash_value = TH::compute_node_hash(&data);
-                // TODO (Aner, 15/4/21): Change to use interior mutability.
-                let mut write_locked_map = output_map.write().map_err(|_| {
-                    UpdatedSkeletonTreeError::PoisonedLock("Cannot write to output map.".to_owned())
-                })?;
-                write_locked_map.insert(
-                    index,
-                    Mutex::new(FilledNode {
-                        hash: hash_value,
-                        data,
-                    }),
-                );
+                // TODO (Aner, 15/4/21): Change to use interior mutability. Doing so will require
+                // inserting all nodes as "unknown nodes" in advance, so they are only edited
+                // here.
+                Self::write_to_output_map_with_interior_mutuability(
+                    output_map, index, hash_value, data,
+                )?;
                 Ok(hash_value)
             }
             UpdatedSkeletonNode::Sibling(hash_result) => Ok(*hash_result),
             UpdatedSkeletonNode::Leaf(node_data) => {
-                let hash_value = TH::compute_node_hash(&NodeData::Leaf(node_data.clone()));
-                // TODO (Aner, 15/4/21): Change to use interior mutability.
-                let mut write_locked_map = output_map.write().map_err(|_| {
-                    UpdatedSkeletonTreeError::PoisonedLock("Cannot write to output map.".to_owned())
-                })?;
-                write_locked_map.insert(
-                    index,
-                    Mutex::new(FilledNode {
-                        hash: hash_value,
-                        data: NodeData::Leaf(node_data.clone()),
-                    }),
-                );
+                let data = NodeData::Leaf(node_data.clone());
+                let hash_value = TH::compute_node_hash(&data);
+                // TODO (Aner, 15/4/21): Change to use interior mutability. Doing so will require
+                // inserting all nodes as "unknown nodes" in advance, so they are only edited
+                // here.
+                Self::write_to_output_map_with_interior_mutuability(
+                    output_map, index, hash_value, data,
+                )?;
                 Ok(hash_value)
             }
         }