Use dedicated PreciseCapturingArg for representing what goes in use<>

6 files changed, 146 insertions, 56 deletions

diff --git a/compiler/rustc_ast/src/ast.rs b/compiler/rustc_ast/src/ast.rs
index 55a934c48a8..623a6437c53 100644
--- a/compiler/rustc_ast/src/ast.rs
+++ b/compiler/rustc_ast/src/ast.rs
@@ -2132,7 +2132,7 @@ pub enum TyKind {
     /// The `NodeId` exists to prevent lowering from having to
     /// generate `NodeId`s on the fly, which would complicate
     /// the generation of opaque `type Foo = impl Trait` items significantly.
-    ImplTrait(NodeId, GenericBounds, Option<P<GenericArgs>>),
+    ImplTrait(NodeId, GenericBounds, Option<ThinVec<PreciseCapturingArg>>),
     /// No-op; kept solely so that we can pretty-print faithfully.
     Paren(P<Ty>),
     /// Unused for now.
@@ -2188,6 +2188,14 @@ pub enum TraitObjectSyntax {
     None,
 }
 
+#[derive(Clone, PartialEq, Encodable, Decodable, Debug)]
+pub enum PreciseCapturingArg {
+    /// Lifetime parameter
+    Lifetime(Lifetime),
+    /// Type or const parameter
+    Arg(Ident, NodeId),
+}
+
 /// Inline assembly operand explicit register or register class.
 ///
 /// E.g., `"eax"` as in `asm!("mov eax, 2", out("eax") result)`.
diff --git a/compiler/rustc_ast/src/mut_visit.rs b/compiler/rustc_ast/src/mut_visit.rs
index b8a5eeebf3a..cb625873937 100644
--- a/compiler/rustc_ast/src/mut_visit.rs
+++ b/compiler/rustc_ast/src/mut_visit.rs
@@ -259,6 +259,14 @@ pub trait MutVisitor: Sized {
         noop_visit_param_bound(tpb, self);
     }
 
+    fn visit_precise_capturing_args(&mut self, args: &mut ThinVec<PreciseCapturingArg>) {
+        noop_visit_precise_capturing_args(args, self);
+    }
+
+    fn visit_precise_capturing_arg(&mut self, arg: &mut PreciseCapturingArg) {
+        noop_visit_precise_capturing_arg(arg, self);
+    }
+
     fn visit_mt(&mut self, mt: &mut MutTy) {
         noop_visit_mt(mt, self);
     }
@@ -522,7 +530,7 @@ pub fn noop_visit_ty<T: MutVisitor>(ty: &mut P<Ty>, vis: &mut T) {
             vis.visit_id(id);
             visit_vec(bounds, |bound| vis.visit_param_bound(bound));
             visit_opt(precise_capturing, |precise_capturing| {
-                vis.visit_generic_args(precise_capturing);
+                vis.visit_precise_capturing_args(precise_capturing);
             });
         }
         TyKind::MacCall(mac) => vis.visit_mac_call(mac),
@@ -917,6 +925,27 @@ pub fn noop_visit_param_bound<T: MutVisitor>(pb: &mut GenericBound, vis: &mut T)
     }
 }
 
