Optimize try_mark_green and eliminate the lock on dep node colors

4 files changed, 117 insertions, 116 deletions

diff --git a/src/librustc/dep_graph/graph.rs b/src/librustc/dep_graph/graph.rs
index 501ef01d74c..f8e426088de 100644
--- a/src/librustc/dep_graph/graph.rs
+++ b/src/librustc/dep_graph/graph.rs
@@ -3,7 +3,7 @@ use rustc_data_structures::stable_hasher::{HashStable, StableHasher};
 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};
+use rustc_data_structures::sync::{Lrc, Lock, AtomicU32, Ordering};
 use std::env;
 use std::hash::Hash;
 use std::collections::hash_map::Entry;
@@ -58,7 +58,7 @@ struct DepGraphData {
     /// nodes and edges as well as all fingerprints of nodes that have them.
     previous: PreviousDepGraph,
 
-    colors: Lock<DepNodeColorMap>,
+    colors: DepNodeColorMap,
 
     /// When we load, there may be `.o` files, cached mir, or other such
     /// things available to us. If we find that they are not dirty, we
@@ -84,7 +84,7 @@ impl DepGraph {
                 dep_node_debug: Default::default(),
                 current: Lock::new(CurrentDepGraph::new(prev_graph_node_count)),
                 previous: prev_graph,
-                colors: Lock::new(DepNodeColorMap::new(prev_graph_node_count)),
+                colors: DepNodeColorMap::new(prev_graph_node_count),
                 loaded_from_cache: Default::default(),
             })),
         }
@@ -282,12 +282,11 @@ impl DepGraph {
                     DepNodeColor::Red
                 };
 
-                let mut colors = data.colors.borrow_mut();
-                debug_assert!(colors.get(prev_index).is_none(),
+                debug_assert!(data.colors.get(prev_index).is_none(),
                               "DepGraph::with_task() - Duplicate DepNodeColor \
                                insertion for {:?}", key);
 
-                colors.insert(prev_index, color);
+                data.colors.insert(prev_index, color);
             }
 
             (result, dep_node_index)
