use with_lock tp avoid drop

Author: SparrowLii

compiler/rustc_query_system/src/cache.rs

@@ -26,17 +26,17 @@ impl<Key, Value> Default for Cache<Key, Value> {
 impl<Key, Value> Cache<Key, Value> {
     /// Actually frees the underlying memory in contrast to what stdlib containers do on `clear`
     pub fn clear(&self) {
-        *self.hashmap.borrow_mut() = Default::default();
+        self.hashmap.with_lock(|map| *map = Default::default());
     }
 }
 
 impl<Key: Eq + Hash, Value: Clone> Cache<Key, Value> {
     pub fn get<Tcx: DepContext>(&self, key: &Key, tcx: Tcx) -> Option<Value> {
-        Some(self.hashmap.borrow().get(key)?.get(tcx))
+        self.hashmap.with_borrow(|map| map.get(key).map(|node| node.get(tcx)))
     }
 
     pub fn insert(&self, key: Key, dep_node: DepNodeIndex, value: Value) {
-        self.hashmap.borrow_mut().insert(key, WithDepNode::new(dep_node, value));
+        self.hashmap.with_lock(|map| map.insert(key, WithDepNode::new(dep_node, value)));
     }
 }
 

compiler/rustc_query_system/src/dep_graph/graph.rs

@@ -2,7 +2,7 @@ use parking_lot::Mutex;
 use rustc_data_structures::fingerprint::Fingerprint;
 use rustc_data_structures::fx::{FxHashMap, FxHashSet};
 use rustc_data_structures::profiling::{EventId, QueryInvocationId, SelfProfilerRef};
-use rustc_data_structures::sharded::{self, Sharded};
+use rustc_data_structures::sharded::Sharded;
 use rustc_data_structures::stable_hasher::{HashStable, StableHasher};
 use rustc_data_structures::steal::Steal;
 use rustc_data_structures::sync::{AtomicU32, AtomicU64, Lock, Lrc, Ordering};
@@ -466,8 +466,42 @@ impl<K: DepKind> DepGraph<K> {
     pub fn read_index(&self, dep_node_index: DepNodeIndex) {
         if let Some(ref data) = self.data {
             K::read_deps(|task_deps| {
-                let mut task_deps = match task_deps {
-                    TaskDepsRef::Allow(deps) => deps.lock(),
+                match task_deps {
+                    TaskDepsRef::Allow(deps) => deps.with_lock(|task_deps| {
+                        // As long as we only have a low number of reads we can avoid doing a hash
+                        // insert and potentially allocating/reallocating the hashmap
+                        let new_read = if task_deps.reads.len() < TASK_DEPS_READS_CAP {
+                            task_deps.reads.iter().all(|other| *other != dep_node_index)
+                        } else {
+                            task_deps.read_set.insert(dep_node_index)
+                        };
+                        if new_read {
+                            task_deps.reads.push(dep_node_index);
+                            if task_deps.reads.len() == TASK_DEPS_READS_CAP {
+                                // Fill `read_set` with what we have so far so we can use the hashset
+                                // next time
+                                task_deps.read_set.extend(task_deps.reads.iter().copied());
+                            }
+
+                            #[cfg(debug_assertions)]
+                            {
+                                if let Some(target) = task_deps.node {
+                                    if let Some(ref forbidden_edge) = data.current.forbidden_edge {
+                                        let src =
+                                            forbidden_edge.index_to_node.lock()[&dep_node_index];
+                                        if forbidden_edge.test(&src, &target) {
+                                            panic!(
+                                                "forbidden edge {:?} -> {:?} created",
+                                                src, target
+                                            )
+                                        }
+                                    }
+                                }
+                            }
+                        } else if cfg!(debug_assertions) {
+                            data.current.total_duplicate_read_count.fetch_add(1, Relaxed);
+                        }
+                    }),
                     TaskDepsRef::EvalAlways => {
                         // We don't need to record dependencies of eval_always
                         // queries. They are re-evaluated unconditionally anyway.
@@ -478,41 +512,6 @@ impl<K: DepKind> DepGraph<K> {
                         panic!("Illegal read of: {dep_node_index:?}")
                     }
                 };
-                let task_deps = &mut *task_deps;
-
-                if cfg!(debug_assertions) {
-                    data.current.total_read_count.fetch_add(1, Relaxed);
-                }
-
-                // As long as we only have a low number of reads we can avoid doing a hash
-                // insert and potentially allocating/reallocating the hashmap
-                let new_read = if task_deps.reads.len() < TASK_DEPS_READS_CAP {
-                    task_deps.reads.iter().all(|other| *other != dep_node_index)
-                } else {
-                    task_deps.read_set.insert(dep_node_index)
-                };
-                if new_read {
-                    task_deps.reads.push(dep_node_index);
-                    if task_deps.reads.len() == TASK_DEPS_READS_CAP {
-                        // Fill `read_set` with what we have so far so we can use the hashset
-                        // next time
-                        task_deps.read_set.extend(task_deps.reads.iter().copied());
-                    }
-
-                    #[cfg(debug_assertions)]
-                    {
-                        if let Some(target) = task_deps.node {
-                            if let Some(ref forbidden_edge) = data.current.forbidden_edge {
-                                let src = forbidden_edge.index_to_node.lock()[&dep_node_index];
-                                if forbidden_edge.test(&src, &target) {
-                                    panic!("forbidden edge {:?} -> {:?} created", src, target)
-                                }
-                            }
-                        }
-                    }
-                } else if cfg!(debug_assertions) {
-                    data.current.total_duplicate_read_count.fetch_add(1, Relaxed);
-                }
             })
         }
     }
@@ -574,7 +573,9 @@ impl<K: DepKind> DepGraph<K> {
 
             let mut edges = SmallVec::new();
             K::read_deps(|task_deps| match task_deps {
-                TaskDepsRef::Allow(deps) => edges.extend(deps.lock().reads.iter().copied()),
+                TaskDepsRef::Allow(deps) => {
+                    deps.with_borrow(|deps| edges.extend(deps.reads.iter().copied()))
+                }
                 TaskDepsRef::EvalAlways => {
                     edges.push(DepNodeIndex::FOREVER_RED_NODE);
                 }
@@ -627,14 +628,11 @@ impl<K: DepKind> DepGraphData<K> {
     #[inline]
     pub fn dep_node_index_of_opt(&self, dep_node: &DepNode<K>) -> Option<DepNodeIndex> {
         if let Some(prev_index) = self.previous.node_to_index_opt(dep_node) {
-            self.current.prev_index_to_index.lock()[prev_index]
+            self.current.prev_index_to_index.with_borrow(|nodes| nodes[prev_index])
         } else {
             self.current
                 .new_node_to_index
-                .get_shard_by_value(dep_node)
-                .lock()
-                .get(dep_node)
-                .copied()
+                .with_get_shard_by_value(dep_node, |node| node.get(dep_node).copied())
         }
     }
 
@@ -670,7 +668,7 @@ impl<K: DepKind> DepGraphData<K> {
     }
 
     pub fn mark_debug_loaded_from_disk(&self, dep_node: DepNode<K>) {
-        self.debug_loaded_from_disk.lock().insert(dep_node);
+        self.debug_loaded_from_disk.with_lock(|node| node.insert(dep_node));
     }
 }
 
@@ -693,7 +691,11 @@ impl<K: DepKind> DepGraph<K> {
     }
 
     pub fn debug_was_loaded_from_disk(&self, dep_node: DepNode<K>) -> bool {
-        self.data.as_ref().unwrap().debug_loaded_from_disk.lock().contains(&dep_node)
+        self.data
+            .as_ref()
+            .unwrap()
+            .debug_loaded_from_disk
+            .with_borrow(|node| node.contains(&dep_node))
     }
 
     #[inline(always)]
@@ -703,15 +705,15 @@ impl<K: DepKind> DepGraph<K> {
     {
         let dep_node_debug = &self.data.as_ref().unwrap().dep_node_debug;
 
-        if dep_node_debug.borrow().contains_key(&dep_node) {
+        if dep_node_debug.with_borrow(|node| node.contains_key(&dep_node)) {
             return;
         }
         let debug_str = self.with_ignore(debug_str_gen);
-        dep_node_debug.borrow_mut().insert(dep_node, debug_str);
+        dep_node_debug.with_lock(|node| node.insert(dep_node, debug_str));
     }
 
     pub fn dep_node_debug_str(&self, dep_node: DepNode<K>) -> Option<String> {
-        self.data.as_ref()?.dep_node_debug.borrow().get(&dep_node).cloned()
+        self.data.as_ref()?.dep_node_debug.with_borrow(|node| node.get(&dep_node).cloned())
     }
 
     fn node_color(&self, dep_node: &DepNode<K>) -> Option<DepNodeColor> {
@@ -1163,10 +1165,7 @@ impl<K: DepKind> CurrentDepGraph<K> {
                 record_stats,
             )),
             new_node_to_index: Sharded::new(|| {
-                FxHashMap::with_capacity_and_hasher(
-                    new_node_count_estimate / sharded::SHARDS,
-                    Default::default(),
-                )
+                FxHashMap::with_capacity_and_hasher(new_node_count_estimate, Default::default())
             }),
             prev_index_to_index: Lock::new(IndexVec::from_elem_n(None, prev_graph_node_count)),
             anon_id_seed,
@@ -1199,16 +1198,16 @@ impl<K: DepKind> CurrentDepGraph<K> {
         edges: EdgesVec,
         current_fingerprint: Fingerprint,
     ) -> DepNodeIndex {
-        let dep_node_index = match self.new_node_to_index.get_shard_by_value(&key).lock().entry(key)
-        {
-            Entry::Occupied(entry) => *entry.get(),
-            Entry::Vacant(entry) => {
-                let dep_node_index =
-                    self.encoder.borrow().send(profiler, key, current_fingerprint, edges);
-                entry.insert(dep_node_index);
-                dep_node_index
-            }
-        };
+        let dep_node_index =
+            self.new_node_to_index.with_get_shard_by_value(&key, |node| match node.entry(key) {
+                Entry::Occupied(entry) => *entry.get(),
+                Entry::Vacant(entry) => {
+                    let dep_node_index =
+                        self.encoder.borrow().send(profiler, key, current_fingerprint, edges);
+                    entry.insert(dep_node_index);
+                    dep_node_index
+                }
+            });
 
         #[cfg(debug_assertions)]
         self.record_edge(dep_node_index, key, current_fingerprint);
@@ -1298,25 +1297,26 @@ impl<K: DepKind> CurrentDepGraph<K> {
     ) -> DepNodeIndex {
         self.debug_assert_not_in_new_nodes(prev_graph, prev_index);
 
-        let mut prev_index_to_index = self.prev_index_to_index.lock();
-
-        match prev_index_to_index[prev_index] {
-            Some(dep_node_index) => dep_node_index,
-            None => {
-                let key = prev_graph.index_to_node(prev_index);
-                let edges = prev_graph
-                    .edge_targets_from(prev_index)
-                    .iter()
-                    .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(profiler, key, fingerprint, edges);
-                prev_index_to_index[prev_index] = Some(dep_node_index);
-                #[cfg(debug_assertions)]
-                self.record_edge(dep_node_index, key, fingerprint);
-                dep_node_index
+        self.prev_index_to_index.with_lock(|prev_index_to_index| {
+            match prev_index_to_index[prev_index] {
+                Some(dep_node_index) => dep_node_index,
+                None => {
+                    let key = prev_graph.index_to_node(prev_index);
+                    let edges = prev_graph
+                        .edge_targets_from(prev_index)
+                        .iter()
+                        .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(profiler, key, fingerprint, edges);
+                    prev_index_to_index[prev_index] = Some(dep_node_index);
+                    #[cfg(debug_assertions)]
+                    self.record_edge(dep_node_index, key, fingerprint);
+                    dep_node_index
+                }
             }
-        }
+        })
     }
 
     #[inline]
@@ -1327,7 +1327,7 @@ impl<K: DepKind> CurrentDepGraph<K> {
     ) {
         let node = &prev_graph.index_to_node(prev_index);
         debug_assert!(
-            !self.new_node_to_index.get_shard_by_value(node).lock().contains_key(node),
+            !self.new_node_to_index.with_get_shard_by_value(node, |lock| lock.contains_key(node)),
             "node from previous graph present in new node collection"
         );
     }

compiler/rustc_query_system/src/dep_graph/serialized.rs

@@ -186,7 +186,7 @@ impl<K: DepKind> EncoderState<K> {
             if let Some(record_graph) = &mut record_graph.try_lock() {
                 record_graph.push(index, node.node, &node.edges);
             }
-        }
+        };
 
         if let Some(stats) = &mut self.stats {
             let kind = node.node.kind;
@@ -242,7 +242,7 @@ impl<K: DepKind + Encodable<FileEncoder>> GraphEncoder<K> {
 
     pub(crate) fn with_query(&self, f: impl Fn(&DepGraphQuery<K>)) {
         if let Some(record_graph) = &self.record_graph {
-            f(&record_graph.lock())
+            record_graph.with_borrow(f)
         }
     }
 
@@ -307,7 +307,7 @@ impl<K: DepKind + Encodable<FileEncoder>> GraphEncoder<K> {
     ) -> DepNodeIndex {
         let _prof_timer = 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.with_lock(|status| status.encode_node(&node, &self.record_graph))
     }
 
     pub fn finish(self, profiler: &SelfProfilerRef) -> FileEncodeResult {

compiler/rustc_query_system/src/ich/impls_syntax.rs

@@ -80,7 +80,7 @@ impl<'a> HashStable<StableHashingContext<'a>> for SourceFile {
         src_hash.hash_stable(hcx, hasher);
 
         // We are always in `Lines` form by the time we reach here.
-        assert!(self.lines.borrow().is_lines());
+        assert!(self.lines.with_borrow(|lines| lines.is_lines()));
         self.lines(|lines| {
             // We only hash the relative position within this source_file
             lines.len().hash_stable(hcx, hasher);