+pub fn noop_visit_precise_capturing_args<T: MutVisitor>(
+    args: &mut ThinVec<PreciseCapturingArg>,
+    vis: &mut T,
+) {
+    for arg in args {
+        vis.visit_precise_capturing_arg(arg);
+    }
+}
+
+pub fn noop_visit_precise_capturing_arg<T: MutVisitor>(arg: &mut PreciseCapturingArg, vis: &mut T) {
+    match arg {
+        PreciseCapturingArg::Lifetime(lt) => {
+            vis.visit_lifetime(lt);
+        }
+        PreciseCapturingArg::Arg(ident, id) => {
+            vis.visit_ident(ident);
+            vis.visit_id(id);
+        }
+    }
+}
+
 pub fn noop_flat_map_generic_param<T: MutVisitor>(
     mut param: GenericParam,
     vis: &mut T,
diff --git a/compiler/rustc_ast/src/visit.rs b/compiler/rustc_ast/src/visit.rs
index 91740021727..b7e4d31edd7 100644
--- a/compiler/rustc_ast/src/visit.rs
+++ b/compiler/rustc_ast/src/visit.rs
@@ -20,6 +20,7 @@ use rustc_span::Span;
 
 pub use rustc_ast_ir::visit::VisitorResult;
 pub use rustc_ast_ir::{try_visit, visit_opt, walk_list, walk_visitable_list};
+use thin_vec::ThinVec;
 
 #[derive(Copy, Clone, Debug, PartialEq)]
 pub enum AssocCtxt {
@@ -184,6 +185,12 @@ pub trait Visitor<'ast>: Sized {
     fn visit_param_bound(&mut self, bounds: &'ast GenericBound, _ctxt: BoundKind) -> Self::Result {
         walk_param_bound(self, bounds)
     }
+    fn visit_precise_capturing_args(&mut self, args: &'ast ThinVec<PreciseCapturingArg>) {
+        walk_precise_capturing_args(self, args);
+    }
+    fn visit_precise_capturing_arg(&mut self, arg: &'ast PreciseCapturingArg) {
+        walk_precise_capturing_arg(self, arg);
+    }
     fn visit_poly_trait_ref(&mut self, t: &'ast PolyTraitRef) -> Self::Result {
         walk_poly_trait_ref(self, t)
     }
@@ -459,7 +466,7 @@ pub fn walk_ty<'a, V: Visitor<'a>>(visitor: &mut V, typ: &'a Ty) -> V::Result {
         }
         TyKind::ImplTrait(_, bounds, precise_capturing) => {
             walk_list!(visitor, visit_param_bound, bounds, BoundKind::Impl);
-            visit_opt!(visitor, visit_generic_args, precise_capturing);
+            visit_opt!(visitor, visit_precise_capturing_args, precise_capturing);
         }
         TyKind::Typeof(expression) => try_visit!(visitor.visit_anon_const(expression)),
         TyKind::Infer | TyKind::ImplicitSelf | TyKind::Dummy | TyKind::Err(_) => {}
@@ -638,6 +645,29 @@ pub fn walk_param_bound<'a, V: Visitor<'a>>(visitor: &mut V, bound: &'a GenericB
     }
 }
 
+pub fn walk_precise_capturing_args<'a, V: Visitor<'a>>(
+    visitor: &mut V,
+    args: &'a ThinVec<PreciseCapturingArg>,
+) {
+    for arg in args {
+        visitor.visit_precise_capturing_arg(arg);
+    }
+}
+
+pub fn walk_precise_capturing_arg<'a, V: Visitor<'a>>(
+    visitor: &mut V,
+    arg: &'a PreciseCapturingArg,
+) {
+    match arg {
+        PreciseCapturingArg::Lifetime(lt) => {
+            visitor.visit_lifetime(lt, LifetimeCtxt::GenericArg);
+        }
+        PreciseCapturingArg::Arg(ident, _) => {
+            visitor.visit_ident(*ident);
+        }
+    }
+}
+
 pub fn walk_generic_param<'a, V: Visitor<'a>>(
     visitor: &mut V,
     param: &'a GenericParam,
diff --git a/compiler/rustc_ast_lowering/src/lib.rs b/compiler/rustc_ast_lowering/src/lib.rs
index 32b2f7c86ca..731dda39e9a 100644
--- a/compiler/rustc_ast_lowering/src/lib.rs
+++ b/compiler/rustc_ast_lowering/src/lib.rs
@@ -1525,7 +1525,7 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
         bounds: &GenericBounds,
         fn_kind: Option<FnDeclKind>,
         itctx: ImplTraitContext,
-        precise_capturing: Option<&ast::GenericArgs>,
+        precise_capturing: Option<&[ast::PreciseCapturingArg]>,
     ) -> hir::TyKind<'hir> {
         // Make sure we know that some funky desugaring has been going on here.
         // This is a first: there is code in other places like for loop
@@ -1534,58 +1534,55 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
         // frequently opened issues show.
         let opaque_ty_span = self.mark_span_with_reason(DesugaringKind::OpaqueTy, span, None);
 
-        let captured_lifetimes_to_duplicate = if let Some(precise_capturing) = precise_capturing {
-            let ast::GenericArgs::AngleBracketed(precise_capturing) = precise_capturing else {
-                panic!("we only parse angle-bracketed args")
-            };
-            // We'll actually validate these later on; all we need is the list of
-            // lifetimes to duplicate during this portion of lowering.
-            precise_capturing
-                .args
-                .iter()
-                .filter_map(|arg| match arg {
-                    ast::AngleBracketedArg::Arg(ast::GenericArg::Lifetime(lt)) => Some(*lt),
-                    _ => None,
-                })
-                .collect()
-        } else {
-            match origin {
-                hir::OpaqueTyOrigin::TyAlias { .. } => {
-                    // type alias impl trait and associated type position impl trait were
-                    // decided to capture all in-scope lifetimes, which we collect for
-                    // all opaques during resolution.
-                    self.resolver
-                        .take_extra_lifetime_params(opaque_ty_node_id)
-                        .into_iter()
-                        .map(|(ident, id, _)| Lifetime { id, ident })
-                        .collect()
-                }
-                hir::OpaqueTyOrigin::FnReturn(..) => {
-                    if matches!(
-                        fn_kind.expect("expected RPITs to be lowered with a FnKind"),
-                        FnDeclKind::Impl | FnDeclKind::Trait
-                    ) || self.tcx.features().lifetime_capture_rules_2024
-                        || span.at_least_rust_2024()
-                    {
-                        // return-position impl trait in trait was decided to capture all
-                        // in-scope lifetimes, which we collect for all opaques during resolution.
+        let captured_lifetimes_to_duplicate =
+            if let Some(precise_capturing) = precise_capturing_args {
+                // We'll actually validate these later on; all we need is the list of
+                // lifetimes to duplicate during this portion of lowering.
+                precise_capturing
+                    .iter()
+                    .filter_map(|arg| match arg {
+                        ast::PreciseCapturingArg::Lifetime(lt) => Some(*lt),
+                        ast::PreciseCapturingArg::Arg(..) => None,
+                    })
+                    .collect()
+            } else {
+                match origin {
+                    hir::OpaqueTyOrigin::TyAlias { .. } => {
+                        // type alias impl trait and associated type position impl trait were
+                        // decided to capture all in-scope lifetimes, which we collect for
+                        // all opaques during resolution.
                         self.resolver
                             .take_extra_lifetime_params(opaque_ty_node_id)
                             .into_iter()
                             .map(|(ident, id, _)| Lifetime { id, ident })
                             .collect()
-                    } else {
-                        // in fn return position, like the `fn test<'a>() -> impl Debug + 'a`
-                        // example, we only need to duplicate lifetimes that appear in the
-                        // bounds, since those are the only ones that are captured by the opaque.
-                        lifetime_collector::lifetimes_in_bounds(self.resolver, bounds)
+                    }
+                    hir::OpaqueTyOrigin::FnReturn(..) => {
+                        if matches!(
+                            fn_kind.expect("expected RPITs to be lowered with a FnKind"),
+                            FnDeclKind::Impl | FnDeclKind::Trait
+                        ) || self.tcx.features().lifetime_capture_rules_2024
+                            || span.at_least_rust_2024()
+                        {
+                            // return-position impl trait in trait was decided to capture all
+                            // in-scope lifetimes, which we collect for all opaques during resolution.
+                            self.resolver
+                                .take_extra_lifetime_params(opaque_ty_node_id)
+                                .into_iter()
+                                .map(|(ident, id, _)| Lifetime { id, ident })
+                                .collect()
+                        } else {
+                            // in fn return position, like the `fn test<'a>() -> impl Debug + 'a`
+                            // example, we only need to duplicate lifetimes that appear in the
+                            // bounds, since those are the only ones that are captured by the opaque.
+                            lifetime_collector::lifetimes_in_bounds(self.resolver, bounds)
+                        }
+                    }
+                    hir::OpaqueTyOrigin::AsyncFn(..) => {
+                        unreachable!("should be using `lower_async_fn_ret_ty`")
                     }
                 }
-                hir::OpaqueTyOrigin::AsyncFn(..) => {
-                    unreachable!("should be using `lower_async_fn_ret_ty`")
-                }
-            }
-        };
+            };
         debug!(?captured_lifetimes_to_duplicate);
 
         self.lower_opaque_inner(
diff --git a/compiler/rustc_parse/src/parser/ty.rs b/compiler/rustc_parse/src/parser/ty.rs
index 5d6c4f39ae1..62f0602e1fa 100644
--- a/compiler/rustc_parse/src/parser/ty.rs
+++ b/compiler/rustc_parse/src/parser/ty.rs
@@ -1,4 +1,4 @@
-use super::{Parser, PathStyle, TokenType, Trailing};
+use super::{Parser, PathStyle, SeqSep, TokenType, Trailing};
 
 use crate::errors::{
     self, DynAfterMut, ExpectedFnPathFoundFnKeyword, ExpectedMutOrConstInRawPointerType,
@@ -14,7 +14,7 @@ use rustc_ast::util::case::Case;
 use rustc_ast::{
     self as ast, BareFnTy, BoundAsyncness, BoundConstness, BoundPolarity, FnRetTy, GenericBound,
     GenericBounds, GenericParam, Generics, Lifetime, MacCall, MutTy, Mutability, PolyTraitRef,
-    TraitBoundModifiers, TraitObjectSyntax, Ty, TyKind, DUMMY_NODE_ID,
+    PreciseCapturingArg, TraitBoundModifiers, TraitObjectSyntax, Ty, TyKind, DUMMY_NODE_ID,
 };
 use rustc_errors::{Applicability, PResult};
 use rustc_span::symbol::{kw, sym, Ident};
@@ -671,13 +671,10 @@ impl<'a> Parser<'a> {
 
         // parse precise captures, if any.
         let precise_capturing = if self.eat_keyword(kw::Use) {
-            self.expect_lt()?;
             let use_span = self.prev_token.span;
             self.psess.gated_spans.gate(sym::precise_capturing, use_span);
-            let lo = self.token.span;
-            let args = self.parse_angle_args(None)?;
-            self.expect_gt()?;
-            Some(ast::AngleBracketedArgs { args, span: lo.to(self.prev_token.span) }.into())
+            let args = self.parse_precise_capturing_args()?;
+            Some(args)
         } else {
             None
         };
@@ -690,6 +687,25 @@ impl<'a> Parser<'a> {
         Ok(TyKind::ImplTrait(ast::DUMMY_NODE_ID, bounds, precise_capturing))
     }
 
+    fn parse_precise_capturing_args(&mut self) -> PResult<'a, ThinVec<PreciseCapturingArg>> {
+        Ok(self
+            .parse_unspanned_seq(
+                &TokenKind::Lt,
+                &TokenKind::Gt,
+                SeqSep::trailing_allowed(token::Comma),
+                |self_| {
+                    if self_.check_ident() {
+                        Ok(PreciseCapturingArg::Arg(self_.parse_ident().unwrap(), DUMMY_NODE_ID))
+                    } else if self_.check_lifetime() {
+                        Ok(PreciseCapturingArg::Lifetime(self_.expect_lifetime()))
+                    } else {
+                        self_.unexpected_any()
+                    }
+                },
+            )?
+            .0)
+    }
+
     /// Is a `dyn B0 + ... + Bn` type allowed here?
     fn is_explicit_dyn_type(&mut self) -> bool {
         self.check_keyword(kw::Dyn)
diff --git a/compiler/rustc_resolve/src/late.rs b/compiler/rustc_resolve/src/late.rs
index 9e881532311..e4bcdc96f44 100644
--- a/compiler/rustc_resolve/src/late.rs
+++ b/compiler/rustc_resolve/src/late.rs
@@ -1047,10 +1047,20 @@ impl<'a: 'ast, 'ast, 'tcx> Visitor<'ast> for LateResolutionVisitor<'a, '_, 'ast,
         });
         self.diag_metadata.current_function = previous_value;
     }
+
     fn visit_lifetime(&mut self, lifetime: &'ast Lifetime, use_ctxt: visit::LifetimeCtxt) {
         self.resolve_lifetime(lifetime, use_ctxt)
     }
 
+    fn visit_precise_capturing_arg(&mut self, arg: &'ast PreciseCapturingArg) {
+        match arg {
+            PreciseCapturingArg::Lifetime(_) => visit::walk_precise_capturing_arg(self, arg),
+            PreciseCapturingArg::Arg(ident, _) => {
+                todo!("cannot resolve args yet: {ident}");
+            }
+        }
+    }
+
     fn visit_generics(&mut self, generics: &'ast Generics) {
         self.visit_generic_params(&generics.params, self.diag_metadata.current_self_item.is_some());
         for p in &generics.where_clause.predicates {