Refactor delegation generics resolution

Author: aerooneqq

compiler/rustc_ast_lowering/src/delegation/generics.rs

@@ -232,123 +232,199 @@ impl<'hir> GenericsGenerationResult<'hir> {
     }
 }
 
+enum ParentSegmentArgs<'a> {
+    /// Parent segment is valid and generic args are specified:
+    /// `reuse Trait::<'static, ()>::foo;`.
+    Specified(&'a AngleBracketedArgs),
+    /// Parent segment is valid and args are not specified:
+    /// `reuse Trait::foo;`.
+    NotSpecified,
+    /// Parent segment doest not exist (`reuse foo`) or we can not
+    /// add generics to it:
+    /// ```rust
+    /// mod to_reuse {
+    ///     fn foo() {}
+    /// }
+    ///
+    /// // Can't add generic args to module.
+    /// reuse to_reuse::foo;
+    /// ```
+    Invalid,
+}
+
+struct GenericsResolutionDto<'a, 'tcx> {
+    trait_impl: bool,
+
+    parent_args: ParentSegmentArgs<'a>,
+    child_args: Option<&'a AngleBracketedArgs>,
+
+    sig_parent_params: &'tcx [ty::GenericParamDef],
+    sig_child_params: &'tcx [ty::GenericParamDef],
+
+    free_to_trait_delegation: bool,
+    /// `reuse Trait::foo;`.
+    qself_is_none: bool,
+    /// `reuse <_ as Trait>::foo;`.
+    qself_is_infer: bool,
+    /// Whether we should generate `Self` generic param.
+    generate_self: bool,
+}
+
 impl<'hir, T: LoweringContextForResolution<'hir>> DelegationResolverForLowering<'_, T> {
-    pub(super) fn resolve_generics(
+    fn create_resolution_dto<'a>(
         &self,
-        delegation: &Delegation,
+        delegation: &'a Delegation,
         sig_id: DefId,
-    ) -> GenericsGenerationResults<'hir> {
+    ) -> GenericsResolutionDto<'a, 'hir> {
         let tcx = self.0.tcx();
-
         let delegation_parent_kind = tcx.def_kind(tcx.local_parent(self.0.owner_id()));
 
-        let segments = &delegation.path.segments;
-        let len = segments.len();
+        let trait_impl = matches!(delegation_parent_kind, DefKind::Impl { of_trait: true });
+
+        let delegation_in_free_ctx =
+            !matches!(delegation_parent_kind, DefKind::Trait | DefKind::Impl { .. });
+
+        let sig_parent = tcx.parent(sig_id);
+        let sig_in_trait = matches!(tcx.def_kind(sig_parent), DefKind::Trait);
+        let free_to_trait_delegation = delegation_in_free_ctx && sig_in_trait;
+
+        let sig_child_params = &tcx.generics_of(sig_id).own_params;
+        let mut sig_parent_params: &[ty::GenericParamDef] = &[];
 
-        let get_user_args = |idx: usize| -> Option<&AngleBracketedArgs> {
-            let segment = &segments[idx];
+        let qself_is_infer =
+            delegation.qself.as_ref().is_some_and(|qself| qself.ty.is_maybe_parenthesised_infer());
 
-            let Some(args) = segment.args.as_ref() else { return None };
-            let GenericArgs::AngleBracketed(args) = args else {
-                tcx.dcx().span_delayed_bug(
-                    segment.span(),
-                    "expected angle-bracketed generic args in delegation segment",
-                );
+        let qself_is_none = delegation.qself.is_none();
+        let generate_self = free_to_trait_delegation && (qself_is_none || qself_is_infer);
 
-                return None;
-            };
+        let parent_args = if let [.., parent_segment, _] = &delegation.path.segments[..] {
+            if let Some(res) = self.0.get_resolution_id(parent_segment.id)
+                && matches!(tcx.def_kind(res), DefKind::Trait | DefKind::TraitAlias)
+            {
+                sig_parent_params = &tcx.generics_of(sig_parent).own_params;
+                self.get_user_args(parent_segment)
+                    .map(|args| ParentSegmentArgs::Specified(args))
+                    .unwrap_or(ParentSegmentArgs::NotSpecified)
+            } else {
+                ParentSegmentArgs::Invalid
+            }
+        } else {
+            ParentSegmentArgs::Invalid
+        };
+
+        let child_segment = delegation.path.segments.last().expect("must be at least one segment");
+        let child_args = self.get_user_args(child_segment);
+
+        GenericsResolutionDto {
+            trait_impl,
+            parent_args,
+            child_args,
+            sig_parent_params,
+            sig_child_params,
+            free_to_trait_delegation,
+            qself_is_none,
+            qself_is_infer,
+            generate_self,
+        }
+    }
+
+    fn get_user_args<'a>(&self, segment: &'a PathSegment) -> Option<&'a AngleBracketedArgs> {
+        let Some(args) = segment.args.as_ref() else { return None };
+        let GenericArgs::AngleBracketed(args) = args else {
+            self.0.tcx().dcx().span_delayed_bug(
+                segment.span(),
+                "expected angle-bracketed generic args in delegation segment",
+            );
 
-            // Treat empty args `reuse foo::<> as bar` as `reuse foo as bar`,
-            // the same logic applied when we call function `fn f<T>(t: T)`
-            // like that `f::<>(())`, in HIR no `<>` will be generated.
-            (!args.args.is_empty()).then(|| args)
+            return None;
         };
 
