Refactor `ensure` and `try_get_with` into `read_node_index`

There was a bit of code shared between `try_get_with` and `ensure`, after I
added `ensure`. I refactored that shared code into a query-agnostic method
called `read_node_index`.

The new method `read_node_index` will attempt to find the node
index (`DepNodeIndex`) of a query. When `read_node_index` finds the
`DepNodeIndex`, it marks the current query as a reader of the node it's
requesting the index of.

This is used by `try_get_with` and `ensure` as it elides the unimportant (to
them) details of if the query is invalidated by previous changed computation (Red)
or new and if they had to mark the query green. For both `try_get_with` and
`ensure`, they just need to know if they can lookup the results or have to
reevaluate.

2 files changed, 43 insertions, 48 deletions

diff --git a/src/librustc/dep_graph/graph.rs b/src/librustc/dep_graph/graph.rs
index 8aff042955c..38a27305c2b 100644
--- a/src/librustc/dep_graph/graph.rs
+++ b/src/librustc/dep_graph/graph.rs
@@ -427,6 +427,40 @@ impl DepGraph {
         self.data.as_ref().and_then(|data| data.colors.borrow().get(dep_node).cloned())
     }
 
+    /// 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
+    /// the a node index can be found for that node.
+    pub fn read_node_index(&self, tcx: TyCtxt, dep_node: &DepNode) -> Option<DepNodeIndex> {
+        match self.node_color(dep_node) {
+            Some(DepNodeColor::Green(dep_node_index)) => {
+                self.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.is_fully_enabled() {
+                    return None;
+                }
+                match self.try_mark_green(tcx, &dep_node) {
+                    Some(dep_node_index) => {
+                        debug_assert!(tcx.dep_graph.is_green(dep_node_index));
+                        tcx.dep_graph.read_index(dep_node_index);
+                        Some(dep_node_index)
+                    }
+                    None => {
+                        None
+                    }
+                }
+            }
+        }
+    }
+
     pub fn try_mark_green(&self,
                           tcx: TyCtxt,
                           dep_node: &DepNode)
diff --git a/src/librustc/ty/maps/plumbing.rs b/src/librustc/ty/maps/plumbing.rs
index d6eaf6d1bc4..13880b37136 100644
--- a/src/librustc/ty/maps/plumbing.rs
+++ b/src/librustc/ty/maps/plumbing.rs
@@ -309,25 +309,8 @@ macro_rules! define_maps {
                 }
 
                 if !dep_node.kind.is_input() {
-                    use dep_graph::DepNodeColor;
-                    if let Some(DepNodeColor::Green(dep_node_index)) = tcx.dep_graph
-                                                                          .node_color(&dep_node) {
+                    if let Some(dep_node_index) = tcx.dep_graph.read_node_index(tcx, &dep_node) {
                         profq_msg!(tcx, ProfileQueriesMsg::CacheHit);
-                        tcx.dep_graph.read_index(dep_node_index);
-                        return Self::load_from_disk_and_cache_in_memory(tcx,
-                                                                        key,
-                                                                        span,
-                                                                        dep_node_index)
-                    }
-
-                    debug!("ty::queries::{}::try_get_with(key={:?}) - running try_mark_green",
-                           stringify!($name),
-                           key);
-
-                    if let Some(dep_node_index) = tcx.dep_graph.try_mark_green(tcx, &dep_node) {
-                        debug_assert!(tcx.dep_graph.is_green(dep_node_index));
-                        profq_msg!(tcx, ProfileQueriesMsg::CacheHit);
-                        tcx.dep_graph.read_index(dep_node_index);
                         return Self::load_from_disk_and_cache_in_memory(tcx,
                                                                         key,
                                                                         span,
@@ -357,36 +340,14 @@ macro_rules! define_maps {
                 // Ensuring an "input" or anonymous query makes no sense
                 assert!(!dep_node.kind.is_anon());
                 assert!(!dep_node.kind.is_input());
-                use dep_graph::DepNodeColor;
-                match tcx.dep_graph.node_color(&dep_node) {
-                    Some(DepNodeColor::Green(dep_node_index)) => {
-                        tcx.dep_graph.read_index(dep_node_index);
-                    }
-                    Some(DepNodeColor::Red) => {
-                        // A DepNodeColor::Red DepNode means that this query was executed
-                        // before. We can not call `dep_graph.read()` here as we don't have
-                        // the DepNodeIndex. Instead, We call the query again to issue the
-                        // appropriate `dep_graph.read()` call. The performance cost this
-                        // introduces should be negligible as we'll immediately hit the
-                        // in-memory cache.
-                        let _ = tcx.$name(key);
-                    }
-                    None => {
-                        // Huh
-                        if !tcx.dep_graph.is_fully_enabled() {
-                            let _ = tcx.$name(key);
-                            return;
-                        }
-                        match tcx.dep_graph.try_mark_green(tcx, &dep_node) {
-                            Some(dep_node_index) => {
-                                debug_assert!(tcx.dep_graph.is_green(dep_node_index));
-                                tcx.dep_graph.read_index(dep_node_index);
-                            }
-                            None => {
-                                let _ = tcx.$name(key);
-                            }
-                        }
-                    }
+                if tcx.dep_graph.read_node_index(tcx, &dep_node).is_none() {
+                    // A None return from `read_node_index` 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
+                    // DepNodeIndex. We must invoke the query itself. The performance cost
+                    // this introduces should be negligible as we'll immediately hit the
+                    // in-memory cache, or another query down the line will.
+                    let _ = tcx.$name(key);
                 }
             }