more with shard

Author: SparrowLii

compiler/rustc_data_structures/src/sharded.rs

@@ -31,13 +31,21 @@ impl<T: Default> Default for Sharded<T> {
     }
 }
 
-impl<T> Sharded<T> {
+impl<T: Default> Sharded<T> {
     #[inline]
     pub fn new(mut value: impl FnMut() -> T) -> Self {
-        Sharded {
-            single_thread: !active(),
-            shard: Lock::new(value()),
-            shards: [(); SHARDS].map(|()| CacheAligned(Lock::new(value()))),
+        if likely(!active()) {
+            Sharded {
+                single_thread: !active(),
+                shard: Lock::new(value()),
+                shards: [(); SHARDS].map(|()| CacheAligned(Lock::new(T::default()))),
+            }
+        } else {
+            Sharded {
+                single_thread: !active(),
+                shard: Lock::new(T::default()),
+                shards: [(); SHARDS].map(|()| CacheAligned(Lock::new(value()))),
+            }
         }
     }
 
@@ -59,16 +67,6 @@ impl<T> Sharded<T> {
         }
     }
 
-    /// The shard is selected by hashing `val` with `FxHasher`.
-    #[inline]
-    pub fn get_shard_by_value<K: Hash + ?Sized>(&self, val: &K) -> &Lock<T> {
-        if likely(self.single_thread) {
-            &self.shard
-        } else {
-            &self.shards[get_shard_index_by_hash(make_hash(val))].0
-        }
-    }
-
     #[inline]
     pub fn with_get_shard_by_hash<F: FnOnce(&mut T) -> R, R>(
         &self,

compiler/rustc_query_system/src/dep_graph/graph.rs

@@ -1333,7 +1333,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/query/plumbing.rs

@@ -161,12 +161,10 @@ where
         cache.complete(key, result, dep_node_index);
 
         let job = {
-            let mut lock = state.active.get_shard_by_value(&key).lock();
-
-            match lock.remove(&key).unwrap() {
+            state.active.with_get_shard_by_value(&key, |lock| match lock.remove(&key).unwrap() {
                 QueryResult::Started(job) => job,
                 QueryResult::Poisoned => panic!(),
-            }
+            })
         };
 
         job.signal_complete();
@@ -184,14 +182,14 @@ where
         // Poison the query so jobs waiting on it panic.
         let state = self.state;
         let job = {
-            let mut shard = state.active.get_shard_by_value(&self.key).lock();
-
-            let job = match shard.remove(&self.key).unwrap() {
-                QueryResult::Started(job) => job,
-                QueryResult::Poisoned => panic!(),
-            };
-            shard.insert(self.key, QueryResult::Poisoned);
-            job
+            state.active.with_get_shard_by_value(&self.key, |shard| {
+                let job = match shard.remove(&self.key).unwrap() {
+                    QueryResult::Started(job) => job,
+                    QueryResult::Poisoned => panic!(),
+                };
+                shard.insert(self.key, QueryResult::Poisoned);
+                job
+            })
         };
         // Also signal the completion of the job, so waiters
         // will continue execution.