Implement def_id based remapping

spastorino · spastorino · commit 1d6cebfd6bed · 2022-08-04T11:26:57.000-03:00
diff --git a/compiler/rustc_ast_lowering/src/lib.rs b/compiler/rustc_ast_lowering/src/lib.rs
@@ -151,6 +151,8 @@ trait ResolverAstLoweringExt {
     fn get_lifetime_res(&self, id: NodeId) -> Option<LifetimeRes>;
     fn take_extra_lifetime_params(&mut self, id: NodeId) -> Vec<(Ident, NodeId, LifetimeRes)>;
     fn decl_macro_kind(&self, def_id: LocalDefId) -> MacroKind;
+    fn record_def_id_remap(&mut self, from: LocalDefId, to: LocalDefId);
+    fn get_remapped_def_id(&self, local_def_id: LocalDefId) -> LocalDefId;
 }
 
 impl ResolverAstLoweringExt for ResolverAstLowering {
@@ -218,6 +220,25 @@ impl ResolverAstLoweringExt for ResolverAstLowering {
     fn decl_macro_kind(&self, def_id: LocalDefId) -> MacroKind {
         self.builtin_macro_kinds.get(&def_id).copied().unwrap_or(MacroKind::Bang)
     }
+
+    /// Push a remapping into the top-most map. Panics if no map has been pushed.
+    #[tracing::instrument(level = "debug", skip(self))]
+    fn record_def_id_remap(&mut self, from: LocalDefId, to: LocalDefId) {
+        self.generics_def_id_map.last_mut().expect("no map pushed").insert(from, to);
+    }
+
+    fn get_remapped_def_id(&self, mut local_def_id: LocalDefId) -> LocalDefId {
+        for map in &self.generics_def_id_map {
+            if let Some(r) = map.get(&local_def_id) {
+                debug!("def_id_remapper: remapping from `{local_def_id:?}` to `{r:?}`");
+                local_def_id = *r;
+            } else {
+                debug!("def_id_remapper: no remapping for `{local_def_id:?}` found in map");
+            }
+        }
+
+        local_def_id
+    }
 }
 
 /// Context of `impl Trait` in code, which determines whether it is allowed in an HIR subtree,
@@ -474,7 +495,10 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
     }
 
     fn opt_local_def_id(&self, node: NodeId) -> Option<LocalDefId> {
-        self.resolver.node_id_to_def_id.get(&node).copied()
+        self.resolver
+            .node_id_to_def_id
+            .get(&node)
+            .map(|local_def_id| self.resolver.get_remapped_def_id(*local_def_id))
     }
 
     fn local_def_id(&self, node: NodeId) -> LocalDefId {
@@ -534,6 +558,17 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
         debug_assert!(_old.is_none())
     }
 
