Allow to feed a value in another query's cache

compiler/rustc_ast_lowering/src/lib.rs

@@ -497,7 +497,7 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
             self.tcx.hir().def_key(self.local_def_id(node_id)),
         );
 
-        let def_id = self.tcx.create_def(parent, data);
+        let def_id = self.tcx.create_def(parent, data).def_id();
 
         debug!("create_def: def_id_to_node_id[{:?}] <-> {:?}", def_id, node_id);
         self.resolver.node_id_to_def_id.insert(node_id, def_id);

compiler/rustc_hir/src/definitions.rs

@@ -368,10 +368,6 @@ impl Definitions {
         LocalDefId { local_def_index: self.table.allocate(key, def_path_hash) }
     }
 
-    pub fn iter_local_def_id(&self) -> impl Iterator<Item = LocalDefId> + '_ {
-        self.table.def_path_hashes.indices().map(|local_def_index| LocalDefId { local_def_index })
-    }
-
     #[inline(always)]
     pub fn local_def_path_hash_to_def_id(
         &self,
@@ -389,6 +385,10 @@ impl Definitions {
     pub fn def_path_hash_to_def_index_map(&self) -> &DefPathHashMap {
         &self.table.def_path_hash_to_index
     }
+
+    pub fn num_definitions(&self) -> usize {
+        self.table.def_path_hashes.len()
+    }
 }
 
 #[derive(Copy, Clone, PartialEq, Debug)]

compiler/rustc_macros/src/query.rs

@@ -114,6 +114,9 @@ struct QueryModifiers {
 
     /// Always remap the ParamEnv's constness before hashing.
     remap_env_constness: Option<Ident>,
+
+    /// Generate a `feed` method to set the query's value from another query.
+    feedable: Option<Ident>,
 }
 
 fn parse_query_modifiers(input: ParseStream<'_>) -> Result<QueryModifiers> {
@@ -128,6 +131,7 @@ fn parse_query_modifiers(input: ParseStream<'_>) -> Result<QueryModifiers> {
     let mut depth_limit = None;
     let mut separate_provide_extern = None;
     let mut remap_env_constness = None;
+    let mut feedable = None;
 
     while !input.is_empty() {
         let modifier: Ident = input.parse()?;
@@ -187,6 +191,8 @@ fn parse_query_modifiers(input: ParseStream<'_>) -> Result<QueryModifiers> {
             try_insert!(separate_provide_extern = modifier);
         } else if modifier == "remap_env_constness" {
             try_insert!(remap_env_constness = modifier);
+        } else if modifier == "feedable" {
+            try_insert!(feedable = modifier);
         } else {
             return Err(Error::new(modifier.span(), "unknown query modifier"));
         }
@@ -206,6 +212,7 @@ fn parse_query_modifiers(input: ParseStream<'_>) -> Result<QueryModifiers> {
         depth_limit,
         separate_provide_extern,
         remap_env_constness,
+        feedable,
     })
 }
 
@@ -296,6 +303,7 @@ pub fn rustc_queries(input: TokenStream) -> TokenStream {
     let mut query_stream = quote! {};
     let mut query_description_stream = quote! {};
     let mut query_cached_stream = quote! {};
+    let mut feedable_queries = quote! {};
 
     for query in queries.0 {
         let Query { name, arg, modifiers, .. } = &query;
@@ -350,6 +358,22 @@ pub fn rustc_queries(input: TokenStream) -> TokenStream {
             [#attribute_stream] fn #name(#arg) #result,
         });
 
+        if modifiers.feedable.is_some() {
+            assert!(modifiers.anon.is_none(), "Query {name} cannot be both `feedable` and `anon`.");
+            assert!(
+                modifiers.eval_always.is_none(),
+                "Query {name} cannot be both `feedable` and `eval_always`."
+            );
+            assert!(
+                modifiers.no_hash.is_none(),
+                "Query {name} cannot be both `feedable` and `no_hash`."
+            );
+            feedable_queries.extend(quote! {
+                #(#doc_comments)*
+                [#attribute_stream] fn #name(#arg) #result,
+            });
+        }
+
         add_query_desc_cached_impl(&query, &mut query_description_stream, &mut query_cached_stream);
     }
 
@@ -363,7 +387,11 @@ pub fn rustc_queries(input: TokenStream) -> TokenStream {
                 }
             }
         }
-
+        macro_rules! rustc_feedable_queries {
+            ( $macro:ident! ) => {
+                $macro!(#feedable_queries);
+            }
+        }
         pub mod descs {
             use super::*;
             #query_description_stream

compiler/rustc_middle/src/lib.rs

@@ -30,8 +30,10 @@
 #![feature(core_intrinsics)]
 #![feature(discriminant_kind)]
 #![feature(exhaustive_patterns)]
+#![feature(generators)]
 #![feature(get_mut_unchecked)]
 #![feature(if_let_guard)]
+#![feature(iter_from_generator)]
 #![feature(negative_impls)]
 #![feature(never_type)]
 #![feature(extern_types)]

compiler/rustc_middle/src/ty/context.rs

@@ -53,6 +53,7 @@ use rustc_hir::{
 use rustc_index::vec::{Idx, IndexVec};
 use rustc_macros::HashStable;
 use rustc_middle::mir::FakeReadCause;
+use rustc_query_system::dep_graph::DepNodeIndex;
 use rustc_query_system::ich::StableHashingContext;
 use rustc_serialize::opaque::{FileEncodeResult, FileEncoder};
 use rustc_session::config::{CrateType, OutputFilenames};
@@ -1009,6 +1010,21 @@ pub struct FreeRegionInfo {
     pub is_impl_item: bool,
 }
 
+/// This struct should only be created by `create_def`.
+#[derive(Copy, Clone)]
+pub struct TyCtxtFeed<'tcx> {
+    pub tcx: TyCtxt<'tcx>,
+    // Do not allow direct access, as downstream code must not mutate this field.
+    def_id: LocalDefId,
+}
+
+impl<'tcx> TyCtxtFeed<'tcx> {
+    #[inline(always)]
+    pub fn def_id(&self) -> LocalDefId {
+        self.def_id
+    }
+}
+
 /// The central data structure of the compiler. It stores references
 /// to the various **arenas** and also houses the results of the
 /// various **compiler queries** that have been performed. See the
@@ -1471,12 +1487,15 @@ impl<'tcx> TyCtxt<'tcx> {
     }
 
     /// Create a new definition within the incr. comp. engine.
-    pub fn create_def(self, parent: LocalDefId, data: hir::definitions::DefPathData) -> LocalDefId {
+    pub fn create_def(
+        self,
+        parent: LocalDefId,
+        data: hir::definitions::DefPathData,
+    ) -> TyCtxtFeed<'tcx> {
         // This function modifies `self.definitions` using a side-effect.
         // We need to ensure that these side effects are re-run by the incr. comp. engine.
         // Depending on the forever-red node will tell the graph that the calling query
         // needs to be re-evaluated.
-        use rustc_query_system::dep_graph::DepNodeIndex;
         self.dep_graph.read_index(DepNodeIndex::FOREVER_RED_NODE);
 
         // The following call has the side effect of modifying the tables inside `definitions`.
@@ -1493,23 +1512,38 @@ impl<'tcx> TyCtxt<'tcx> {
         // This is fine because:
         // - those queries are `eval_always` so we won't miss their result changing;
         // - this write will have happened before these queries are called.
-        self.definitions.write().create_def(parent, data)
+        let def_id = self.definitions.write().create_def(parent, data);
+
+        TyCtxtFeed { tcx: self, def_id }
     }
 
     pub fn iter_local_def_id(self) -> impl Iterator<Item = LocalDefId> + 'tcx {
-        // Create a dependency to the crate to be sure we re-execute this when the amount of
+        // Create a dependency to the red node to be sure we re-execute this when the amount of
         // definitions change.
-        self.ensure().hir_crate(());
-        // Leak a read lock once we start iterating on definitions, to prevent adding new ones
-        // while iterating.  If some query needs to add definitions, it should be `ensure`d above.
-        let definitions = self.definitions.leak();
-        definitions.iter_local_def_id()
+        self.dep_graph.read_index(DepNodeIndex::FOREVER_RED_NODE);
+
+        let definitions = &self.definitions;
+        std::iter::from_generator(|| {
+            let mut i = 0;
+
+            // Recompute the number of definitions each time, because our caller may be creating
+            // new ones.
+            while i < { definitions.read().num_definitions() } {
+                let local_def_index = rustc_span::def_id::DefIndex::from_usize(i);
+                yield LocalDefId { local_def_index };
+                i += 1;
+            }
+
+            // Leak a read lock once we finish iterating on definitions, to prevent adding new ones.
+            definitions.leak();
+        })
     }
 
     pub fn def_path_table(self) -> &'tcx rustc_hir::definitions::DefPathTable {
         // Create a dependency to the crate to be sure we re-execute this when the amount of
         // definitions change.
-        self.ensure().hir_crate(());
+        self.dep_graph.read_index(DepNodeIndex::FOREVER_RED_NODE);
+
         // Leak a read lock once we start iterating on definitions, to prevent adding new ones
         // while iterating.  If some query needs to add definitions, it should be `ensure`d above.
         let definitions = self.definitions.leak();

compiler/rustc_middle/src/ty/query.rs

@@ -28,6 +28,7 @@ use crate::traits::query::{
 };
 use crate::traits::specialization_graph;
 use crate::traits::{self, ImplSource};
+use crate::ty::context::TyCtxtFeed;
 use crate::ty::fast_reject::SimplifiedType;
 use crate::ty::layout::TyAndLayout;
 use crate::ty::subst::{GenericArg, SubstsRef};
@@ -327,6 +328,46 @@ macro_rules! define_callbacks {
     };
 }
 
+macro_rules! define_feedable {
+    ($($(#[$attr:meta])* [$($modifiers:tt)*] fn $name:ident($($K:tt)*) -> $V:ty,)*) => {
+        impl<'tcx> TyCtxtFeed<'tcx> {
+            $($(#[$attr])*
+            #[inline(always)]
+            pub fn $name(self, value: $V) -> query_stored::$name<'tcx> {
+                let key = self.def_id().into_query_param();
+                opt_remap_env_constness!([$($modifiers)*][key]);
+
+                let tcx = self.tcx;
+                let cache = &tcx.query_caches.$name;
+
+                let cached = try_get_cached(tcx, cache, &key, copy);
+
+                match cached {
+                    Ok(old) => {
+                        assert_eq!(
+                            value, old,
+                            "Trying to feed an already recorded value for query {} key={key:?}",
+                            stringify!($name),
+                        );
+                        return old;
+                    }
+                    Err(()) => (),
+                }
+
+                let dep_node = dep_graph::DepNode::construct(tcx, dep_graph::DepKind::$name, &key);
+                let dep_node_index = tcx.dep_graph.with_feed_task(
+                    dep_node,
+                    tcx,
+                    key,
+                    &value,
+                    dep_graph::hash_result,
+                );
+                cache.complete(key, value, dep_node_index)
+            })*
+        }
+    }
+}
+
 // Each of these queries corresponds to a function pointer field in the
 // `Providers` struct for requesting a value of that type, and a method
 // on `tcx: TyCtxt` (and `tcx.at(span)`) for doing that request in a way
@@ -340,6 +381,7 @@ macro_rules! define_callbacks {
 // as they will raise an fatal error on query cycles instead.
 
 rustc_query_append! { define_callbacks! }
+rustc_feedable_queries! { define_feedable! }
 
 mod sealed {
     use super::{DefId, LocalDefId, OwnerId};

compiler/rustc_query_impl/src/plumbing.rs

@@ -252,6 +252,18 @@ macro_rules! depth_limit {
     };
 }
 
+macro_rules! feedable {
+    ([]) => {{
+        false
+    }};
+    ([(feedable) $($rest:tt)*]) => {{
+        true
+    }};
+    ([$other:tt $($modifiers:tt)*]) => {
+        feedable!([$($modifiers)*])
+    };
+}
+
 macro_rules! hash_result {
     ([]) => {{
         Some(dep_graph::hash_result)
@@ -309,7 +321,7 @@ pub(crate) fn create_query_frame<
         ty::print::with_forced_impl_filename_line!(do_describe(tcx.tcx, key))
     );
     let description =
-        if tcx.sess.verbose() { format!("{} [{}]", description, name) } else { description };
+        if tcx.sess.verbose() { format!("{} [{:?}]", description, name) } else { description };
     let span = if kind == dep_graph::DepKind::def_span {
         // The `def_span` query is used to calculate `default_span`,
         // so exit to avoid infinite recursion.
@@ -491,6 +503,7 @@ macro_rules! define_queries {
                     anon: is_anon!([$($modifiers)*]),
                     eval_always: is_eval_always!([$($modifiers)*]),
                     depth_limit: depth_limit!([$($modifiers)*]),
+                    feedable: feedable!([$($modifiers)*]),
                     dep_kind: dep_graph::DepKind::$name,
                     hash_result: hash_result!([$($modifiers)*]),
                     handle_cycle_error: handle_cycle_error!([$($modifiers)*]),