Use a sharded dep node to dep node index map

Author: Zoxc

src/librustc/dep_graph/graph.rs

@@ -4,6 +4,7 @@ use rustc_data_structures::fx::{FxHashMap, FxHashSet};
 use rustc_data_structures::indexed_vec::{Idx, IndexVec};
 use smallvec::SmallVec;
 use rustc_data_structures::sync::{Lrc, Lock, AtomicU32, AtomicU64, Ordering};
+use rustc_data_structures::sharded::{self, Sharded};
 use std::sync::atomic::Ordering::{Acquire, SeqCst};
 use std::env;
 use std::hash::Hash;
@@ -395,7 +396,7 @@ impl DepGraph {
     #[inline]
     pub fn read(&self, v: DepNode) {
         if let Some(ref data) = self.data {
-            let map = data.current.node_to_node_index.lock();
+            let map = data.current.node_to_node_index.get_shard_by_value(&v).lock();
             if let Some(dep_node_index) = map.get(&v).copied() {
                 std::mem::drop(map);
                 data.read_index(dep_node_index);
@@ -419,6 +420,7 @@ impl DepGraph {
             .unwrap()
             .current
             .node_to_node_index
+            .get_shard_by_value(dep_node)
             .lock()
             .get(dep_node)
             .cloned()
@@ -428,7 +430,11 @@ impl DepGraph {
     #[inline]
     pub fn dep_node_exists(&self, dep_node: &DepNode) -> bool {
         if let Some(ref data) = self.data {
-            data.current.node_to_node_index.lock().contains_key(dep_node)
+            data.current
+                .node_to_node_index
+                .get_shard_by_value(&dep_node)
+                .lock()
+                .contains_key(dep_node)
         } else {
             false
         }
@@ -609,7 +615,11 @@ impl DepGraph {
 
         #[cfg(not(parallel_compiler))]
         {
-            debug_assert!(!data.current.borrow().node_to_node_index.contains_key(dep_node));
+            debug_assert!(!data.current
+                               .node_to_node_index
+                               .get_shard_by_value(dep_node)
+                               .lock()
+                               .contains_key(dep_node));
             debug_assert!(data.colors.get(prev_dep_node_index).is_none());
         }
 
@@ -955,7 +965,7 @@ struct DepNodeData {
 
 pub(super) struct CurrentDepGraph {
     data: Lock<IndexVec<DepNodeIndex, DepNodeData>>,
-    node_to_node_index: Lock<FxHashMap<DepNode, DepNodeIndex>>,
+    node_to_node_index: Sharded<FxHashMap<DepNode, DepNodeIndex>>,
     #[allow(dead_code)]
     forbidden_edge: Option<EdgeFilter>,
 
@@ -1008,8 +1018,8 @@ impl CurrentDepGraph {
 
         CurrentDepGraph {
             data: Lock::new(IndexVec::with_capacity(new_node_count_estimate)),
-            node_to_node_index: Lock::new(FxHashMap::with_capacity_and_hasher(
-                new_node_count_estimate,
+            node_to_node_index: Sharded::new(|| FxHashMap::with_capacity_and_hasher(
+                new_node_count_estimate / sharded::SHARDS,
                 Default::default(),
             )),
             anon_id_seed: stable_hasher.finish(),
@@ -1058,7 +1068,10 @@ impl CurrentDepGraph {
         edges: SmallVec<[DepNodeIndex; 8]>,
         fingerprint: Fingerprint
     ) -> DepNodeIndex {
-        debug_assert!(!self.node_to_node_index.lock().contains_key(&dep_node));
+        debug_assert!(!self.node_to_node_index
+                           .get_shard_by_value(&dep_node)
+                           .lock()
+                           .contains_key(&dep_node));
         self.intern_node(dep_node, edges, fingerprint)
     }
 
@@ -1068,7 +1081,7 @@ impl CurrentDepGraph {
         edges: SmallVec<[DepNodeIndex; 8]>,
         fingerprint: Fingerprint
     ) -> DepNodeIndex {
-        match self.node_to_node_index.lock().entry(dep_node) {
+        match self.node_to_node_index.get_shard_by_value(&dep_node).lock().entry(dep_node) {
             Entry::Occupied(entry) => *entry.get(),
             Entry::Vacant(entry) => {
                 let mut data = self.data.lock();

src/librustc_data_structures/sharded.rs

@@ -2,6 +2,7 @@ use std::hash::{Hasher, Hash};
 use std::mem;
 use std::borrow::Borrow;
 use std::collections::hash_map::RawEntryMut;
+use smallvec::SmallVec;
 use crate::fx::{FxHasher, FxHashMap};
 use crate::sync::{Lock, LockGuard};
 
@@ -15,7 +16,7 @@ const SHARD_BITS: usize = 5;
 #[cfg(not(parallel_compiler))]
 const SHARD_BITS: usize = 0;
 
-const SHARDS: usize = 1 << SHARD_BITS;
+pub const SHARDS: usize = 1 << SHARD_BITS;
 
 /// An array of cache-line aligned inner locked structures with convenience methods.
 #[derive(Clone)]
@@ -26,21 +27,36 @@ pub struct Sharded<T> {
 impl<T: Default> Default for Sharded<T> {
     #[inline]
     fn default() -> Self {
+        Self::new(|| T::default())
+    }
+}
+
+impl<T> Sharded<T> {
+    #[inline]
+    pub fn new(mut value: impl FnMut() -> T) -> Self {
+        // Create a vector of the values we want
+        let mut values: SmallVec<[_; SHARDS]> = (0..SHARDS).map(|_| {
+            CacheAligned(Lock::new(value()))
+        }).collect();
+
+        // Create an unintialized array
         let mut shards: mem::MaybeUninit<[CacheAligned<Lock<T>>; SHARDS]> =
             mem::MaybeUninit::uninit();
-        let first = shards.as_mut_ptr() as *mut CacheAligned<Lock<T>>;
+
         unsafe {
-            for i in 0..SHARDS {
-                first.add(i).write(CacheAligned(Lock::new(T::default())));
-            }
+            // Copy the values into our array
+            let first = shards.as_mut_ptr() as *mut CacheAligned<Lock<T>>;
+            values.as_ptr().copy_to_nonoverlapping(first, SHARDS);
+
+            // Ignore the content of the vector
+            values.set_len(0);
+
             Sharded {
                 shards: shards.assume_init(),
             }
         }
     }
-}
 
-impl<T> Sharded<T> {
     #[inline]
     pub fn get_shard_by_value<K: Hash + ?Sized>(&self, val: &K) -> &Lock<T> {
         if SHARDS == 1 {