Auto merge of #45312 - theotherjimmy:refactor-ensure, r=michaelwoerister

Refactor `ensure` and `try_get_with`

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.

@nikomatsakis this is the [refactor we discussed](#45228 (comment)) in the comment thread of #45228

Author: bors

src/librustc/ty/maps/plumbing.rs

@@ -12,7 +12,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};
+use dep_graph::{DepNodeIndex, DepNode, DepKind, DepNodeColor};
 use errors::{Diagnostic, DiagnosticBuilder};
 use ty::{TyCtxt};
 use ty::maps::Query; // NB: actually generated by the macros in this file
@@ -133,6 +133,40 @@ impl<'a, 'gcx, 'tcx> TyCtxt<'a, 'gcx, 'tcx> {
 
         Ok(result)
     }
+
+    /// 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(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, &dep_node) {
+                    Some(dep_node_index) => {
+                        debug_assert!(self.dep_graph.is_green(dep_node_index));
+                        self.dep_graph.read_index(dep_node_index);
+                        Some(dep_node_index)
+                    }
+                    None => {
+                        None
+                    }
+                }
+            }
+        }
+    }
 }
 
 // If enabled, send a message to the profile-queries thread
@@ -309,25 +343,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.try_mark_green_and_read(&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 +374,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.try_mark_green_and_read(&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
+                    // 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);
                 }
             }