From 67b3e7cc435931330cae5ae104ec7a7d37df89dd Mon Sep 17 00:00:00 2001
From: Jimmy Brisson <theotherjimmy@gmail.com>
Date: Tue, 10 Oct 2017 22:54:06 -0500
Subject: [PATCH] Implement query ensure

`$query::ensure()` guarantees that one of two things is true after it returns:
  - The query has all green inputs.
  - The query has been executed.
 and counts as a read from the source query
---
 src/librustc/ty/maps/plumbing.rs | 42 ++++++++++++++++++++++++++++++++
 1 file changed, 42 insertions(+)

diff --git a/src/librustc/ty/maps/plumbing.rs b/src/librustc/ty/maps/plumbing.rs
index 88b619558d9..28397ac33d9 100644
--- a/src/librustc/ty/maps/plumbing.rs
+++ b/src/librustc/ty/maps/plumbing.rs
@@ -344,6 +344,48 @@ macro_rules! define_maps {
                 }
             }
 
+            /// Ensure that either this query has all green inputs or been executed.
+            /// Executing query::ensure(D) is considered a read of the dep-node D.
+            ///
+            /// This function is particularly useful when executing passes for their
+            /// side-effects -- e.g., in order to report errors for erroneous programs.
+            ///
+            /// Note: The optimization is only available during incr. comp.
+            pub fn ensure(tcx: TyCtxt<'a, $tcx, 'lcx>, key: $K) -> () {
+                let dep_node = Self::to_dep_node(tcx, &key);
+
+                // 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)) => {
+                        profq_msg!(tcx, ProfileQueriesMsg::CacheHit);
+                        tcx.dep_graph.read_index(dep_node_index);
+                    }
+                    Some(DepNodeColor::Red) => {
+                        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));
+                                profq_msg!(tcx, ProfileQueriesMsg::CacheHit);
+                                tcx.dep_graph.read_index(dep_node_index);
+                            }
+                            None => {
+                                let _ = tcx.$name(key);
+                            }
+                        }
+                    }
+                }
+            }
+
             fn compute_result(tcx: TyCtxt<'a, $tcx, 'lcx>, key: $K) -> $V {
                 let provider = tcx.maps.providers[key.map_crate()].$name;
                 provider(tcx.global_tcx(), key)