@@ -502,7 +501,7 @@ impl DepGraph {
     pub fn node_color(&self, dep_node: &DepNode) -> Option<DepNodeColor> {
         if let Some(ref data) = self.data {
             if let Some(prev_index) = data.previous.node_to_index_opt(dep_node) {
-                return data.colors.borrow().get(prev_index)
+                return data.colors.get(prev_index)
             } else {
                 // This is a node that did not exist in the previous compilation
                 // session, so we consider it to be red.
@@ -513,56 +512,86 @@ impl DepGraph {
         None
     }
 
-    pub fn try_mark_green<'tcx>(&self,
-                                tcx: TyCtxt<'_, 'tcx, 'tcx>,
-                                dep_node: &DepNode)
-                                -> Option<DepNodeIndex> {
-        debug!("try_mark_green({:?}) - BEGIN", dep_node);
-        let data = self.data.as_ref().unwrap();
+    /// Try to read a node index for the node dep_node.
+    /// A node will have an index, when it's already been marked green, or when we can mark it
+    /// green. This function will mark the current task as a reader of the specified node, when
+    /// a node index can be found for that node.
+    pub fn try_mark_green_and_read(
+        &self,
+        tcx: TyCtxt<'_, '_, '_>,
+        dep_node: &DepNode
+    ) -> Option<(SerializedDepNodeIndex, DepNodeIndex)> {
+        self.try_mark_green(tcx, dep_node).map(|(prev_index, dep_node_index)| {
+            debug_assert!(self.is_green(&dep_node));
+            self.read_index(dep_node_index);
+            (prev_index, dep_node_index)
+        })
+    }
 
-        #[cfg(not(parallel_queries))]
-        debug_assert!(!data.current.borrow().node_to_node_index.contains_key(dep_node));
-
-        if dep_node.kind.is_input() {
-            // We should only hit try_mark_green() for inputs that do not exist
-            // anymore in the current compilation session. Existing inputs are
-            // eagerly marked as either red/green before any queries are
-            // executed.
-            debug_assert!(dep_node.extract_def_id(tcx).is_none());
-            debug!("try_mark_green({:?}) - END - DepNode is deleted input", dep_node);
-            return None;
-        }
+    pub fn try_mark_green(
+        &self,
+        tcx: TyCtxt<'_, '_, '_>,
+        dep_node: &DepNode
+    ) -> Option<(SerializedDepNodeIndex, DepNodeIndex)> {
+        debug_assert!(!dep_node.kind.is_input());
 
-        let (prev_deps, prev_dep_node_index) = match data.previous.edges_from(dep_node) {
-            Some(prev) => {
-                // This DepNode and the corresponding query invocation existed
-                // in the previous compilation session too, so we can try to
-                // mark it as green by recursively marking all of its
-                // dependencies green.
-                prev
-            }
+        // Return None if the dep graph is disabled
+        let data = self.data.as_ref()?;
+
+        // Return None if the dep node didn't exist in the previous session
+        let prev_index = data.previous.node_to_index_opt(dep_node)?;
+
+        match data.colors.get(prev_index) {
+            Some(DepNodeColor::Green(dep_node_index)) => Some((prev_index, dep_node_index)),
+            Some(DepNodeColor::Red) => None,
             None => {
-                // This DepNode did not exist in the previous compilation session,
-                // so we cannot mark it as green.
-                debug!("try_mark_green({:?}) - END - DepNode does not exist in \
-                        current compilation session anymore", dep_node);
-                return None
+                self.try_mark_previous_green(
+                    tcx.global_tcx(),
+                    data,
+                    prev_index,
+                    &dep_node
+                ).map(|dep_node_index| {
+                    (prev_index, dep_node_index)
+                })
             }
-        };
+        }
+    }
 
-        debug_assert!(data.colors.borrow().get(prev_dep_node_index).is_none());
+    fn try_mark_previous_green<'tcx>(
+        &self,
+        tcx: TyCtxt<'_, 'tcx, 'tcx>,
+        data: &DepGraphData,
+        prev_dep_node_index: SerializedDepNodeIndex,
+        dep_node: &DepNode
+    ) -> Option<DepNodeIndex> {
+        debug!("try_mark_previous_green({:?}) - BEGIN", dep_node);
+
+        #[cfg(not(parallel_queries))]
+        {
+            debug_assert!(!data.current.borrow().node_to_node_index.contains_key(dep_node));
+            debug_assert!(data.colors.get(prev_dep_node_index).is_none());
+        }
+
+        debug_assert!(!dep_node.kind.is_input());
+        debug_assert_eq!(data.previous.index_to_node(prev_dep_node_index), *dep_node);
+
+        // We never try to mark inputs as green
+        // FIXME: Make an debug_assert!
+        assert!(!dep_node.kind.is_input());
+
+        let prev_deps = data.previous.edge_targets_from(prev_dep_node_index);
 
         let mut current_deps = SmallVec::new();
 
         for &dep_dep_node_index in prev_deps {
-            let dep_dep_node_color = data.colors.borrow().get(dep_dep_node_index);
+            let dep_dep_node_color = data.colors.get(dep_dep_node_index);
 
             match dep_dep_node_color {
                 Some(DepNodeColor::Green(node_index)) => {
                     // This dependency has been marked as green before, we are
                     // still fine and can continue with checking the other
                     // dependencies.
-                    debug!("try_mark_green({:?}) --- found dependency {:?} to \
+                    debug!("try_mark_previous_green({:?}) --- found dependency {:?} to \
                             be immediately green",
                             dep_node,
                             data.previous.index_to_node(dep_dep_node_index));
@@ -573,7 +602,7 @@ impl DepGraph {
                     // compared to the previous compilation session. We cannot
                     // mark the DepNode as green and also don't need to bother
                     // with checking any of the other dependencies.
-                    debug!("try_mark_green({:?}) - END - dependency {:?} was \
+                    debug!("try_mark_previous_green({:?}) - END - dependency {:?} was \
                             immediately red",
                             dep_node,
                             data.previous.index_to_node(dep_dep_node_index));
@@ -585,12 +614,18 @@ impl DepGraph {
                     // We don't know the state of this dependency. If it isn't
                     // an input node, let's try to mark it green recursively.
                     if !dep_dep_node.kind.is_input() {
-                         debug!("try_mark_green({:?}) --- state of dependency {:?} \
+                         debug!("try_mark_previous_green({:?}) --- state of dependency {:?} \
                                  is unknown, trying to mark it green", dep_node,
                                  dep_dep_node);
 
-                        if let Some(node_index) = self.try_mark_green(tcx, dep_dep_node) {
-                            debug!("try_mark_green({:?}) --- managed to MARK \
+                        let node_index = self.try_mark_previous_green(
+                            tcx,
+                            data,
+                            dep_dep_node_index,
+                            dep_dep_node
+                        );
+                        if let Some(node_index) = node_index {
+                            debug!("try_mark_previous_green({:?}) --- managed to MARK \
                                     dependency {:?} as green", dep_node, dep_dep_node);
                             current_deps.push(node_index);
                             continue;
@@ -620,20 +655,20 @@ impl DepGraph {
                     }
 
                     // We failed to mark it green, so we try to force the query.
-                    debug!("try_mark_green({:?}) --- trying to force \
+                    debug!("try_mark_previous_green({:?}) --- trying to force \
                             dependency {:?}", dep_node, dep_dep_node);
                     if ::ty::query::force_from_dep_node(tcx, dep_dep_node) {
-                        let dep_dep_node_color = data.colors.borrow().get(dep_dep_node_index);
+                        let dep_dep_node_color = data.colors.get(dep_dep_node_index);
 
                         match dep_dep_node_color {
                             Some(DepNodeColor::Green(node_index)) => {
-                                debug!("try_mark_green({:?}) --- managed to \
+                                debug!("try_mark_previous_green({:?}) --- managed to \
                                         FORCE dependency {:?} to green",
                                         dep_node, dep_dep_node);
                                 current_deps.push(node_index);
                             }
                             Some(DepNodeColor::Red) => {
-                                debug!("try_mark_green({:?}) - END - \
+                                debug!("try_mark_previous_green({:?}) - END - \
                                         dependency {:?} was red after forcing",
                                        dep_node,
                                        dep_dep_node);
@@ -641,7 +676,7 @@ impl DepGraph {
                             }
                             None => {
                                 if !tcx.sess.has_errors() {
-                                    bug!("try_mark_green() - Forcing the DepNode \
+                                    bug!("try_mark_previous_green() - Forcing the DepNode \
                                           should have set its color")
                                 } else {
                                     // If the query we just forced has resulted
@@ -653,7 +688,7 @@ impl DepGraph {
                         }
                     } else {
                         // The DepNode could not be forced.
-                        debug!("try_mark_green({:?}) - END - dependency {:?} \
+                        debug!("try_mark_previous_green({:?}) - END - dependency {:?} \
                                 could not be forced", dep_node, dep_dep_node);
                         return None
                     }
@@ -705,16 +740,15 @@ impl DepGraph {
         }
 
         // ... and finally storing a "Green" entry in the color map.
-        let mut colors = data.colors.borrow_mut();
         // Multiple threads can all write the same color here
         #[cfg(not(parallel_queries))]
-        debug_assert!(colors.get(prev_dep_node_index).is_none(),
-                      "DepGraph::try_mark_green() - Duplicate DepNodeColor \
+        debug_assert!(data.colors.get(prev_dep_node_index).is_none(),
+                      "DepGraph::try_mark_previous_green() - Duplicate DepNodeColor \
                       insertion for {:?}", dep_node);
 
-        colors.insert(prev_dep_node_index, DepNodeColor::Green(dep_node_index));
+        data.colors.insert(prev_dep_node_index, DepNodeColor::Green(dep_node_index));
 
-        debug!("try_mark_green({:?}) - END - successfully marked as green", dep_node);
+        debug!("try_mark_previous_green({:?}) - END - successfully marked as green", dep_node);
         Some(dep_node_index)
     }
 
@@ -735,9 +769,8 @@ impl DepGraph {
     pub fn exec_cache_promotions<'a, 'tcx>(&self, tcx: TyCtxt<'a, 'tcx, 'tcx>) {
         let green_nodes: Vec<DepNode> = {
             let data = self.data.as_ref().unwrap();
-            let colors = data.colors.borrow();
-            colors.values.indices().filter_map(|prev_index| {
-                match colors.get(prev_index) {
+            data.colors.values.indices().filter_map(|prev_index| {
+                match data.colors.get(prev_index) {
                     Some(DepNodeColor::Green(_)) => {
                         let dep_node = data.previous.index_to_node(prev_index);
                         if dep_node.cache_on_disk(tcx) {
@@ -1035,7 +1068,7 @@ pub struct TaskDeps {
 // A data structure that stores Option<DepNodeColor> values as a contiguous
 // array, using one u32 per entry.
 struct DepNodeColorMap {
-    values: IndexVec<SerializedDepNodeIndex, u32>,
+    values: IndexVec<SerializedDepNodeIndex, AtomicU32>,
 }
 
 const COMPRESSED_NONE: u32 = 0;
@@ -1045,12 +1078,12 @@ const COMPRESSED_FIRST_GREEN: u32 = 2;
 impl DepNodeColorMap {
     fn new(size: usize) -> DepNodeColorMap {
         DepNodeColorMap {
-            values: IndexVec::from_elem_n(COMPRESSED_NONE, size)
+            values: (0..size).map(|_| AtomicU32::new(COMPRESSED_NONE)).collect(),
         }
     }
 
     fn get(&self, index: SerializedDepNodeIndex) -> Option<DepNodeColor> {
-        match self.values[index] {
+        match self.values[index].load(Ordering::Acquire) {
             COMPRESSED_NONE => None,
             COMPRESSED_RED => Some(DepNodeColor::Red),
             value => Some(DepNodeColor::Green(DepNodeIndex::from_u32(
@@ -1059,10 +1092,10 @@ impl DepNodeColorMap {
         }
     }
 
-    fn insert(&mut self, index: SerializedDepNodeIndex, color: DepNodeColor) {
-        self.values[index] = match color {
+    fn insert(&self, index: SerializedDepNodeIndex, color: DepNodeColor) {
+        self.values[index].store(match color {
             DepNodeColor::Red => COMPRESSED_RED,
             DepNodeColor::Green(index) => index.as_u32() + COMPRESSED_FIRST_GREEN,
-        }
+        }, Ordering::Release)
     }
 }
diff --git a/src/librustc/dep_graph/prev.rs b/src/librustc/dep_graph/prev.rs
index 76d2954b4e3..ea5350ac97f 100644
--- a/src/librustc/dep_graph/prev.rs
+++ b/src/librustc/dep_graph/prev.rs
@@ -19,14 +19,11 @@ impl PreviousDepGraph {
     }
 
     #[inline]
-    pub fn edges_from(&self,
-                      dep_node: &DepNode)
-                      -> Option<(&[SerializedDepNodeIndex], SerializedDepNodeIndex)> {
-        self.index
-            .get(dep_node)
-            .map(|&node_index| {
-                (self.data.edge_targets_from(node_index), node_index)
-            })
+    pub fn edge_targets_from(
+        &self,
+        dep_node_index: SerializedDepNodeIndex
+    ) -> &[SerializedDepNodeIndex] {
+        self.data.edge_targets_from(dep_node_index)
     }
 
     #[inline]
diff --git a/src/librustc/ty/query/plumbing.rs b/src/librustc/ty/query/plumbing.rs
index 562cd75a75f..5d827e07c59 100644
--- a/src/librustc/ty/query/plumbing.rs
+++ b/src/librustc/ty/query/plumbing.rs
@@ -2,7 +2,7 @@
 //! that generate the actual methods on tcx which find and execute the
 //! provider, manage the caches, and so forth.
 
-use dep_graph::{DepNodeIndex, DepNode, DepKind, DepNodeColor};
+use dep_graph::{DepNodeIndex, DepNode, DepKind, SerializedDepNodeIndex};
 use errors::DiagnosticBuilder;
 use errors::Level;
 use errors::Diagnostic;
@@ -335,40 +335,6 @@ impl<'a, 'gcx, 'tcx> TyCtxt<'a, 'gcx, 'tcx> {
         eprintln!("end of query stack");
     }
 
-    /// Try to read a node index for the node dep_node.
-    /// A node will have an index, when it's already been marked green, or when we can mark it
-    /// green. This function will mark the current task as a reader of the specified node, when
-    /// a node index can be found for that node.
-    pub(super) fn try_mark_green_and_read(self, dep_node: &DepNode) -> Option<DepNodeIndex> {
-        match self.dep_graph.node_color(dep_node) {
-            Some(DepNodeColor::Green(dep_node_index)) => {
-                self.dep_graph.read_index(dep_node_index);
-                Some(dep_node_index)
-            }
-            Some(DepNodeColor::Red) => {
-                None
-            }
-            None => {
-                // try_mark_green (called below) will panic when full incremental
-                // compilation is disabled. If that's the case, we can't try to mark nodes
-                // as green anyway, so we can safely return None here.
-                if !self.dep_graph.is_fully_enabled() {
-                    return None;
-                }
-                match self.dep_graph.try_mark_green(self.global_tcx(), &dep_node) {
-                    Some(dep_node_index) => {
-                        debug_assert!(self.dep_graph.is_green(&dep_node));
-                        self.dep_graph.read_index(dep_node_index);
-                        Some(dep_node_index)
-                    }
-                    None => {
-                        None
-                    }
-                }
-            }
-        }
-    }
-
     #[inline(never)]
     fn try_get_with<Q: QueryDescription<'gcx>>(
         self,
@@ -435,10 +401,13 @@ impl<'a, 'gcx, 'tcx> TyCtxt<'a, 'gcx, 'tcx> {
         }
 
         if !dep_node.kind.is_input() {
-            if let Some(dep_node_index) = self.try_mark_green_and_read(&dep_node) {
+            if let Some((prev_dep_node_index,
+                         dep_node_index)) = self.dep_graph.try_mark_green_and_read(self,
+                                                                                   &dep_node) {
                 return Ok(self.load_from_disk_and_cache_in_memory::<Q>(
                     key,
                     job,
+                    prev_dep_node_index,
                     dep_node_index,
                     &dep_node
                 ))
@@ -454,6 +423,7 @@ impl<'a, 'gcx, 'tcx> TyCtxt<'a, 'gcx, 'tcx> {
         self,
         key: Q::Key,
         job: JobOwner<'a, 'gcx, Q>,
+        prev_dep_node_index: SerializedDepNodeIndex,
         dep_node_index: DepNodeIndex,
         dep_node: &DepNode
     ) -> Q::Value
@@ -466,10 +436,7 @@ impl<'a, 'gcx, 'tcx> TyCtxt<'a, 'gcx, 'tcx> {
         // First we try to load the result from the on-disk cache
         let result = if Q::cache_on_disk(key.clone()) &&
                         self.sess.opts.debugging_opts.incremental_queries {
-            let prev_dep_node_index =
-                self.dep_graph.prev_dep_node_index_of(dep_node);
-            let result = Q::try_load_from_disk(self.global_tcx(),
-                                               prev_dep_node_index);
+            let result = Q::try_load_from_disk(self.global_tcx(), prev_dep_node_index);
 
             // We always expect to find a cached result for things that
             // can be forced from DepNode.
@@ -624,7 +591,7 @@ impl<'a, 'gcx, 'tcx> TyCtxt<'a, 'gcx, 'tcx> {
         // Ensuring an "input" or anonymous query makes no sense
         assert!(!dep_node.kind.is_anon());
         assert!(!dep_node.kind.is_input());
-        if self.try_mark_green_and_read(&dep_node).is_none() {
+        if self.dep_graph.try_mark_green_and_read(self, &dep_node).is_none() {
             // A None return from `try_mark_green_and_read` means that this is either
             // a new dep node or that the dep node has already been marked red.
             // Either way, we can't call `dep_graph.read()` as we don't have the
diff --git a/src/librustc_data_structures/sync.rs b/src/librustc_data_structures/sync.rs
index f9f94f0be7b..0253eef4dfa 100644
--- a/src/librustc_data_structures/sync.rs
+++ b/src/librustc_data_structures/sync.rs
@@ -70,6 +70,7 @@ cfg_if! {
         pub struct Atomic<T: Copy>(Cell<T>);
 
         impl<T: Copy> Atomic<T> {
+            #[inline]
             pub fn new(v: T) -> Self {
                 Atomic(Cell::new(v))
             }
@@ -80,10 +81,12 @@ cfg_if! {
                 self.0.into_inner()
             }
 
+            #[inline]
             pub fn load(&self, _: Ordering) -> T {
                 self.0.get()
             }
 
+            #[inline]
             pub fn store(&self, val: T, _: Ordering) {
                 self.0.set(val)
             }
@@ -118,6 +121,7 @@ cfg_if! {
 
         pub type AtomicUsize = Atomic<usize>;
         pub type AtomicBool = Atomic<bool>;
+        pub type AtomicU32 = Atomic<u32>;
         pub type AtomicU64 = Atomic<u64>;
 
         pub use self::serial_join as join;
@@ -223,7 +227,7 @@ cfg_if! {
         pub use parking_lot::MutexGuard as LockGuard;
         pub use parking_lot::MappedMutexGuard as MappedLockGuard;
 
-        pub use std::sync::atomic::{AtomicBool, AtomicUsize, AtomicU64};
+        pub use std::sync::atomic::{AtomicBool, AtomicUsize, AtomicU32, AtomicU64};
 
         pub use std::sync::Arc as Lrc;
         pub use std::sync::Weak as Weak;