Lower and resolve precise captures in HIR

compiler/rustc_ast/src/ast.rs

@@ -63,7 +63,7 @@ impl fmt::Debug for Label {
 
 /// A "Lifetime" is an annotation of the scope in which variable
 /// can be used, e.g. `'a` in `&'a i32`.
-#[derive(Clone, Encodable, Decodable, Copy, PartialEq, Eq)]
+#[derive(Clone, Encodable, Decodable, Copy, PartialEq, Eq, Hash)]
 pub struct Lifetime {
     pub id: NodeId,
     pub ident: Ident,

compiler/rustc_ast/src/mut_visit.rs

@@ -259,10 +259,6 @@ 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);
     }
@@ -530,7 +526,9 @@ 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_precise_capturing_args(precise_capturing);
+                for arg in precise_capturing {
+                    vis.visit_precise_capturing_arg(arg);
+                }
             });
         }
         TyKind::MacCall(mac) => vis.visit_mac_call(mac),
@@ -925,15 +923,6 @@ 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) => {

compiler/rustc_ast/src/visit.rs

@@ -20,7 +20,6 @@ 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 {
@@ -185,9 +184,6 @@ 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);
     }
@@ -466,7 +462,11 @@ 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_precise_capturing_args, precise_capturing);
+            if let Some(precise_capturing) = precise_capturing {
+                for arg in precise_capturing {
+                    try_visit!(visitor.visit_precise_capturing_arg(arg));
+                }
+            }
         }
         TyKind::Typeof(expression) => try_visit!(visitor.visit_anon_const(expression)),
         TyKind::Infer | TyKind::ImplicitSelf | TyKind::Dummy | TyKind::Err(_) => {}
@@ -645,15 +645,6 @@ 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,

compiler/rustc_ast_lowering/src/lib.rs

@@ -48,6 +48,7 @@ use rustc_ast::{self as ast, *};
 use rustc_ast_pretty::pprust;
 use rustc_data_structures::captures::Captures;
 use rustc_data_structures::fingerprint::Fingerprint;
+use rustc_data_structures::fx::FxIndexSet;
 use rustc_data_structures::sorted_map::SortedMap;
 use rustc_data_structures::stable_hasher::{HashStable, StableHasher};
 use rustc_data_structures::sync::Lrc;
@@ -1525,7 +1526,7 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
         bounds: &GenericBounds,
         fn_kind: Option<FnDeclKind>,
         itctx: ImplTraitContext,
-        precise_capturing: Option<&[ast::PreciseCapturingArg]>,
+        precise_capturing_args: Option<&[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
@@ -1541,9 +1542,13 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
                 precise_capturing
                     .iter()
                     .filter_map(|arg| match arg {
-                        ast::PreciseCapturingArg::Lifetime(lt) => Some(*lt),
+                        PreciseCapturingArg::Lifetime(lt) => Some(*lt),
-                        ast::PreciseCapturingArg::Arg(..) => None,
+                        PreciseCapturingArg::Arg(..) => None,
                     })
+                    // Add in all the lifetimes mentioned in the bounds. We will error
+                    // them out later, but capturing them here is important to make sure
+                    // they actually get resolved in resolve_bound_vars.
+                    .chain(lifetime_collector::lifetimes_in_bounds(self.resolver, bounds))
                     .collect()
             } else {
                 match origin {
@@ -1592,6 +1597,7 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
             captured_lifetimes_to_duplicate,
             span,
             opaque_ty_span,
+            precise_capturing_args,
             |this| this.lower_param_bounds(bounds, itctx),
         )
     }
@@ -1601,9 +1607,10 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
         opaque_ty_node_id: NodeId,
         origin: hir::OpaqueTyOrigin,
         in_trait: bool,
-        captured_lifetimes_to_duplicate: Vec<Lifetime>,
+        captured_lifetimes_to_duplicate: FxIndexSet<Lifetime>,
         span: Span,
         opaque_ty_span: Span,
+        precise_capturing_args: Option<&[PreciseCapturingArg]>,
         lower_item_bounds: impl FnOnce(&mut Self) -> &'hir [hir::GenericBound<'hir>],
     ) -> hir::TyKind<'hir> {
         let opaque_ty_def_id = self.create_def(
@@ -1690,8 +1697,15 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
             // Install the remapping from old to new (if any). This makes sure that
             // any lifetimes that would have resolved to the def-id of captured
             // lifetimes are remapped to the new *synthetic* lifetimes of the opaque.
-            let bounds = this
+            let (bounds, precise_capturing_args) =
-                .with_remapping(captured_to_synthesized_mapping, |this| lower_item_bounds(this));
+                this.with_remapping(captured_to_synthesized_mapping, |this| {
+                    (
+                        lower_item_bounds(this),
+                        precise_capturing_args.map(|precise_capturing| {
+                            this.lower_precise_capturing_args(precise_capturing)
+                        }),
+                    )
+                });
 
             let generic_params =
                 this.arena.alloc_from_iter(synthesized_lifetime_definitions.iter().map(
@@ -1736,6 +1750,7 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
                 origin,
                 lifetime_mapping,
                 in_trait,
+                precise_capturing_args,
             };
 
             // Generate an `type Foo = impl Trait;` declaration.
@@ -1768,6 +1783,23 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
         )
     }
 
