Avoid the double lock around `EncoderState`

2 files changed, 18 insertions, 17 deletions

diff --git a/compiler/rustc_query_system/src/dep_graph/graph.rs b/compiler/rustc_query_system/src/dep_graph/graph.rs
index 03e435506ab..9f067273f35 100644
--- a/compiler/rustc_query_system/src/dep_graph/graph.rs
+++ b/compiler/rustc_query_system/src/dep_graph/graph.rs
@@ -3,7 +3,6 @@ use rustc_data_structures::fx::{FxHashMap, FxHashSet};
 use rustc_data_structures::profiling::{QueryInvocationId, SelfProfilerRef};
 use rustc_data_structures::sharded::{self, Sharded};
 use rustc_data_structures::stable_hasher::{HashStable, StableHasher};
-use rustc_data_structures::steal::Steal;
 use rustc_data_structures::sync::{AtomicU32, AtomicU64, Lock, Lrc};
 use rustc_data_structures::unord::UnordMap;
 use rustc_index::IndexVec;
@@ -194,7 +193,7 @@ impl<D: Deps> DepGraph<D> {
 
     pub fn with_query(&self, f: impl Fn(&DepGraphQuery)) {
         if let Some(data) = &self.data {
-            data.current.encoder.borrow().with_query(f)
+            data.current.encoder.with_query(f)
         }
     }
 
@@ -954,7 +953,7 @@ impl<D: Deps> DepGraph<D> {
 
     pub fn print_incremental_info(&self) {
         if let Some(data) = &self.data {
-            data.current.encoder.borrow().print_incremental_info(
+            data.current.encoder.print_incremental_info(
                 data.current.total_read_count.load(Ordering::Relaxed),
                 data.current.total_duplicate_read_count.load(Ordering::Relaxed),
             )
@@ -962,7 +961,7 @@ impl<D: Deps> DepGraph<D> {
     }
 
     pub fn finish_encoding(&self) -> FileEncodeResult {
-        if let Some(data) = &self.data { data.current.encoder.steal().finish() } else { Ok(0) }
+        if let Some(data) = &self.data { data.current.encoder.finish() } else { Ok(0) }
     }
 
     pub(crate) fn next_virtual_depnode_index(&self) -> DepNodeIndex {
@@ -1045,7 +1044,7 @@ rustc_index::newtype_index! {
 /// manipulating both, we acquire `new_node_to_index` or `prev_index_to_index`
 /// first, and `data` second.
 pub(super) struct CurrentDepGraph<D: Deps> {
-    encoder: Steal<GraphEncoder<D>>,
+    encoder: GraphEncoder<D>,
     new_node_to_index: Sharded<FxHashMap<DepNode, DepNodeIndex>>,
     prev_index_to_index: Lock<IndexVec<SerializedDepNodeIndex, Option<DepNodeIndex>>>,
 
@@ -1111,13 +1110,13 @@ impl<D: Deps> CurrentDepGraph<D> {
         let new_node_count_estimate = 102 * prev_graph_node_count / 100 + 200;
 
         CurrentDepGraph {
-            encoder: Steal::new(GraphEncoder::new(
+            encoder: GraphEncoder::new(
                 encoder,
                 prev_graph_node_count,
                 record_graph,
                 record_stats,
                 profiler,
-            )),
+            ),
             new_node_to_index: Sharded::new(|| {
                 FxHashMap::with_capacity_and_hasher(
                     new_node_count_estimate / sharded::shards(),
@@ -1156,7 +1155,7 @@ impl<D: Deps> CurrentDepGraph<D> {
         let dep_node_index = match self.new_node_to_index.lock_shard_by_value(&key).entry(key) {
             Entry::Occupied(entry) => *entry.get(),
             Entry::Vacant(entry) => {
-                let dep_node_index = self.encoder.borrow().send(key, current_fingerprint, edges);
+                let dep_node_index = self.encoder.send(key, current_fingerprint, edges);
                 entry.insert(dep_node_index);
                 dep_node_index
             }
@@ -1182,7 +1181,7 @@ impl<D: Deps> CurrentDepGraph<D> {
                 let dep_node_index = match prev_index_to_index[prev_index] {
                     Some(dep_node_index) => dep_node_index,
                     None => {
-                        let dep_node_index = self.encoder.borrow().send(key, fingerprint, edges);
+                        let dep_node_index = self.encoder.send(key, fingerprint, edges);
                         prev_index_to_index[prev_index] = Some(dep_node_index);
                         dep_node_index
                     }
@@ -1243,7 +1242,7 @@ impl<D: Deps> CurrentDepGraph<D> {
                     .map(|i| prev_index_to_index[i].unwrap())
                     .collect();
                 let fingerprint = prev_graph.fingerprint_by_index(prev_index);
-                let dep_node_index = self.encoder.borrow().send(key, fingerprint, edges);
+                let dep_node_index = self.encoder.send(key, fingerprint, edges);
                 prev_index_to_index[prev_index] = Some(dep_node_index);
                 #[cfg(debug_assertions)]
                 self.record_edge(dep_node_index, key, fingerprint);
diff --git a/compiler/rustc_query_system/src/dep_graph/serialized.rs b/compiler/rustc_query_system/src/dep_graph/serialized.rs
index d8db94defbd..cc823917ce8 100644
--- a/compiler/rustc_query_system/src/dep_graph/serialized.rs
+++ b/compiler/rustc_query_system/src/dep_graph/serialized.rs
@@ -505,7 +505,7 @@ impl<D: Deps> EncoderState<D> {
 
 pub struct GraphEncoder<D: Deps> {
     profiler: SelfProfilerRef,
-    status: Lock<EncoderState<D>>,
+    status: Lock<Option<EncoderState<D>>>,
     record_graph: Option<Lock<DepGraphQuery>>,
 }
 
@@ -518,7 +518,7 @@ impl<D: Deps> GraphEncoder<D> {
         profiler: &SelfProfilerRef,
     ) -> Self {
         let record_graph = record_graph.then(|| Lock::new(DepGraphQuery::new(prev_node_count)));
-        let status = Lock::new(EncoderState::new(encoder, record_stats));
+        let status = Lock::new(Some(EncoderState::new(encoder, record_stats)));
         GraphEncoder { status, record_graph, profiler: profiler.clone() }
     }
 
@@ -533,7 +533,8 @@ impl<D: Deps> GraphEncoder<D> {
         total_read_count: u64,
         total_duplicate_read_count: u64,
     ) {
-        let status = self.status.lock();
+        let mut status = self.status.lock();
+        let status = status.as_mut().unwrap();
         if let Some(record_stats) = &status.stats {
             let mut stats: Vec<_> = record_stats.values().collect();
             stats.sort_by_key(|s| -(s.node_counter as i64));
@@ -588,11 +589,12 @@ impl<D: Deps> GraphEncoder<D> {
     ) -> DepNodeIndex {
         let _prof_timer = self.profiler.generic_activity("incr_comp_encode_dep_graph");
         let node = NodeInfo { node, fingerprint, edges };
-        self.status.lock().encode_node(&node, &self.record_graph)
+        self.status.lock().as_mut().unwrap().encode_node(&node, &self.record_graph)
     }
 
-    pub fn finish(self) -> FileEncodeResult {
-        let _prof_timer = self.profiler.generic_activity("incr_comp_encode_dep_graph");
-        self.status.into_inner().finish(&self.profiler)
+    pub fn finish(&self) -> FileEncodeResult {
+        let _prof_timer = self.profiler.generic_activity("incr_comp_encode_dep_graph_finish");
+
+        self.status.lock().take().unwrap().finish(&self.profiler)
     }
 }