+    fn with_remapping<R>(
+        &mut self,
+        remap: FxHashMap<LocalDefId, LocalDefId>,
+        f: impl FnOnce(&mut Self) -> R,
+    ) -> R {
+        self.resolver.generics_def_id_map.push(remap);
+        let res = f(self);
+        self.resolver.generics_def_id_map.pop();
+        res
+    }
+
     fn make_owner_info(&mut self, node: hir::OwnerNode<'hir>) -> &'hir hir::OwnerInfo<'hir> {
         let attrs = std::mem::take(&mut self.attrs);
         let mut bodies = std::mem::take(&mut self.bodies);
@@ -1325,9 +1360,7 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
         let mut new_remapping = FxHashMap::default();
 
         self.with_hir_id_owner(opaque_ty_node_id, |lctx| {
-            let hir_bounds = if origin == hir::OpaqueTyOrigin::TyAlias {
-                lctx.lower_param_bounds(bounds, itctx)
-            } else {
+            if origin != hir::OpaqueTyOrigin::TyAlias {
                 debug!(?lctx.captured_lifetimes);
 
                 let lifetime_stash = std::mem::replace(
@@ -1347,53 +1380,57 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
                     &mut new_remapping,
                 );
 
-                let ret = lctx.lower_param_bounds(bounds, itctx);
-
                 let ctxt = std::mem::replace(&mut lctx.captured_lifetimes, lifetime_stash).unwrap();
 
                 collected_lifetimes = ctxt.captures;
-
-                ret
             };
+            debug!(?new_remapping);
             debug!(?collected_lifetimes);
 
-            let lifetime_defs =
-                lctx.arena.alloc_from_iter(collected_lifetimes.iter().map(|&(lifetime, _)| {
-                    let hir_id = lctx.lower_node_id(lifetime.id);
-                    debug_assert_ne!(lctx.opt_local_def_id(lifetime.id), None);
+            lctx.with_remapping(new_remapping, |lctx| {
+                let hir_bounds = lctx.lower_param_bounds(bounds, itctx);
+
+                let lifetime_defs =
+                    lctx.arena.alloc_from_iter(collected_lifetimes.iter().map(|&(lifetime, _)| {
+                        let hir_id = lctx.lower_node_id(lifetime.id);
+                        debug_assert_ne!(lctx.opt_local_def_id(lifetime.id), None);
+
+                        let (name, kind) = if lifetime.ident.name == kw::UnderscoreLifetime {
+                            (hir::ParamName::Fresh, hir::LifetimeParamKind::Elided)
+                        } else {
+                            (
+                                hir::ParamName::Plain(lifetime.ident),
+                                hir::LifetimeParamKind::Explicit,
+                            )
+                        };
 
-                    let (name, kind) = if lifetime.ident.name == kw::UnderscoreLifetime {
-                        (hir::ParamName::Fresh, hir::LifetimeParamKind::Elided)
-                    } else {
-                        (hir::ParamName::Plain(lifetime.ident), hir::LifetimeParamKind::Explicit)
-                    };
+                        hir::GenericParam {
+                            hir_id,
+                            name,
+                            span: lifetime.ident.span,
+                            pure_wrt_drop: false,
+                            kind: hir::GenericParamKind::Lifetime { kind },
+                            colon_span: None,
+                        }
+                    }));
 
-                    hir::GenericParam {
-                        hir_id,
-                        name,
-                        span: lifetime.ident.span,
-                        pure_wrt_drop: false,
-                        kind: hir::GenericParamKind::Lifetime { kind },
-                        colon_span: None,
-                    }
-                }));
-
-            debug!("lower_opaque_impl_trait: lifetime_defs={:#?}", lifetime_defs);
-
-            let opaque_ty_item = hir::OpaqueTy {
-                generics: self.arena.alloc(hir::Generics {
-                    params: lifetime_defs,
-                    predicates: &[],
-                    has_where_clause_predicates: false,
-                    where_clause_span: lctx.lower_span(span),
-                    span: lctx.lower_span(span),
-                }),
-                bounds: hir_bounds,
-                origin,
-            };
+                debug!("lower_opaque_impl_trait: lifetime_defs={:#?}", lifetime_defs);
+
+                let opaque_ty_item = hir::OpaqueTy {
+                    generics: self.arena.alloc(hir::Generics {
+                        params: lifetime_defs,
+                        predicates: &[],
+                        has_where_clause_predicates: false,
+                        where_clause_span: lctx.lower_span(span),
+                        span: lctx.lower_span(span),
+                    }),
+                    bounds: hir_bounds,
+                    origin,
+                };
 
-            trace!("lower_opaque_impl_trait: {:#?}", opaque_ty_def_id);
-            lctx.generate_opaque_type(opaque_ty_def_id, opaque_ty_item, span, opaque_ty_span)
+                trace!("lower_opaque_impl_trait: {:#?}", opaque_ty_def_id);
+                lctx.generate_opaque_type(opaque_ty_def_id, opaque_ty_item, span, opaque_ty_span)
+            })
         });
 
         let lifetimes =
@@ -1746,58 +1783,62 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
                 &mut new_remapping,
             );
 
-            // We have to be careful to get elision right here. The
-            // idea is that we create a lifetime parameter for each
-            // lifetime in the return type.  So, given a return type
-            // like `async fn foo(..) -> &[&u32]`, we lower to `impl
-            // Future<Output = &'1 [ &'2 u32 ]>`.
-            //
-            // Then, we will create `fn foo(..) -> Foo<'_, '_>`, and
-            // hence the elision takes place at the fn site.
-            let ret = this.lower_async_fn_output_type_to_future_bound(output, fn_def_id, span);
-
             let ctxt = std::mem::replace(&mut this.captured_lifetimes, lifetime_stash).unwrap();
 
             captures = ctxt.captures;
 
-            let future_bound = ret;
-
-            let generic_params =
-                this.arena.alloc_from_iter(captures.iter().map(|&(lifetime, _)| {
-                    let hir_id = this.lower_node_id(lifetime.id);
-                    debug_assert_ne!(this.opt_local_def_id(lifetime.id), None);
-
-                    let (name, kind) = if lifetime.ident.name == kw::UnderscoreLifetime {
-                        (hir::ParamName::Fresh, hir::LifetimeParamKind::Elided)
-                    } else {
-                        (hir::ParamName::Plain(lifetime.ident), hir::LifetimeParamKind::Explicit)
-                    };
+            this.with_remapping(new_remapping, |this| {
+                // We have to be careful to get elision right here. The
+                // idea is that we create a lifetime parameter for each
+                // lifetime in the return type.  So, given a return type
+                // like `async fn foo(..) -> &[&u32]`, we lower to `impl
+                // Future<Output = &'1 [ &'2 u32 ]>`.
+                //
+                // Then, we will create `fn foo(..) -> Foo<'_, '_>`, and
+                // hence the elision takes place at the fn site.
+                let future_bound =
+                    this.lower_async_fn_output_type_to_future_bound(output, fn_def_id, span);
+
+                let generic_params =
+                    this.arena.alloc_from_iter(captures.iter().map(|&(lifetime, _)| {
+                        let hir_id = this.lower_node_id(lifetime.id);
+                        debug_assert_ne!(this.opt_local_def_id(lifetime.id), None);
+
+                        let (name, kind) = if lifetime.ident.name == kw::UnderscoreLifetime {
+                            (hir::ParamName::Fresh, hir::LifetimeParamKind::Elided)
+                        } else {
+                            (
+                                hir::ParamName::Plain(lifetime.ident),
+                                hir::LifetimeParamKind::Explicit,
+                            )
+                        };
 
-                    hir::GenericParam {
-                        hir_id,
-                        name,
-                        span: lifetime.ident.span,
-                        pure_wrt_drop: false,
-                        kind: hir::GenericParamKind::Lifetime { kind },
-                        colon_span: None,
-                    }
-                }));
-            debug!("lower_async_fn_ret_ty: generic_params={:#?}", generic_params);
-
-            let opaque_ty_item = hir::OpaqueTy {
-                generics: this.arena.alloc(hir::Generics {
-                    params: generic_params,
-                    predicates: &[],
-                    has_where_clause_predicates: false,
-                    where_clause_span: this.lower_span(span),
-                    span: this.lower_span(span),
-                }),
-                bounds: arena_vec![this; future_bound],
-                origin: hir::OpaqueTyOrigin::AsyncFn(fn_def_id),
-            };
+                        hir::GenericParam {
+                            hir_id,
+                            name,
+                            span: lifetime.ident.span,
+                            pure_wrt_drop: false,
+                            kind: hir::GenericParamKind::Lifetime { kind },
+                            colon_span: None,
+                        }
+                    }));
+                debug!("lower_async_fn_ret_ty: generic_params={:#?}", generic_params);
+
+                let opaque_ty_item = hir::OpaqueTy {
+                    generics: this.arena.alloc(hir::Generics {
+                        params: generic_params,
+                        predicates: &[],
+                        has_where_clause_predicates: false,
+                        where_clause_span: this.lower_span(span),
+                        span: this.lower_span(span),
+                    }),
+                    bounds: arena_vec![this; future_bound],
+                    origin: hir::OpaqueTyOrigin::AsyncFn(fn_def_id),
+                };
 
-            trace!("exist ty from async fn def id: {:#?}", opaque_ty_def_id);
-            this.generate_opaque_type(opaque_ty_def_id, opaque_ty_item, span, opaque_ty_span)
+                trace!("exist ty from async fn def id: {:#?}", opaque_ty_def_id);
+                this.generate_opaque_type(opaque_ty_def_id, opaque_ty_item, span, opaque_ty_span)
+            })
         });
 
         // As documented above, we need to create the lifetime
