delegation: split resolution and lowering

compiler/rustc_ast/src/ast.rs

@@ -3921,6 +3921,12 @@ pub struct Delegation {
     pub from_glob: bool,
 }
 
+impl Delegation {
+    pub fn last_segment_span(&self) -> Span {
+        self.path.segments.last().unwrap().ident.span
+    }
+}
+
 #[derive(Clone, Encodable, Decodable, Debug, Walkable)]
 pub enum DelegationSuffixes {
     List(ThinVec<(Ident, Option<Ident>)>),

compiler/rustc_ast_lowering/src/delegation.rs

@@ -44,22 +44,20 @@ use hir::{BodyId, HirId};
 use rustc_abi::ExternAbi;
 use rustc_ast as ast;
 use rustc_ast::*;
-use rustc_data_structures::fx::FxHashSet;
-use rustc_errors::ErrorGuaranteed;
+use rustc_data_structures::fx::{FxHashSet, FxIndexMap};
 use rustc_hir::attrs::{AttributeKind, InlineAttr};
-use rustc_hir::def_id::DefId;
+use rustc_hir::def_id::{DefId, LocalDefId};
 use rustc_hir::{self as hir, FnDeclFlags};
 use rustc_middle::span_bug;
-use rustc_middle::ty::Asyncness;
+use rustc_middle::ty::{Asyncness, TyCtxt};
 use rustc_span::symbol::kw;
-use rustc_span::{Ident, Span, Symbol};
-use smallvec::SmallVec;
+use rustc_span::{ErrorGuaranteed, Ident, Span, Symbol};
 
 use crate::delegation::generics::{GenericsGenerationResult, GenericsGenerationResults};
 use crate::errors::{CycleInDelegationSignatureResolution, UnresolvedDelegationCallee};
 use crate::{
     AllowReturnTypeNotation, ImplTraitContext, ImplTraitPosition, LoweringContext, ParamMode,
-    ResolverAstLoweringExt,
+    ResolverAstLoweringExt, index_crate,
 };
 
 mod generics;
@@ -105,6 +103,66 @@ static ATTRS_ADDITIONS: &[AttrAdditionInfo] = &[
     },
 ];
 
