Make node_id_to_hir_id owner-local.

Author: cjgillot

compiler/rustc_ast_lowering/src/item.rs

@@ -473,10 +473,6 @@ impl<'hir> LoweringContext<'_, 'hir> {
                     let new_id = self.resolver.local_def_id(new_node_id);
                     let Some(res) = resolutions.next() else {
                         // Associate an HirId to both ids even if there is no resolution.
-                        let _old = self
-                            .node_id_to_hir_id
-                            .insert(new_node_id, hir::HirId::make_owner(new_id));
-                        debug_assert!(_old.is_none());
                         self.owners.ensure_contains_elem(new_id, || hir::MaybeOwner::Phantom);
                         let _old = std::mem::replace(
                             &mut self.owners[new_id],

compiler/rustc_ast_lowering/src/lib.rs

@@ -44,7 +44,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::FxHashSet;
+use rustc_data_structures::fx::{FxHashMap, FxHashSet};
 use rustc_data_structures::sorted_map::SortedMap;
 use rustc_data_structures::stable_hasher::{HashStable, StableHasher};
 use rustc_data_structures::sync::Lrc;
@@ -67,6 +67,7 @@ use rustc_span::symbol::{kw, sym, Ident, Symbol};
 use rustc_span::{Span, DUMMY_SP};
 
 use smallvec::SmallVec;
+use std::collections::hash_map::Entry;
 use tracing::{debug, trace};
 
 macro_rules! arena_vec {
@@ -154,10 +155,9 @@ struct LoweringContext<'a, 'hir: 'a> {
 
     current_hir_id_owner: LocalDefId,
     item_local_id_counter: hir::ItemLocalId,
-    node_id_to_hir_id: IndexVec<NodeId, Option<hir::HirId>>,
 
     /// NodeIds that are lowered inside the current HIR owner.
-    local_node_ids: Vec<NodeId>,
+    node_id_to_local_id: FxHashMap<NodeId, hir::ItemLocalId>,
 
     allow_try_trait: Option<Lrc<[Symbol]>>,
     allow_gen_future: Option<Lrc<[Symbol]>>,
@@ -311,8 +311,7 @@ pub fn lower_crate<'a, 'hir>(
         anonymous_lifetime_mode: AnonymousLifetimeMode::PassThrough,
         current_hir_id_owner: CRATE_DEF_ID,
         item_local_id_counter: hir::ItemLocalId::new(0),
-        node_id_to_hir_id: IndexVec::new(),
-        local_node_ids: Vec::new(),
+        node_id_to_local_id: FxHashMap::default(),
         generator_kind: None,
         task_context: None,
         current_item: None,
@@ -439,23 +438,22 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
 
         let current_attrs = std::mem::take(&mut self.attrs);
         let current_bodies = std::mem::take(&mut self.bodies);
-        let current_node_ids = std::mem::take(&mut self.local_node_ids);
+        let current_node_ids = std::mem::take(&mut self.node_id_to_local_id);
         let current_owner = std::mem::replace(&mut self.current_hir_id_owner, def_id);
         let current_local_counter =
             std::mem::replace(&mut self.item_local_id_counter, hir::ItemLocalId::new(1));
 
         // Always allocate the first `HirId` for the owner itself.
-        let _old = self.node_id_to_hir_id.insert(owner, hir::HirId::make_owner(def_id));
+        let _old = self.node_id_to_local_id.insert(owner, hir::ItemLocalId::new(0));
         debug_assert_eq!(_old, None);
-        self.local_node_ids.push(owner);
 
         let item = f(self);
         debug_assert_eq!(def_id, item.def_id());
         let info = self.make_owner_info(item);
 
         self.attrs = current_attrs;
         self.bodies = current_bodies;
-        self.local_node_ids = current_node_ids;
+        self.node_id_to_local_id = current_node_ids;
         self.current_hir_id_owner = current_owner;
         self.item_local_id_counter = current_local_counter;
 
@@ -468,32 +466,32 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
     fn make_owner_info(&mut self, node: hir::OwnerNode<'hir>) -> hir::OwnerInfo<'hir> {
         let attrs = std::mem::take(&mut self.attrs);
         let mut bodies = std::mem::take(&mut self.bodies);
-        let local_node_ids = std::mem::take(&mut self.local_node_ids);
+        let node_id_to_local_id = std::mem::take(&mut self.node_id_to_local_id);
 
-        let local_id_to_def_id = local_node_ids
+        let local_id_to_def_id = node_id_to_local_id
             .iter()
-            .filter_map(|&node_id| {
-                let hir_id = self.node_id_to_hir_id[node_id]?;
-                if hir_id.local_id == hir::ItemLocalId::new(0) {
+            .filter_map(|(&node_id, &local_id)| {
+                if local_id == hir::ItemLocalId::new(0) {
                     None
                 } else {
                     let def_id = self.resolver.opt_local_def_id(node_id)?;
+
                     self.owners.ensure_contains_elem(def_id, || hir::MaybeOwner::Phantom);
                     if let o @ hir::MaybeOwner::Phantom = &mut self.owners[def_id] {
                         // Do not override a `MaybeOwner::Owner` that may already here.
+                        let hir_id = hir::HirId { owner: self.current_hir_id_owner, local_id };
                         *o = hir::MaybeOwner::NonOwner(hir_id);
                     }
-                    Some((hir_id.local_id, def_id))
+                    Some((local_id, def_id))
                 }
             })
             .collect();
 
-        let trait_map = local_node_ids
+        let trait_map = node_id_to_local_id
             .into_iter()
-            .filter_map(|node_id| {
-                let hir_id = self.node_id_to_hir_id[node_id]?;
+            .filter_map(|(node_id, local_id)| {
                 let traits = self.resolver.take_trait_map(node_id)?;
-                Some((hir_id.local_id, traits.into_boxed_slice()))
+                Some((local_id, traits.into_boxed_slice()))
             })
             .collect();
 
@@ -558,14 +556,18 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
     fn lower_node_id(&mut self, ast_node_id: NodeId) -> hir::HirId {
         assert_ne!(ast_node_id, DUMMY_NODE_ID);
 
-        *self.node_id_to_hir_id.get_or_insert_with(ast_node_id, || {
-            // Generate a new `HirId`.
-            let owner = self.current_hir_id_owner;
-            let local_id = self.item_local_id_counter;
-            self.item_local_id_counter.increment_by(1);
-            self.local_node_ids.push(ast_node_id);
-            hir::HirId { owner, local_id }
-        })
+        let owner = self.current_hir_id_owner;
+        let local_id = match self.node_id_to_local_id.entry(ast_node_id) {
+            Entry::Occupied(o) => *o.get(),
+            Entry::Vacant(v) => {
+                // Generate a new `HirId`.
+                let local_id = self.item_local_id_counter;
+                self.item_local_id_counter.increment_by(1);
+                v.insert(local_id);
+                local_id
+            }
+        };
+        hir::HirId { owner, local_id }
     }
 
     fn next_id(&mut self) -> hir::HirId {
@@ -574,11 +576,17 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
     }
 
     fn lower_res(&mut self, res: Res<NodeId>) -> Res {
-        res.map_id(|id| {
-            self.node_id_to_hir_id.get(id).copied().flatten().unwrap_or_else(|| {
-                panic!("expected `NodeId` to be lowered already for res {:#?}", res);
-            })
-        })
+        let res: Result<Res, ()> = res.apply_id(|id| {
+            let owner = self.current_hir_id_owner;
+            let local_id = self.node_id_to_local_id.get(&id).copied().ok_or(())?;
+            Ok(hir::HirId { owner, local_id })
+        });
+        // We may fail to find a HirId when the Res points to a Local from an enclosing HIR owner.
+        // This can happen when trying to lower the return type `x` in erroneous code like
+        //   async fn foo(x: u8) -> x {}
+        // In that case, `x` is lowered as a function parameter, and the return type is lowered as
+        // an opaque type as a synthetized HIR owner.
+        res.unwrap_or(Res::Err)
     }
 
     fn expect_full_res(&mut self, id: NodeId) -> Res<NodeId> {

compiler/rustc_hir/src/def.rs

@@ -611,6 +611,19 @@ impl<Id> Res<Id> {
         }
     }
 
+    pub fn apply_id<R, E>(self, mut map: impl FnMut(Id) -> Result<R, E>) -> Result<Res<R>, E> {
+        Ok(match self {
+            Res::Def(kind, id) => Res::Def(kind, id),
+            Res::SelfCtor(id) => Res::SelfCtor(id),
+            Res::PrimTy(id) => Res::PrimTy(id),
+            Res::Local(id) => Res::Local(map(id)?),
+            Res::SelfTy { trait_, alias_to } => Res::SelfTy { trait_, alias_to },
+            Res::ToolMod => Res::ToolMod,
+            Res::NonMacroAttr(attr_kind) => Res::NonMacroAttr(attr_kind),
+            Res::Err => Res::Err,
+        })
+    }
+
     #[track_caller]
     pub fn expect_non_local<OtherId>(self) -> Res<OtherId> {
         self.map_id(|_| panic!("unexpected `Res::Local`"))