Auto merge of #104940 - cjgillot:query-feed-simple, r=oli-obk

Allow to feed a value in another query's cache

Restricted version of #96840

A query can create new definitions.

If those definitions are created after HIR lowering, they do not appear in the initial HIR map, and information for them cannot be provided in the normal pull-based way.

In order to make those definitions useful, we allow to feed values as query results for the newly created definition.

The API is as follows:
```rust
let feed = tcx.create_def(<parent def id>, <DefPathData>);
// `feed` is a TyCtxtFeed<'tcx>.

// Access the created definition.
let def_id: LocalDefId = feed.def_id;

// Assign `my_query(def_id) := my_value`.
feed.my_query(my_value).
```

This PR keeps the consistency checks introduced by #96840, even if they are not reachable. This allows to extend the behaviour later without forgetting them.

cc `@oli-obk` `@spastorino`

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};
@@ -1031,6 +1032,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
@@ -1493,12 +1509,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`.
@@ -1515,23 +1534,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)*]),