-        let sig_params = &tcx.generics_of(sig_id).own_params[..];
+        // Treat empty args `reuse foo::<> as bar` as `reuse foo as bar`,
+        // the same logic applied when we call function `fn f<T>(t: T)`
+        // like that `f::<>(())`, in HIR no `<>` will be generated.
+        (!args.args.is_empty()).then(|| args)
+    }
+
+    pub(super) fn resolve_generics(
+        &self,
+        delegation: &Delegation,
+        sig_id: DefId,
+    ) -> GenericsGenerationResults<'hir> {
+        let dto @ GenericsResolutionDto { trait_impl, generate_self, sig_child_params, .. } =
+            self.create_resolution_dto(delegation, sig_id);
 
         // If we are in trait impl always generate function whose generics matches
         // those that are defined in trait.
-        if matches!(delegation_parent_kind, DefKind::Impl { of_trait: true }) {
+        if trait_impl {
             // Considering parent generics, during signature inheritance
             // we will take those args that are in trait impl header trait ref.
             let parent =
                 DelegationGenerics { data: vec![], pos: GenericsPosition::Child, trait_impl: true };
 
             let parent = GenericsGenerationResult::new(parent);
 
-            let child = DelegationGenerics::generate_all(sig_params, GenericsPosition::Child, true);
+            let child = DelegationGenerics::generate_all(
+                dto.sig_child_params,
+                GenericsPosition::Child,
+                true,
+            );
+
             let child = GenericsGenerationResult::new(child);
 
             return GenericsGenerationResults { parent, child, self_ty_propagation_kind: None };
         }
 
-        let delegation_in_free_ctx =
-            !matches!(delegation_parent_kind, DefKind::Trait | DefKind::Impl { .. });
-
-        let sig_parent = tcx.parent(sig_id);
-        let sig_in_trait = matches!(tcx.def_kind(sig_parent), DefKind::Trait);
-        let free_to_trait_delegation = delegation_in_free_ctx && sig_in_trait;
-
-        let qself_is_infer =
-            delegation.qself.as_ref().is_some_and(|qself| qself.ty.is_maybe_parenthesised_infer());
-
-        let qself_is_none = delegation.qself.is_none();
-
-        let generate_self = free_to_trait_delegation && (qself_is_none || qself_is_infer);
-
-        let can_add_generics_to_parent = len >= 2
-            && self.0.get_resolution_id(segments[len - 2].id).is_some_and(|def_id| {
-                matches!(tcx.def_kind(def_id), DefKind::Trait | DefKind::TraitAlias)
-            });
-
-        let parent_generics = if can_add_generics_to_parent {
-            let sig_parent_params = &tcx.generics_of(sig_parent).own_params;
-
-            if let Some(args) = get_user_args(len - 2) {
-                DelegationGenerics {
-                    data: Self::create_slots_from_args(
-                        tcx,
-                        args,
-                        &sig_parent_params[usize::from(!generate_self)..],
-                        generate_self,
-                    ),
-                    pos: GenericsPosition::Parent,
-                    trait_impl: false,
-                }
-            } else {
-                DelegationGenerics::generate_all(
-                    &sig_parent_params[usize::from(!generate_self)..],
-                    GenericsPosition::Parent,
-                    false,
-                )
+        let tcx = self.0.tcx();
+        let parent_generics = match dto.parent_args {
+            ParentSegmentArgs::Specified(args) => DelegationGenerics {
+                data: Self::create_slots_from_args(
+                    tcx,
+                    args,
+                    &dto.sig_parent_params[usize::from(!generate_self)..],
+                    generate_self,
+                ),
+                pos: GenericsPosition::Parent,
+                trait_impl,
+            },
+            ParentSegmentArgs::NotSpecified => DelegationGenerics::generate_all(
+                &dto.sig_parent_params[usize::from(!generate_self)..],
+                GenericsPosition::Parent,
+                trait_impl,
+            ),
+            ParentSegmentArgs::Invalid => {
+                DelegationGenerics { data: vec![], pos: GenericsPosition::Parent, trait_impl }
             }
-        } else {
-            DelegationGenerics { data: vec![], pos: GenericsPosition::Parent, trait_impl: false }
         };
 
-        let child_generics = if let Some(args) = get_user_args(len - 1) {
-            let synth_params_index =
-                sig_params.iter().position(|p| p.kind.is_synthetic()).unwrap_or(sig_params.len());
+        let child_generics = if let Some(args) = dto.child_args {
+            let synth_params_index = sig_child_params
+                .iter()
+                .position(|p| p.kind.is_synthetic())
+                .unwrap_or(sig_child_params.len());
 
-            let mut slots =
-                Self::create_slots_from_args(tcx, args, &sig_params[..synth_params_index], false);
+            let mut slots = Self::create_slots_from_args(
+                tcx,
+                args,
+                &dto.sig_child_params[..synth_params_index],
+                trait_impl,
+            );
 
-            for synth_param in &sig_params[synth_params_index..] {
+            for synth_param in &sig_child_params[synth_params_index..] {
                 slots.push(GenericArgSlot::Generate(synth_param, None));
             }
 
-            DelegationGenerics { data: slots, pos: GenericsPosition::Child, trait_impl: false }
+            DelegationGenerics { data: slots, pos: GenericsPosition::Child, trait_impl }
         } else {
-            DelegationGenerics::generate_all(sig_params, GenericsPosition::Child, false)
+            DelegationGenerics::generate_all(sig_child_params, GenericsPosition::Child, trait_impl)
         };
 
         GenericsGenerationResults {
             parent: GenericsGenerationResult::new(parent_generics),
             child: GenericsGenerationResult::new(child_generics),
-            self_ty_propagation_kind: match free_to_trait_delegation {
-                true => Some(match qself_is_none {
+            self_ty_propagation_kind: match dto.free_to_trait_delegation {
+                true => Some(match dto.qself_is_none {
                     true => hir::DelegationSelfTyPropagationKind::SelfParam,
-                    false => match qself_is_infer {
+                    false => match dto.qself_is_infer {
                         true => hir::DelegationSelfTyPropagationKind::SelfParam,
                         // HirId is filled during generic args propagation.
                         false => hir::DelegationSelfTyPropagationKind::SelfTy(HirId::INVALID),