+    fn lower_precise_capturing_args(
+        &mut self,
+        precise_capturing_args: &[PreciseCapturingArg],
+    ) -> &'hir [hir::PreciseCapturingArg<'hir>] {
+        self.arena.alloc_from_iter(precise_capturing_args.iter().map(|arg| match arg {
+            PreciseCapturingArg::Lifetime(lt) => {
+                hir::PreciseCapturingArg::Lifetime(self.lower_lifetime(lt))
+            }
+            PreciseCapturingArg::Arg(_, node_id) => {
+                let res = self.resolver.get_partial_res(*node_id).map_or(Res::Err, |partial_res| {
+                    partial_res.full_res().expect("no partial res expected for precise capture arg")
+                });
+                hir::PreciseCapturingArg::Param(self.lower_res(res), self.lower_node_id(*node_id))
+            }
+        }))
+    }
+
     fn lower_fn_params_to_names(&mut self, decl: &FnDecl) -> &'hir [Ident] {
         self.arena.alloc_from_iter(decl.inputs.iter().map(|param| match param.pat.kind {
             PatKind::Ident(_, ident, _) => self.lower_ident(ident),
@@ -1908,7 +1940,7 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
         let opaque_ty_span =
             self.mark_span_with_reason(DesugaringKind::Async, span, allowed_features);
 
-        let captured_lifetimes: Vec<_> = self
+        let captured_lifetimes = self
             .resolver
             .take_extra_lifetime_params(opaque_ty_node_id)
             .into_iter()
@@ -1922,6 +1954,7 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
             captured_lifetimes,
             span,
             opaque_ty_span,
+            None,
             |this| {
                 let bound = this.lower_coroutine_fn_output_type_to_bound(
                     output,

compiler/rustc_ast_lowering/src/lifetime_collector.rs

@@ -1,6 +1,7 @@
 use super::ResolverAstLoweringExt;
 use rustc_ast::visit::{self, BoundKind, LifetimeCtxt, Visitor};
 use rustc_ast::{GenericBounds, Lifetime, NodeId, PathSegment, PolyTraitRef, Ty, TyKind};
+use rustc_data_structures::fx::FxIndexSet;
 use rustc_hir::def::{DefKind, LifetimeRes, Res};
 use rustc_middle::span_bug;
 use rustc_middle::ty::ResolverAstLowering;
@@ -10,27 +11,23 @@ use rustc_span::Span;
 struct LifetimeCollectVisitor<'ast> {
     resolver: &'ast ResolverAstLowering,
     current_binders: Vec<NodeId>,
-    collected_lifetimes: Vec<Lifetime>,
+    collected_lifetimes: FxIndexSet<Lifetime>,
 }
 
 impl<'ast> LifetimeCollectVisitor<'ast> {
     fn new(resolver: &'ast ResolverAstLowering) -> Self {
-        Self { resolver, current_binders: Vec::new(), collected_lifetimes: Vec::new() }
+        Self { resolver, current_binders: Vec::new(), collected_lifetimes: FxIndexSet::default() }
     }
 
     fn record_lifetime_use(&mut self, lifetime: Lifetime) {
         match self.resolver.get_lifetime_res(lifetime.id).unwrap_or(LifetimeRes::Error) {
             LifetimeRes::Param { binder, .. } | LifetimeRes::Fresh { binder, .. } => {
                 if !self.current_binders.contains(&binder) {
-                    if !self.collected_lifetimes.contains(&lifetime) {
+                    self.collected_lifetimes.insert(lifetime);
-                        self.collected_lifetimes.push(lifetime);
-                    }
                 }
             }
             LifetimeRes::Static | LifetimeRes::Error => {
-                if !self.collected_lifetimes.contains(&lifetime) {
+                self.collected_lifetimes.insert(lifetime);
-                    self.collected_lifetimes.push(lifetime);
-                }
             }
             LifetimeRes::Infer => {}
             res => {
@@ -111,7 +108,7 @@ impl<'ast> Visitor<'ast> for LifetimeCollectVisitor<'ast> {
 pub(crate) fn lifetimes_in_bounds(
     resolver: &ResolverAstLowering,
     bounds: &GenericBounds,
-) -> Vec<Lifetime> {
+) -> FxIndexSet<Lifetime> {
     let mut visitor = LifetimeCollectVisitor::new(resolver);
     for bound in bounds {
         visitor.visit_param_bound(bound, BoundKind::Bound);

compiler/rustc_hir/src/hir.rs

@@ -2557,6 +2557,15 @@ pub struct OpaqueTy<'hir> {
     /// originating from a trait method. This makes it so that the opaque is
     /// lowered as an associated type.
     pub in_trait: bool,
+    /// List of arguments captured via `impl use<'a, P, ...> Trait` syntax.
+    pub precise_capturing_args: Option<&'hir [PreciseCapturingArg<'hir>]>,
+}
+
+#[derive(Debug, Clone, Copy, HashStable_Generic)]
+pub enum PreciseCapturingArg<'hir> {
+    Lifetime(&'hir Lifetime),
+    /// Non-lifetime argument (type or const)
+    Param(Res, HirId),
 }
 
 /// From whence the opaque type came.

compiler/rustc_hir/src/intravisit.rs

@@ -413,6 +413,9 @@ pub trait Visitor<'v>: Sized {
     fn visit_param_bound(&mut self, bounds: &'v GenericBound<'v>) -> Self::Result {
         walk_param_bound(self, bounds)
     }
+    fn visit_precise_capturing_arg(&mut self, arg: &'v PreciseCapturingArg<'v>) -> Self::Result {
+        walk_precise_capturing_arg(self, arg)
+    }
     fn visit_poly_trait_ref(&mut self, t: &'v PolyTraitRef<'v>) -> Self::Result {
         walk_poly_trait_ref(self, t)
     }
@@ -526,10 +529,15 @@ pub fn walk_item<'v, V: Visitor<'v>>(visitor: &mut V, item: &'v Item<'v>) -> V::
             try_visit!(visitor.visit_ty(ty));
             try_visit!(visitor.visit_generics(generics));
         }
-        ItemKind::OpaqueTy(&OpaqueTy { generics, bounds, .. }) => {
+        ItemKind::OpaqueTy(&OpaqueTy { generics, bounds, precise_capturing_args, .. }) => {
             try_visit!(visitor.visit_id(item.hir_id()));
             try_visit!(walk_generics(visitor, generics));
             walk_list!(visitor, visit_param_bound, bounds);
+            if let Some(precise_capturing_args) = precise_capturing_args {
+                for arg in precise_capturing_args {
+                    try_visit!(visitor.visit_precise_capturing_arg(arg));
+                }
+            }
         }
         ItemKind::Enum(ref enum_definition, ref generics) => {
             try_visit!(visitor.visit_generics(generics));
@@ -1137,6 +1145,16 @@ pub fn walk_param_bound<'v, V: Visitor<'v>>(
     }
 }
 
+pub fn walk_precise_capturing_arg<'v, V: Visitor<'v>>(
+    visitor: &mut V,
+    arg: &'v PreciseCapturingArg<'v>,
+) -> V::Result {
+    match *arg {
+        PreciseCapturingArg::Lifetime(lt) => visitor.visit_lifetime(lt),
+        PreciseCapturingArg::Param(_, hir_id) => visitor.visit_id(hir_id),
+    }
+}
+
 pub fn walk_poly_trait_ref<'v, V: Visitor<'v>>(
     visitor: &mut V,
     trait_ref: &'v PolyTraitRef<'v>,

compiler/rustc_hir_analysis/src/collect/resolve_bound_vars.rs

@@ -557,6 +557,31 @@ impl<'a, 'tcx> Visitor<'tcx> for BoundVarContext<'a, 'tcx> {
         }
     }
 
+    fn visit_precise_capturing_arg(
+        &mut self,
+        arg: &'tcx hir::PreciseCapturingArg<'tcx>,
+    ) -> Self::Result {
+        match *arg {
+            hir::PreciseCapturingArg::Lifetime(lt) => match lt.res {
+                LifetimeName::Param(def_id) => {
+                    self.resolve_lifetime_ref(def_id, lt);
+                }
+                LifetimeName::Error => {}
+                LifetimeName::ImplicitObjectLifetimeDefault
+                | LifetimeName::Infer
+                | LifetimeName::Static => todo!("TODO: Error on invalid lifetime"),
+            },
+            hir::PreciseCapturingArg::Param(res, hir_id) => match res {
+                Res::Def(DefKind::TyParam | DefKind::ConstParam, def_id)
+                | Res::SelfTyParam { trait_: def_id } => {
+                    self.resolve_type_ref(def_id.expect_local(), hir_id);
+                }
+                Res::Err => {}
+                _ => todo!("TODO: Error on invalid param res"),
+            },
+        }
+    }
+
     fn visit_foreign_item(&mut self, item: &'tcx hir::ForeignItem<'tcx>) {
         match item.kind {
             hir::ForeignItemKind::Fn(_, _, generics) => {