+pub(crate) fn delegations_resolutions(
+    tcx: TyCtxt<'_>,
+    _: (),
+) -> FxIndexMap<LocalDefId, Result<DefId, ErrorGuaranteed>> {
+    let krate = tcx.hir_crate(());
+
+    let (resolver, ast_crate) = &*krate.delayed_resolver.borrow();
+
+    // FIXME!!!(fn_delegation): make ast index lifetime same as resolver,
+    // as it is too bad to reindex whole crate on each delegation lowering.
+    let ast_index = index_crate(resolver, ast_crate);
+
+    let mut result = FxIndexMap::<LocalDefId, Result<DefId, ErrorGuaranteed>>::default();
+
+    for &def_id in &krate.delayed_ids {
+        let delegation = ast_index[def_id].delegation().expect("processing delegations");
+        let span = delegation.last_segment_span();
+
+        if let Some(info) = resolver.delegation_info(def_id) {
+            let res = info.resolution_id.map(|id| check_for_cycles(tcx, id, span).map(|_| id));
+            result.insert(def_id, res.flatten());
+        } else {
+            tcx.dcx().span_delayed_bug(
+                span,
+                format!("delegation resolution record was not found for {def_id:?}"),
+            );
+        }
+    }
+
+    result
+}
+
+fn check_for_cycles(tcx: TyCtxt<'_>, mut def_id: DefId, span: Span) -> Result<(), ErrorGuaranteed> {
+    let mut visited: FxHashSet<DefId> = Default::default();
+
+    let (resolver, _) = &*tcx.hir_crate(()).delayed_resolver.borrow();
+
+    loop {
+        visited.insert(def_id);
+
+        // If def_id is in local crate and it corresponds to another delegation
+        // it means that we refer to another delegation as a callee, so in order to obtain
+        // a signature DefId we obtain NodeId of the callee delegation and try to get signature from it.
+        if let Some(local_id) = def_id.as_local()
+            && let Some(info) = resolver.delegation_info(local_id)
+            && let Ok(id) = info.resolution_id
+        {
+            def_id = id;
+            if visited.contains(&def_id) {
+                return Err(match visited.len() {
+                    1 => tcx.dcx().emit_err(UnresolvedDelegationCallee { span }),
+                    _ => tcx.dcx().emit_err(CycleInDelegationSignatureResolution { span }),
+                });
+            }
+        } else {
+            return Ok(());
+        }
+    }
+}
+
 impl<'hir> LoweringContext<'_, 'hir> {
     fn is_method(&self, def_id: DefId, span: Span) -> bool {
         match self.tcx.def_kind(def_id) {
@@ -119,13 +177,13 @@ impl<'hir> LoweringContext<'_, 'hir> {
         delegation: &Delegation,
         item_id: NodeId,
     ) -> DelegationResults<'hir> {
-        let span = self.lower_span(delegation.path.segments.last().unwrap().ident.span);
+        let span = self.lower_span(delegation.last_segment_span());
 
-        // Delegation can be unresolved in illegal places such as function bodies in extern blocks (see #151356)
-        let sig_id = if let Some(delegation_info) = self.resolver.delegation_info(self.owner.def_id)
-        {
-            self.get_sig_id(delegation_info.resolution_id, span)
-        } else {
+        let sig_id = self.tcx.delegations_resolutions(()).get(&self.owner.def_id).copied();
+
+        // Delegation can be missing from the `delegations_resolutions` table
+        // in illegal places such as function bodies in extern blocks (see #151356).
+        let Some(Ok(sig_id)) = sig_id else {
             self.dcx().span_delayed_bug(
                 span,
                 format!("LoweringContext: the delegation {:?} is unresolved", item_id),
@@ -134,49 +192,39 @@ impl<'hir> LoweringContext<'_, 'hir> {
             return self.generate_delegation_error(span, delegation);
         };
 
-        match sig_id {
-            Ok(sig_id) => {
-                self.add_attrs_if_needed(span, sig_id);
+        self.add_attrs_if_needed(span, sig_id);
 
-                let is_method = self.is_method(sig_id, span);
+        let is_method = self.is_method(sig_id, span);
 
-                let (param_count, c_variadic) = self.param_count(sig_id);
+        let (param_count, c_variadic) = self.param_count(sig_id);
 
-                let mut generics = self.uplift_delegation_generics(delegation, sig_id, is_method);
+        let mut generics = self.uplift_delegation_generics(delegation, sig_id, is_method);
 
-                let (body_id, call_expr_id) = self.lower_delegation_body(
-                    delegation,
-                    is_method,
-                    param_count,
-                    &mut generics,
-                    span,
-                );
+        let (body_id, call_expr_id) =
+            self.lower_delegation_body(delegation, is_method, param_count, &mut generics, span);
 
-                let decl = self.lower_delegation_decl(
-                    sig_id,
-                    param_count,
-                    c_variadic,
-                    span,
-                    &generics,
-                    delegation.id,
-                    call_expr_id,
-                );
+        let decl = self.lower_delegation_decl(
+            sig_id,
+            param_count,
+            c_variadic,
+            span,
+            &generics,
+            delegation.id,
+            call_expr_id,
+        );
 
-                let sig = self.lower_delegation_sig(sig_id, decl, span);
-                let ident = self.lower_ident(delegation.ident);
+        let sig = self.lower_delegation_sig(sig_id, decl, span);
+        let ident = self.lower_ident(delegation.ident);
 
-                let generics = self.arena.alloc(hir::Generics {
-                    has_where_clause_predicates: false,
-                    params: self.arena.alloc_from_iter(generics.all_params()),
-                    predicates: self.arena.alloc_from_iter(generics.all_predicates()),
-                    span,
-                    where_clause_span: span,
-                });
+        let generics = self.arena.alloc(hir::Generics {
+            has_where_clause_predicates: false,
+            params: self.arena.alloc_from_iter(generics.all_params()),
+            predicates: self.arena.alloc_from_iter(generics.all_predicates()),
+            span,
+            where_clause_span: span,
+        });
 
-                DelegationResults { body_id, sig, ident, generics }
-            }
-            Err(_) => self.generate_delegation_error(span, delegation),
-        }
+        DelegationResults { body_id, sig, ident, generics }
     }
 
     fn add_attrs_if_needed(&mut self, span: Span, sig_id: DefId) {
@@ -230,36 +278,6 @@ impl<'hir> LoweringContext<'_, 'hir> {
             .collect::<Vec<_>>()
     }
 
-    fn get_sig_id(&self, mut def_id: DefId, span: Span) -> Result<DefId, ErrorGuaranteed> {
-        let mut visited: FxHashSet<DefId> = Default::default();
-        let mut path: SmallVec<[DefId; 1]> = Default::default();
-
-        loop {
-            visited.insert(def_id);
-
-            path.push(def_id);
-
-            // If def_id is in local crate and it corresponds to another delegation
-            // it means that we refer to another delegation as a callee, so in order to obtain
-            // a signature DefId we obtain NodeId of the callee delegation and try to get signature from it.
-            if let Some(local_id) = def_id.as_local()
-                && let Some(delegation_info) = self.resolver.delegation_info(local_id)
-            {
-                def_id = delegation_info.resolution_id;
-                if visited.contains(&def_id) {
-                    // We encountered a cycle in the resolution, or delegation callee refers to non-existent
-                    // entity, in this case emit an error.
-                    return Err(match visited.len() {
-                        1 => self.dcx().emit_err(UnresolvedDelegationCallee { span }),
-                        _ => self.dcx().emit_err(CycleInDelegationSignatureResolution { span }),
-                    });
-                }
-            } else {
-                return Ok(path[0]);
-            }
-        }
-    }
-
     fn get_resolution_id(&self, node_id: NodeId) -> Option<DefId> {
         self.get_partial_res(node_id).and_then(|r| r.expect_full_res().opt_def_id())
     }

compiler/rustc_ast_lowering/src/lib.rs

@@ -60,6 +60,7 @@ use rustc_hir::{
 use rustc_index::{Idx, IndexSlice, IndexVec};
 use rustc_macros::extension;
 use rustc_middle::hir::{self as mid_hir};
+use rustc_middle::queries::Providers;
 use rustc_middle::span_bug;
 use rustc_middle::ty::{DelegationInfo, PerOwnerResolverData, ResolverAstLowering, TyCtxt};
 use rustc_session::errors::add_feature_diagnostics;
@@ -91,6 +92,12 @@ mod pat;
 mod path;
 pub mod stability;
 
+pub fn provide(providers: &mut Providers) {
+    providers.hir_crate = lower_to_hir;
+    providers.lower_delayed_owner = lower_delayed_owner;
+    providers.delegations_resolutions = delegation::delegations_resolutions;
+}
+
 struct LoweringContext<'a, 'hir> {
     tcx: TyCtxt<'hir>,
     resolver: &'a ResolverAstLowering<'hir>,
@@ -433,6 +440,16 @@ enum AstOwner<'a> {
     ForeignItem(&'a ast::ForeignItem),
 }
 
+impl AstOwner<'_> {
+    fn delegation(&self) -> Option<&ast::Delegation> {
+        match self {
+            AstOwner::Item(Item { kind: ItemKind::Delegation(d), .. })
+            | AstOwner::AssocItem(Item { kind: AssocItemKind::Delegation(d), .. }, _) => Some(d),
+            _ => None,
+        }
+    }
+}
+
 #[derive(Copy, Clone, Debug)]
 enum TryBlockScope {
     /// There isn't a `try` block, so a `?` will use `return`.
@@ -512,7 +529,7 @@ fn compute_hir_hash(
     })
 }
 
-pub fn lower_to_hir(tcx: TyCtxt<'_>, (): ()) -> mid_hir::Crate<'_> {
+fn lower_to_hir(tcx: TyCtxt<'_>, (): ()) -> mid_hir::Crate<'_> {
     // Queries that borrow `resolver_for_lowering`.
     tcx.ensure_done().output_filenames(());
     tcx.ensure_done().early_lint_checks(());
@@ -536,13 +553,11 @@ pub fn lower_to_hir(tcx: TyCtxt<'_>, (): ()) -> mid_hir::Crate<'_> {
     let mut delayed_ids: FxIndexSet<LocalDefId> = Default::default();
 
     for def_id in ast_index.indices() {
-        match &ast_index[def_id] {
-            AstOwner::Item(Item { kind: ItemKind::Delegation { .. }, .. })
-            | AstOwner::AssocItem(Item { kind: AssocItemKind::Delegation { .. }, .. }, _) => {
-                delayed_ids.insert(def_id);
-            }
-            _ => lowerer.lower_node(def_id),
-        };
+        if ast_index[def_id].delegation().is_some() {
+            delayed_ids.insert(def_id);
+        } else {
+            lowerer.lower_node(def_id);
+        }
     }
 
     // Don't hash unless necessary, because it's expensive.
@@ -554,7 +569,7 @@ pub fn lower_to_hir(tcx: TyCtxt<'_>, (): ()) -> mid_hir::Crate<'_> {
 }
 
 /// Lowers an AST owner corresponding to `def_id`, now only delegations are lowered this way.
-pub fn lower_delayed_owner(tcx: TyCtxt<'_>, def_id: LocalDefId) {
+fn lower_delayed_owner(tcx: TyCtxt<'_>, def_id: LocalDefId) {
     let krate = tcx.hir_crate(());
 
     let (resolver, krate) = &*krate.delayed_resolver.borrow();

compiler/rustc_interface/src/passes.rs

@@ -877,8 +877,6 @@ pub fn write_interface<'tcx>(tcx: TyCtxt<'tcx>) {
 pub static DEFAULT_QUERY_PROVIDERS: LazyLock<Providers> = LazyLock::new(|| {
     let providers = &mut Providers::default();
     providers.queries.analysis = analysis;
-    providers.queries.hir_crate = rustc_ast_lowering::lower_to_hir;
-    providers.queries.lower_delayed_owner = rustc_ast_lowering::lower_delayed_owner;
     // `hir_delayed_owner` is fed during `lower_delayed_owner`, by default it returns phantom,
     // as if this query was not fed it means that `MaybeOwner` does not exist for provided LocalDefId.
     providers.queries.hir_delayed_owner = |_, _| MaybeOwner::Phantom;
@@ -887,6 +885,7 @@ pub static DEFAULT_QUERY_PROVIDERS: LazyLock<Providers> = LazyLock::new(|| {
     providers.queries.resolutions = |tcx, ()| tcx.resolver_for_lowering_raw(()).1;
     providers.queries.early_lint_checks = early_lint_checks;
     providers.queries.env_var_os = env_var_os;
+    rustc_ast_lowering::provide(&mut providers.queries);
     limits::provide(&mut providers.queries);
     proc_macro_decls::provide(&mut providers.queries);
     rustc_expand::provide(&mut providers.queries);

compiler/rustc_middle/src/queries.rs

@@ -2069,6 +2069,12 @@ rustc_queries! {
         desc { "getting delegation user-specified args" }
     }
 
+    query delegations_resolutions(_: ()) -> &'tcx FxIndexMap<LocalDefId, Result<DefId, ErrorGuaranteed>> {
+        arena_cache
+        eval_always
+        desc { "getting delegations resolutions" }
+    }
+
     /// Does lifetime resolution on items. Importantly, we can't resolve
     /// lifetimes directly on things like trait methods, because of trait params.
     /// See `rustc_resolve::late::lifetimes` for details.

compiler/rustc_middle/src/ty/mod.rs

@@ -270,7 +270,7 @@ pub struct DelegationInfo {
     /// Refers to the next element in a delegation resolution chain.
     /// Usually points to the final resolution, as most "chains" are just
     /// one step to a trait or an impl.
-    pub resolution_id: DefId,
+    pub resolution_id: Result<DefId, ErrorGuaranteed>,
 }
 
 #[derive(Clone, Copy, Debug, StableHash)]

compiler/rustc_resolve/src/late.rs

@@ -3899,24 +3899,24 @@ impl<'a, 'ast, 'ra, 'tcx> LateResolutionVisitor<'a, 'ast, 'ra, 'tcx> {
             this.visit_path(&delegation.path);
         });
 
-        let resolution_id = if is_in_trait_impl { item_id } else { delegation.id };
+        let resolution_node_id = if is_in_trait_impl { item_id } else { delegation.id };
         let def_id = self
             .r
             .partial_res_map
-            .get(&resolution_id)
+            .get(&resolution_node_id)
             .and_then(|r| r.expect_full_res().opt_def_id());
-        if let Some(resolution_id) = def_id {
-            self.r
-                .delegation_infos
-                .insert(self.r.current_owner.def_id, DelegationInfo { resolution_id });
-        } else {
+
+        let resolution_id = def_id.ok_or_else(|| {
             self.r.tcx.dcx().span_delayed_bug(
                 delegation.path.span,
                 format!(
-                    "LoweringContext: couldn't resolve node {resolution_id:?} in delegation item",
+                    "LateResolutionVisitor: couldn't resolve node {resolution_node_id:?} in delegation item",
                 ),
-            );
-        };
+            )
+        });
+
+        let info = DelegationInfo { resolution_id };
+        self.r.delegation_infos.insert(self.r.current_owner.def_id, info);
 
         let Some(body) = &delegation.body else { return };
         self.with_rib(ValueNS, RibKind::FnOrCoroutine, |this| {