@@ -1910,40 +1951,25 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
         ident: Ident,
         res: LifetimeRes,
     ) -> hir::Lifetime {
-        debug!(?self.captured_lifetimes);
-
         let name = match res {
-            LifetimeRes::Param { mut param, .. } => {
+            LifetimeRes::Param { param, .. } => {
                 let p_name = ParamName::Plain(ident);
-                if let Some(mut captured_lifetimes) = self.captured_lifetimes.take() {
-                    if let Entry::Occupied(o) = captured_lifetimes.captures.entry(param) {
-                        param = self.local_def_id(o.get().0.id);
-                    }
-
-                    self.captured_lifetimes = Some(captured_lifetimes);
-                }
+                let param = self.resolver.get_remapped_def_id(param);
 
                 hir::LifetimeName::Param(param, p_name)
             }
             LifetimeRes::Fresh { param, .. } => {
                 debug_assert_eq!(ident.name, kw::UnderscoreLifetime);
+                let param = self.local_def_id(param);
 
-                let mut param = self.local_def_id(param);
-                if let Some(mut captured_lifetimes) = self.captured_lifetimes.take() {
-                    if let Entry::Occupied(o) = captured_lifetimes.captures.entry(param) {
-                        param = self.local_def_id(o.get().0.id);
-                    }
-
-                    self.captured_lifetimes = Some(captured_lifetimes);
-                }
                 hir::LifetimeName::Param(param, ParamName::Fresh)
             }
             LifetimeRes::Infer => hir::LifetimeName::Infer,
             LifetimeRes::Static => hir::LifetimeName::Static,
             LifetimeRes::Error => hir::LifetimeName::Error,
             res => panic!("Unexpected lifetime resolution {:?} for {:?} at {:?}", res, ident, span),
         };
-        debug!(?self.captured_lifetimes);
+
         debug!(?name);
         hir::Lifetime { hir_id: self.lower_node_id(id), span: self.lower_span(span), name }
     }
diff --git a/compiler/rustc_middle/src/ty/mod.rs b/compiler/rustc_middle/src/ty/mod.rs
@@ -177,6 +177,8 @@ pub struct ResolverAstLowering {
     pub label_res_map: NodeMap<ast::NodeId>,
     /// Resolutions for lifetimes.
     pub lifetimes_res_map: NodeMap<LifetimeRes>,
+    /// Mapping from generics def-id to RPIT copied generic def-id
+    pub generics_def_id_map: Vec<FxHashMap<LocalDefId, LocalDefId>>,
     /// Lifetime parameters that lowering will have to introduce.
     pub extra_lifetime_params_map: NodeMap<Vec<(Ident, ast::NodeId, LifetimeRes)>>,
 
diff --git a/compiler/rustc_resolve/src/lib.rs b/compiler/rustc_resolve/src/lib.rs
@@ -913,6 +913,8 @@ pub struct Resolver<'a> {
     label_res_map: NodeMap<NodeId>,
     /// Resolutions for lifetimes.
     lifetimes_res_map: NodeMap<LifetimeRes>,
+    /// Mapping from generics def-id to RPIT copied generic def-id
+    generics_def_id_map: Vec<FxHashMap<LocalDefId, LocalDefId>>,
     /// Lifetime parameters that lowering will have to introduce.
     extra_lifetime_params_map: NodeMap<Vec<(Ident, NodeId, LifetimeRes)>>,
 
@@ -1277,6 +1279,7 @@ impl<'a> Resolver<'a> {
             import_res_map: Default::default(),
             label_res_map: Default::default(),
             lifetimes_res_map: Default::default(),
+            generics_def_id_map: Vec::new(),
             extra_lifetime_params_map: Default::default(),
             extern_crate_map: Default::default(),
             reexport_map: FxHashMap::default(),
@@ -1444,6 +1447,7 @@ impl<'a> Resolver<'a> {
             import_res_map: self.import_res_map,
             label_res_map: self.label_res_map,
             lifetimes_res_map: self.lifetimes_res_map,
+            generics_def_id_map: self.generics_def_id_map,
             extra_lifetime_params_map: self.extra_lifetime_params_map,
             next_node_id: self.next_node_id,
             node_id_to_def_id: self.node_id_to_def_id,
@@ -1488,6 +1492,7 @@ impl<'a> Resolver<'a> {
             import_res_map: self.import_res_map.clone(),
             label_res_map: self.label_res_map.clone(),
             lifetimes_res_map: self.lifetimes_res_map.clone(),
+            generics_def_id_map: self.generics_def_id_map.clone(),
             extra_lifetime_params_map: self.extra_lifetime_params_map.clone(),
             next_node_id: self.next_node_id.clone(),
             node_id_to_def_id: self.node_id_to_def_id.clone(),
