From 2a0ed1fdbcc002027672761cffc1b682afe4cb8a Mon Sep 17 00:00:00 2001
From: Camille GILLOT <gillot.camille@gmail.com>
Date: Wed, 22 Sep 2021 19:28:20 +0200
Subject: [PATCH 1/3] Store def_id_to_hir_id as variant in hir_owner.

If hir_owner is Owner(_), the LocalDefId is pointing to an owner, so the ItemLocalId is 0.
If the HIR node does not exist, we store Phantom.
Otherwise, we store the HirId associated to the LocalDefId.
---
 compiler/rustc_ast_lowering/src/item.rs       | 17 +++--
 compiler/rustc_ast_lowering/src/lib.rs        | 33 ++++-----
 compiler/rustc_hir/src/definitions.rs         | 31 +--------
 compiler/rustc_hir/src/hir.rs                 | 48 +++++++++++--
 compiler/rustc_hir/src/hir_id.rs              |  5 ++
 compiler/rustc_middle/src/hir/map/mod.rs      | 69 ++++++++++---------
 compiler/rustc_middle/src/hir/mod.rs          | 18 ++---
 compiler/rustc_middle/src/query/mod.rs        |  4 +-
 compiler/rustc_middle/src/ty/context.rs       |  2 -
 .../rustc_mir_transform/src/coverage/mod.rs   |  2 +-
 compiler/rustc_resolve/src/late/lifetimes.rs  | 10 +--
 11 files changed, 126 insertions(+), 113 deletions(-)

diff --git a/compiler/rustc_ast_lowering/src/item.rs b/compiler/rustc_ast_lowering/src/item.rs
index ed3abbd5b4d3..8b4a0840df58 100644
--- a/compiler/rustc_ast_lowering/src/item.rs
+++ b/compiler/rustc_ast_lowering/src/item.rs
@@ -105,12 +105,11 @@ impl<'hir> LoweringContext<'_, 'hir> {
     ) -> T {
         let old_len = self.in_scope_lifetimes.len();
 
-        let parent_generics =
-            match self.owners[parent_hir_id].as_ref().unwrap().node().expect_item().kind {
-                hir::ItemKind::Impl(hir::Impl { ref generics, .. })
-                | hir::ItemKind::Trait(_, _, ref generics, ..) => generics.params,
-                _ => &[],
-            };
+        let parent_generics = match self.owners[parent_hir_id].unwrap().node().expect_item().kind {
+            hir::ItemKind::Impl(hir::Impl { ref generics, .. })
+            | hir::ItemKind::Trait(_, _, ref generics, ..) => generics.params,
+            _ => &[],
+        };
         let lt_def_names = parent_generics.iter().filter_map(|param| match param.kind {
             hir::GenericParamKind::Lifetime { .. } => Some(param.name.normalize_to_macros_2_0()),
             _ => None,
@@ -480,6 +479,12 @@ impl<'hir> LoweringContext<'_, 'hir> {
                             .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],
+                            hir::MaybeOwner::NonOwner(hir::HirId::make_owner(new_id)),
+                        );
+                        debug_assert!(matches!(_old, hir::MaybeOwner::Phantom));
                         continue;
                     };
                     let ident = *ident;
diff --git a/compiler/rustc_ast_lowering/src/lib.rs b/compiler/rustc_ast_lowering/src/lib.rs
index 32cec3a295a4..a00471038947 100644
--- a/compiler/rustc_ast_lowering/src/lib.rs
+++ b/compiler/rustc_ast_lowering/src/lib.rs
@@ -98,7 +98,7 @@ struct LoweringContext<'a, 'hir: 'a> {
     arena: &'hir Arena<'hir>,
 
     /// The items being lowered are collected here.
-    owners: IndexVec<LocalDefId, Option<hir::OwnerInfo<'hir>>>,
+    owners: IndexVec<LocalDefId, hir::MaybeOwner<&'hir hir::OwnerInfo<'hir>>>,
     /// Bodies inside the owner being lowered.
     bodies: Vec<(hir::ItemLocalId, &'hir hir::Body<'hir>)>,
     /// Attributes inside the owner being lowered.
@@ -291,7 +291,8 @@ pub fn lower_crate<'a, 'hir>(
 ) -> &'hir hir::Crate<'hir> {
     let _prof_timer = sess.prof.verbose_generic_activity("hir_lowering");
 
-    let owners = IndexVec::from_fn_n(|_| None, resolver.definitions().def_index_count());
+    let owners =
+        IndexVec::from_fn_n(|_| hir::MaybeOwner::Phantom, resolver.definitions().def_index_count());
     LoweringContext {
         sess,
         resolver,
@@ -402,19 +403,6 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
 
         let hir_hash = self.compute_hir_hash();
 
-        let mut def_id_to_hir_id = IndexVec::default();
-
-        for (node_id, hir_id) in self.node_id_to_hir_id.into_iter_enumerated() {
-            if let Some(def_id) = self.resolver.opt_local_def_id(node_id) {
-                if def_id_to_hir_id.len() <= def_id.index() {
-                    def_id_to_hir_id.resize(def_id.index() + 1, None);
-                }
-                def_id_to_hir_id[def_id] = hir_id;
-            }
-        }
-
-        self.resolver.definitions().init_def_id_to_hir_id_mapping(def_id_to_hir_id);
-
         let krate = hir::Crate { owners: self.owners, hir_hash };
         self.arena.alloc(krate)
     }
@@ -427,7 +415,7 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
             .owners
             .iter_enumerated()
             .filter_map(|(def_id, info)| {
-                let info = info.as_ref()?;
+                let info = info.as_owner()?;
                 let def_path_hash = definitions.def_path_hash(def_id);
                 Some((def_path_hash, info))
             })
@@ -469,8 +457,8 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
         self.current_hir_id_owner = current_owner;
         self.item_local_id_counter = current_local_counter;
 
-        let _old = self.owners.insert(def_id, info);
-        debug_assert!(_old.is_none());
+        self.owners.ensure_contains_elem(def_id, || hir::MaybeOwner::Phantom);
+        self.owners[def_id] = hir::MaybeOwner::Owner(self.arena.alloc(info));
 
         def_id
     }
@@ -483,11 +471,18 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
         let local_id_to_def_id = local_node_ids
             .iter()
             .filter_map(|&node_id| {
+                let def_id = self.resolver.opt_local_def_id(node_id)?;
                 let hir_id = self.node_id_to_hir_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.
+                    *o = hir::MaybeOwner::NonOwner(hir_id);
+                }
+
                 if hir_id.local_id == hir::ItemLocalId::new(0) {
                     None
                 } else {
-                    let def_id = self.resolver.opt_local_def_id(node_id)?;
                     Some((hir_id.local_id, def_id))
                 }
             })
diff --git a/compiler/rustc_hir/src/definitions.rs b/compiler/rustc_hir/src/definitions.rs
index 74d6b05ca5fc..d655f12f5e1d 100644
--- a/compiler/rustc_hir/src/definitions.rs
+++ b/compiler/rustc_hir/src/definitions.rs
@@ -7,7 +7,6 @@
 pub use crate::def_id::DefPathHash;
 use crate::def_id::{CrateNum, DefIndex, LocalDefId, StableCrateId, CRATE_DEF_INDEX, LOCAL_CRATE};
 use crate::def_path_hash_map::DefPathHashMap;
-use crate::hir;
 
 use rustc_data_structures::fx::FxHashMap;
 use rustc_data_structures::stable_hasher::StableHasher;
@@ -101,13 +100,6 @@ impl DefPathTable {
 pub struct Definitions {
     table: DefPathTable,
 
-    /// Only [`LocalDefId`]s for items and item-like are HIR owners.
-    /// The associated `HirId` has a `local_id` of `0`.
-    /// Generic parameters and closures are also assigned a `LocalDefId` but are not HIR owners.
-    /// Their `HirId`s are defined by their position while lowering the enclosing owner.
-    // FIXME(cjgillot) Some `LocalDefId`s from `use` items are dropped during lowering and lack a `HirId`.
-    pub(super) def_id_to_hir_id: IndexVec<LocalDefId, Option<hir::HirId>>,
-
     /// Item with a given `LocalDefId` was defined during macro expansion with ID `ExpnId`.
     expansions_that_defined: FxHashMap<LocalDefId, ExpnId>,
 
@@ -322,12 +314,6 @@ impl Definitions {
         })
     }
 
-    #[inline]
-    #[track_caller]
-    pub fn local_def_id_to_hir_id(&self, id: LocalDefId) -> hir::HirId {
-        self.def_id_to_hir_id[id].unwrap()
-    }
-
     /// Adds a root definition (no parent) and a few other reserved definitions.
     pub fn new(stable_crate_id: StableCrateId, crate_span: Span) -> Definitions {
         let key = DefKey {
@@ -354,7 +340,6 @@ impl Definitions {
 
         Definitions {
             table,
-            def_id_to_hir_id: Default::default(),
             expansions_that_defined: Default::default(),
             def_id_to_span,
             stable_crate_id,
@@ -406,20 +391,6 @@ impl Definitions {
         def_id
     }
 
-    /// Initializes the `LocalDefId` to `HirId` mapping once it has been generated during
-    /// AST to HIR lowering.
-    pub fn init_def_id_to_hir_id_mapping(
-        &mut self,
-        mapping: IndexVec<LocalDefId, Option<hir::HirId>>,
-    ) {
-        assert!(
-            self.def_id_to_hir_id.is_empty(),
-            "trying to initialize `LocalDefId` <-> `HirId` mappings twice"
-        );
-
-        self.def_id_to_hir_id = mapping;
-    }
-
     pub fn expansion_that_defined(&self, id: LocalDefId) -> ExpnId {
         self.expansions_that_defined.get(&id).copied().unwrap_or_else(ExpnId::root)
     }
@@ -431,7 +402,7 @@ impl Definitions {
     }
 
     pub fn iter_local_def_id(&self) -> impl Iterator<Item = LocalDefId> + '_ {
-        self.def_id_to_hir_id.iter_enumerated().map(|(k, _)| k)
+        self.table.def_path_hashes.indices().map(|local_def_index| LocalDefId { local_def_index })
     }
 
     #[inline(always)]
diff --git a/compiler/rustc_hir/src/hir.rs b/compiler/rustc_hir/src/hir.rs
index a0ed72c9e9e5..67398c80f360 100644
--- a/compiler/rustc_hir/src/hir.rs
+++ b/compiler/rustc_hir/src/hir.rs
@@ -711,6 +711,15 @@ pub struct OwnerNodes<'tcx> {
     pub local_id_to_def_id: SortedMap<ItemLocalId, LocalDefId>,
 }
 
+impl<'tcx> OwnerNodes<'tcx> {
+    pub fn node(&self) -> OwnerNode<'tcx> {
+        use rustc_index::vec::Idx;
+        let node = self.nodes[ItemLocalId::new(0)].as_ref().unwrap().node;
+        let node = node.as_owner().unwrap(); // Indexing must ensure it is an OwnerNode.
+        node
+    }
+}
+
 /// Full information resulting from lowering an AST node.
 #[derive(Debug, HashStable_Generic)]
 pub struct OwnerInfo<'hir> {
@@ -728,10 +737,39 @@ pub struct OwnerInfo<'hir> {
 impl<'tcx> OwnerInfo<'tcx> {
     #[inline]
     pub fn node(&self) -> OwnerNode<'tcx> {
-        use rustc_index::vec::Idx;
-        let node = self.nodes.nodes[ItemLocalId::new(0)].as_ref().unwrap().node;
-        let node = node.as_owner().unwrap(); // Indexing must ensure it is an OwnerNode.
-        node
+        self.nodes.node()
+    }
+}
+
+#[derive(Copy, Clone, Debug, HashStable_Generic)]
+pub enum MaybeOwner<T> {
+    Owner(T),
+    NonOwner(HirId),
+    /// Used as a placeholder for unused LocalDefId.
+    Phantom,
+}
+
+impl<T> MaybeOwner<T> {
+    pub fn as_owner(self) -> Option<T> {
+        match self {
+            MaybeOwner::Owner(i) => Some(i),
+            MaybeOwner::NonOwner(_) | MaybeOwner::Phantom => None,
+        }
+    }
+
+    pub fn map<U>(self, f: impl FnOnce(T) -> U) -> MaybeOwner<U> {
+        match self {
+            MaybeOwner::Owner(i) => MaybeOwner::Owner(f(i)),
+            MaybeOwner::NonOwner(hir_id) => MaybeOwner::NonOwner(hir_id),
+            MaybeOwner::Phantom => MaybeOwner::Phantom,
+        }
+    }
+
+    pub fn unwrap(self) -> T {
+        match self {
+            MaybeOwner::Owner(i) => i,
+            MaybeOwner::NonOwner(_) | MaybeOwner::Phantom => panic!("Not a HIR owner"),
+        }
     }
 }
 
@@ -743,7 +781,7 @@ impl<'tcx> OwnerInfo<'tcx> {
 /// [rustc dev guide]: https://rustc-dev-guide.rust-lang.org/hir.html
 #[derive(Debug)]
 pub struct Crate<'hir> {
-    pub owners: IndexVec<LocalDefId, Option<OwnerInfo<'hir>>>,
+    pub owners: IndexVec<LocalDefId, MaybeOwner<&'hir OwnerInfo<'hir>>>,
     pub hir_hash: Fingerprint,
 }
 
diff --git a/compiler/rustc_hir/src/hir_id.rs b/compiler/rustc_hir/src/hir_id.rs
index 5b486819c35e..dee391b9cce2 100644
--- a/compiler/rustc_hir/src/hir_id.rs
+++ b/compiler/rustc_hir/src/hir_id.rs
@@ -30,6 +30,11 @@ impl HirId {
         if self.local_id.index() == 0 { Some(self.owner) } else { None }
     }
 
+    #[inline]
+    pub fn is_owner(self) -> bool {
+        self.local_id.index() == 0
+    }
+
     #[inline]
     pub fn make_owner(owner: LocalDefId) -> Self {
         Self { owner, local_id: ItemLocalId::from_u32(0) }
diff --git a/compiler/rustc_middle/src/hir/map/mod.rs b/compiler/rustc_middle/src/hir/map/mod.rs
index 1885df6ac5d9..e75c16337ee4 100644
--- a/compiler/rustc_middle/src/hir/map/mod.rs
+++ b/compiler/rustc_middle/src/hir/map/mod.rs
@@ -123,7 +123,7 @@ impl<'hir> Iterator for ParentOwnerIterator<'hir> {
     fn next(&mut self) -> Option<Self::Item> {
         if self.current_id.local_id.index() != 0 {
             self.current_id.local_id = ItemLocalId::new(0);
-            if let Some(node) = self.map.tcx.hir_owner(self.current_id.owner) {
+            if let MaybeOwner::Owner(node) = self.map.tcx.hir_owner(self.current_id.owner) {
                 return Some((self.current_id.owner, node.node));
             }
         }
@@ -141,7 +141,7 @@ impl<'hir> Iterator for ParentOwnerIterator<'hir> {
             self.current_id = HirId::make_owner(parent_id);
 
             // If this `HirId` doesn't have an entry, skip it and look for its `parent_id`.
-            if let Some(node) = self.map.tcx.hir_owner(self.current_id.owner) {
+            if let MaybeOwner::Owner(node) = self.map.tcx.hir_owner(self.current_id.owner) {
                 return Some((self.current_id.owner, node.node));
             }
         }
@@ -155,14 +155,14 @@ impl<'hir> Map<'hir> {
 
     pub fn root_module(&self) -> &'hir Mod<'hir> {
         match self.tcx.hir_owner(CRATE_DEF_ID).map(|o| o.node) {
-            Some(OwnerNode::Crate(item)) => item,
+            MaybeOwner::Owner(OwnerNode::Crate(item)) => item,
             _ => bug!(),
         }
     }
 
     pub fn items(&self) -> impl Iterator<Item = &'hir Item<'hir>> + 'hir {
         let krate = self.krate();
-        krate.owners.iter().filter_map(|owner| match owner.as_ref()?.node() {
+        krate.owners.iter().filter_map(|owner| match owner.as_owner()?.node() {
             OwnerNode::Item(item) => Some(item),
             _ => None,
         })
@@ -205,7 +205,8 @@ impl<'hir> Map<'hir> {
             Some(hir_id.owner)
         } else {
             self.tcx
-                .hir_owner_nodes(hir_id.owner)?
+                .hir_owner_nodes(hir_id.owner)
+                .as_owner()?
                 .local_id_to_def_id
                 .get(&hir_id.local_id)
                 .copied()
@@ -214,8 +215,12 @@ impl<'hir> Map<'hir> {
 
     #[inline]
     pub fn local_def_id_to_hir_id(&self, def_id: LocalDefId) -> HirId {
-        // FIXME(#85914) is this access safe for incr. comp.?
-        self.tcx.untracked_resolutions.definitions.local_def_id_to_hir_id(def_id)
+        let owner = self.tcx.hir_owner(def_id);
+        match owner {
+            MaybeOwner::Owner(_) => HirId::make_owner(def_id),
+            MaybeOwner::Phantom => bug!("No HirId for {:?}", def_id),
+            MaybeOwner::NonOwner(hir_id) => hir_id,
+        }
     }
 
     pub fn iter_local_def_id(&self) -> impl Iterator<Item = LocalDefId> + '_ {
@@ -321,7 +326,7 @@ impl<'hir> Map<'hir> {
         if id.local_id == ItemLocalId::from_u32(0) {
             Some(self.tcx.hir_owner_parent(id.owner))
         } else {
-            let owner = self.tcx.hir_owner_nodes(id.owner)?;
+            let owner = self.tcx.hir_owner_nodes(id.owner).as_owner()?;
             let node = owner.nodes[id.local_id].as_ref()?;
             let hir_id = HirId { owner: id.owner, local_id: node.parent };
             Some(hir_id)
@@ -335,10 +340,10 @@ impl<'hir> Map<'hir> {
     /// Retrieves the `Node` corresponding to `id`, returning `None` if cannot be found.
     pub fn find(&self, id: HirId) -> Option<Node<'hir>> {
         if id.local_id == ItemLocalId::from_u32(0) {
-            let owner = self.tcx.hir_owner(id.owner)?;
+            let owner = self.tcx.hir_owner(id.owner).as_owner()?;
             Some(owner.node.into())
         } else {
-            let owner = self.tcx.hir_owner_nodes(id.owner)?;
+            let owner = self.tcx.hir_owner_nodes(id.owner).as_owner()?;
             let node = owner.nodes[id.local_id].as_ref()?;
             Some(node.node)
         }
@@ -366,7 +371,7 @@ impl<'hir> Map<'hir> {
     }
 
     pub fn get_generics(&self, id: LocalDefId) -> Option<&'hir Generics<'hir>> {
-        let node = self.tcx.hir_owner(id)?;
+        let node = self.tcx.hir_owner(id).as_owner()?;
         match node.node {
             OwnerNode::ImplItem(impl_item) => Some(&impl_item.generics),
             OwnerNode::TraitItem(trait_item) => Some(&trait_item.generics),
@@ -522,7 +527,7 @@ impl<'hir> Map<'hir> {
             .owners
             .iter_enumerated()
             .flat_map(move |(owner, owner_info)| {
-                let bodies = &owner_info.as_ref()?.nodes.bodies;
+                let bodies = &owner_info.as_owner()?.nodes.bodies;
                 Some(bodies.iter().map(move |&(local_id, _)| {
                     let hir_id = HirId { owner, local_id };
                     let body_id = BodyId { hir_id };
@@ -539,7 +544,7 @@ impl<'hir> Map<'hir> {
 
         par_iter(&self.krate().owners.raw).enumerate().for_each(|(owner, owner_info)| {
             let owner = LocalDefId::new(owner);
-            if let Some(owner_info) = owner_info {
+            if let MaybeOwner::Owner(owner_info) = owner_info {
                 par_iter(owner_info.nodes.bodies.range(..)).for_each(|(local_id, _)| {
                     let hir_id = HirId { owner, local_id: *local_id };
                     let body_id = BodyId { hir_id };
@@ -583,10 +588,10 @@ impl<'hir> Map<'hir> {
     pub fn get_module(&self, module: LocalDefId) -> (&'hir Mod<'hir>, Span, HirId) {
         let hir_id = HirId::make_owner(module);
         match self.tcx.hir_owner(module).map(|o| o.node) {
-            Some(OwnerNode::Item(&Item { span, kind: ItemKind::Mod(ref m), .. })) => {
-                (m, span, hir_id)
-            }
-            Some(OwnerNode::Crate(item)) => (item, item.inner, hir_id),
+            MaybeOwner::Owner(OwnerNode::Item(&Item {
+                span, kind: ItemKind::Mod(ref m), ..
+            })) => (m, span, hir_id),
+            MaybeOwner::Owner(OwnerNode::Crate(item)) => (item, item.inner, hir_id),
             node => panic!("not a module: {:?}", node),
         }
     }
@@ -601,7 +606,7 @@ impl<'hir> Map<'hir> {
     pub fn walk_attributes(self, visitor: &mut impl Visitor<'hir>) {
         let krate = self.krate();
         for (owner, info) in krate.owners.iter_enumerated() {
-            if let Some(info) = info {
+            if let MaybeOwner::Owner(info) = info {
                 for (local_id, attrs) in info.attrs.map.iter() {
                     let id = HirId { owner, local_id: *local_id };
                     for a in *attrs {
@@ -625,7 +630,7 @@ impl<'hir> Map<'hir> {
         V: itemlikevisit::ItemLikeVisitor<'hir>,
     {
         let krate = self.krate();
-        for owner in krate.owners.iter().filter_map(Option::as_ref) {
+        for owner in krate.owners.iter().filter_map(|i| i.as_owner()) {
             match owner.node() {
                 OwnerNode::Item(item) => visitor.visit_item(item),
                 OwnerNode::ForeignItem(item) => visitor.visit_foreign_item(item),
@@ -642,12 +647,14 @@ impl<'hir> Map<'hir> {
         V: itemlikevisit::ParItemLikeVisitor<'hir> + Sync + Send,
     {
         let krate = self.krate();
-        par_for_each_in(&krate.owners.raw, |owner| match owner.as_ref().map(OwnerInfo::node) {
-            Some(OwnerNode::Item(item)) => visitor.visit_item(item),
-            Some(OwnerNode::ForeignItem(item)) => visitor.visit_foreign_item(item),
-            Some(OwnerNode::ImplItem(item)) => visitor.visit_impl_item(item),
-            Some(OwnerNode::TraitItem(item)) => visitor.visit_trait_item(item),
-            Some(OwnerNode::Crate(_)) | None => {}
+        par_for_each_in(&krate.owners.raw, |owner| match owner.map(OwnerInfo::node) {
+            MaybeOwner::Owner(OwnerNode::Item(item)) => visitor.visit_item(item),
+            MaybeOwner::Owner(OwnerNode::ForeignItem(item)) => visitor.visit_foreign_item(item),
+            MaybeOwner::Owner(OwnerNode::ImplItem(item)) => visitor.visit_impl_item(item),
+            MaybeOwner::Owner(OwnerNode::TraitItem(item)) => visitor.visit_trait_item(item),
+            MaybeOwner::Owner(OwnerNode::Crate(_))
+            | MaybeOwner::NonOwner(_)
+            | MaybeOwner::Phantom => {}
         })
     }
 
@@ -878,7 +885,7 @@ impl<'hir> Map<'hir> {
 
     pub fn get_foreign_abi(&self, hir_id: HirId) -> Abi {
         let parent = self.get_parent_item(hir_id);
-        if let Some(node) = self.tcx.hir_owner(parent) {
+        if let MaybeOwner::Owner(node) = self.tcx.hir_owner(parent) {
             if let OwnerNode::Item(Item { kind: ItemKind::ForeignMod { abi, .. }, .. }) = node.node
             {
                 return *abi;
@@ -892,21 +899,21 @@ impl<'hir> Map<'hir> {
 
     pub fn expect_item(&self, id: LocalDefId) -> &'hir Item<'hir> {
         match self.tcx.hir_owner(id) {
-            Some(Owner { node: OwnerNode::Item(item), .. }) => item,
+            MaybeOwner::Owner(Owner { node: OwnerNode::Item(item), .. }) => item,
             _ => bug!("expected item, found {}", self.node_to_string(HirId::make_owner(id))),
         }
     }
 
     pub fn expect_impl_item(&self, id: LocalDefId) -> &'hir ImplItem<'hir> {
         match self.tcx.hir_owner(id) {
-            Some(Owner { node: OwnerNode::ImplItem(item), .. }) => item,
+            MaybeOwner::Owner(Owner { node: OwnerNode::ImplItem(item), .. }) => item,
             _ => bug!("expected impl item, found {}", self.node_to_string(HirId::make_owner(id))),
         }
     }
 
     pub fn expect_trait_item(&self, id: LocalDefId) -> &'hir TraitItem<'hir> {
         match self.tcx.hir_owner(id) {
-            Some(Owner { node: OwnerNode::TraitItem(item), .. }) => item,
+            MaybeOwner::Owner(Owner { node: OwnerNode::TraitItem(item), .. }) => item,
             _ => bug!("expected trait item, found {}", self.node_to_string(HirId::make_owner(id))),
         }
     }
@@ -920,7 +927,7 @@ impl<'hir> Map<'hir> {
 
     pub fn expect_foreign_item(&self, id: LocalDefId) -> &'hir ForeignItem<'hir> {
         match self.tcx.hir_owner(id) {
-            Some(Owner { node: OwnerNode::ForeignItem(item), .. }) => item,
+            MaybeOwner::Owner(Owner { node: OwnerNode::ForeignItem(item), .. }) => item,
             _ => {
                 bug!("expected foreign item, found {}", self.node_to_string(HirId::make_owner(id)))
             }
@@ -1121,7 +1128,7 @@ pub(super) fn crate_hash(tcx: TyCtxt<'_>, crate_num: CrateNum) -> Svh {
             .owners
             .iter_enumerated()
             .filter_map(|(def_id, info)| {
-                let _ = info.as_ref()?;
+                let _ = info.as_owner()?;
                 let def_path_hash = definitions.def_path_hash(def_id);
                 let span = definitions.def_span(def_id);
                 debug_assert_eq!(span.parent(), None);
diff --git a/compiler/rustc_middle/src/hir/mod.rs b/compiler/rustc_middle/src/hir/mod.rs
index b4c7bb7eba79..5543e96036b5 100644
--- a/compiler/rustc_middle/src/hir/mod.rs
+++ b/compiler/rustc_middle/src/hir/mod.rs
@@ -65,20 +65,20 @@ pub fn provide(providers: &mut Providers) {
     providers.crate_hash = map::crate_hash;
     providers.hir_module_items = map::hir_module_items;
     providers.hir_owner = |tcx, id| {
-        let owner = tcx.hir_crate(()).owners.get(id)?.as_ref()?;
-        let node = owner.node();
-        Some(Owner { node, hash_without_bodies: owner.nodes.hash_without_bodies })
+        tcx.hir_crate(()).owners[id].map(|owner| {
+            let node = owner.nodes.node();
+            Owner { node, hash_without_bodies: owner.nodes.hash_without_bodies }
+        })
     };
-    providers.hir_owner_nodes = |tcx, id| tcx.hir_crate(()).owners[id].as_ref().map(|i| &i.nodes);
+    providers.hir_owner_nodes = |tcx, id| tcx.hir_crate(()).owners[id].map(|i| &i.nodes);
     providers.hir_owner_parent = |tcx, id| {
         // Accessing the def_key is ok since its value is hashed as part of `id`'s DefPathHash.
         let parent = tcx.untracked_resolutions.definitions.def_key(id).parent;
         let parent = parent.map_or(CRATE_HIR_ID, |local_def_index| {
             let def_id = LocalDefId { local_def_index };
-            let mut parent_hir_id =
-                tcx.untracked_resolutions.definitions.local_def_id_to_hir_id(def_id);
+            let mut parent_hir_id = tcx.hir().local_def_id_to_hir_id(def_id);
             if let Some(local_id) =
-                tcx.hir_crate(()).owners[parent_hir_id.owner].as_ref().unwrap().parenting.get(&id)
+                tcx.hir_crate(()).owners[parent_hir_id.owner].unwrap().parenting.get(&id)
             {
                 parent_hir_id.local_id = *local_id;
             }
@@ -87,7 +87,7 @@ pub fn provide(providers: &mut Providers) {
         parent
     };
     providers.hir_attrs =
-        |tcx, id| tcx.hir_crate(()).owners[id].as_ref().map_or(AttributeMap::EMPTY, |o| &o.attrs);
+        |tcx, id| tcx.hir_crate(()).owners[id].as_owner().map_or(AttributeMap::EMPTY, |o| &o.attrs);
     providers.source_span = |tcx, def_id| tcx.resolutions(()).definitions.def_span(def_id);
     providers.def_span = |tcx, def_id| tcx.hir().span_if_local(def_id).unwrap_or(DUMMY_SP);
     providers.fn_arg_names = |tcx, id| {
@@ -111,4 +111,6 @@ pub fn provide(providers: &mut Providers) {
         let id = id.expect_local();
         tcx.resolutions(()).definitions.expansion_that_defined(id)
     };
+    providers.in_scope_traits_map =
+        |tcx, id| tcx.hir_crate(()).owners[id].as_owner().map(|owner_info| &owner_info.trait_map);
 }
diff --git a/compiler/rustc_middle/src/query/mod.rs b/compiler/rustc_middle/src/query/mod.rs
index 715a1fa25a1c..f5b4925fbb3e 100644
--- a/compiler/rustc_middle/src/query/mod.rs
+++ b/compiler/rustc_middle/src/query/mod.rs
@@ -52,7 +52,7 @@ rustc_queries! {
     ///
     /// This can be conveniently accessed by methods on `tcx.hir()`.
     /// Avoid calling this query directly.
-    query hir_owner(key: LocalDefId) -> Option<crate::hir::Owner<'tcx>> {
+    query hir_owner(key: LocalDefId) -> hir::MaybeOwner<crate::hir::Owner<'tcx>> {
         desc { |tcx| "HIR owner of `{}`", tcx.def_path_str(key.to_def_id()) }
     }
 
@@ -68,7 +68,7 @@ rustc_queries! {
     ///
     /// This can be conveniently accessed by methods on `tcx.hir()`.
     /// Avoid calling this query directly.
-    query hir_owner_nodes(key: LocalDefId) -> Option<&'tcx hir::OwnerNodes<'tcx>> {
+    query hir_owner_nodes(key: LocalDefId) -> hir::MaybeOwner<&'tcx hir::OwnerNodes<'tcx>> {
         desc { |tcx| "HIR owner items in `{}`", tcx.def_path_str(key.to_def_id()) }
     }
 
diff --git a/compiler/rustc_middle/src/ty/context.rs b/compiler/rustc_middle/src/ty/context.rs
index 6c8573805cb9..75f1318e3631 100644
--- a/compiler/rustc_middle/src/ty/context.rs
+++ b/compiler/rustc_middle/src/ty/context.rs
@@ -2875,8 +2875,6 @@ fn ptr_eq<T, U>(t: *const T, u: *const U) -> bool {
 }
 
 pub fn provide(providers: &mut ty::query::Providers) {
-    providers.in_scope_traits_map =
-        |tcx, id| tcx.hir_crate(()).owners[id].as_ref().map(|owner_info| &owner_info.trait_map);
     providers.resolutions = |tcx, ()| &tcx.untracked_resolutions;
     providers.module_reexports =
         |tcx, id| tcx.resolutions(()).reexport_map.get(&id).map(|v| &v[..]);
diff --git a/compiler/rustc_mir_transform/src/coverage/mod.rs b/compiler/rustc_mir_transform/src/coverage/mod.rs
index 82455654a886..2bb9f48f9b7c 100644
--- a/compiler/rustc_mir_transform/src/coverage/mod.rs
+++ b/compiler/rustc_mir_transform/src/coverage/mod.rs
@@ -579,7 +579,7 @@ fn hash_mir_source<'tcx>(tcx: TyCtxt<'tcx>, hir_body: &'tcx rustc_hir::Body<'tcx
     let mut hcx = tcx.create_no_span_stable_hashing_context();
     let mut stable_hasher = StableHasher::new();
     let owner = hir_body.id().hir_id.owner;
-    let bodies = &tcx.hir_owner_nodes(owner).as_ref().unwrap().bodies;
+    let bodies = &tcx.hir_owner_nodes(owner).unwrap().bodies;
     hcx.with_hir_bodies(false, owner, bodies, |hcx| {
         hir_body.value.hash_stable(hcx, &mut stable_hasher)
     });
diff --git a/compiler/rustc_resolve/src/late/lifetimes.rs b/compiler/rustc_resolve/src/late/lifetimes.rs
index 4c7bdb33fb87..6ef85c426be7 100644
--- a/compiler/rustc_resolve/src/late/lifetimes.rs
+++ b/compiler/rustc_resolve/src/late/lifetimes.rs
@@ -1019,15 +1019,7 @@ impl<'a, 'tcx> Visitor<'tcx> for LifetimeContext<'a, 'tcx> {
                     let hir_id = self.tcx.hir().local_def_id_to_hir_id(def_id);
                     // Ensure that the parent of the def is an item, not HRTB
                     let parent_id = self.tcx.hir().get_parent_node(hir_id);
-                    // FIXME(cjgillot) Can this check be replaced by
-                    // `let parent_is_item = parent_id.is_owner();`?
-                    let parent_is_item = if let Some(parent_def_id) = parent_id.as_owner() {
-                        matches!(self.tcx.hir().krate().owners.get(parent_def_id), Some(Some(_)),)
-                    } else {
-                        false
-                    };
-
-                    if !parent_is_item {
+                    if !parent_id.is_owner() {
                         if !self.trait_definition_only {
                             struct_span_err!(
                                 self.tcx.sess,

From 70ad73b21bef9508909a75e6fd3316f854cd10f2 Mon Sep 17 00:00:00 2001
From: Camille GILLOT <gillot.camille@gmail.com>
Date: Sun, 18 Jul 2021 13:03:53 +0200
Subject: [PATCH 2/3] Make node_id_to_hir_id owner-local.

---
 compiler/rustc_ast_lowering/src/item.rs |  4 --
 compiler/rustc_ast_lowering/src/lib.rs  | 73 +++++++++++++------------
 compiler/rustc_hir/src/def.rs           | 13 +++++
 3 files changed, 51 insertions(+), 39 deletions(-)

diff --git a/compiler/rustc_ast_lowering/src/item.rs b/compiler/rustc_ast_lowering/src/item.rs
index 8b4a0840df58..399a9e63987a 100644
--- a/compiler/rustc_ast_lowering/src/item.rs
+++ b/compiler/rustc_ast_lowering/src/item.rs
@@ -475,10 +475,6 @@ impl<'hir> LoweringContext<'_, 'hir> {
                         res
                     } 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],
diff --git a/compiler/rustc_ast_lowering/src/lib.rs b/compiler/rustc_ast_lowering/src/lib.rs
index a00471038947..50af46e75831 100644
--- a/compiler/rustc_ast_lowering/src/lib.rs
+++ b/compiler/rustc_ast_lowering/src/lib.rs
@@ -42,7 +42,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;
@@ -65,6 +65,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 {
@@ -152,10 +153,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]>>,
@@ -309,8 +309,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,
@@ -437,15 +436,14 @@ 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());
@@ -453,7 +451,7 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
 
         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;
 
@@ -466,34 +464,29 @@ 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| {
+            .filter_map(|(&node_id, &local_id)| {
                 let def_id = self.resolver.opt_local_def_id(node_id)?;
-                let hir_id = self.node_id_to_hir_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);
                 }
 
-                if hir_id.local_id == hir::ItemLocalId::new(0) {
-                    None
-                } else {
-                    Some((hir_id.local_id, def_id))
-                }
+                if local_id == hir::ItemLocalId::new(0) { None } else { 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 +551,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 +571,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> {
diff --git a/compiler/rustc_hir/src/def.rs b/compiler/rustc_hir/src/def.rs
index a43cb0203dd2..b86cfc331066 100644
--- a/compiler/rustc_hir/src/def.rs
+++ b/compiler/rustc_hir/src/def.rs
@@ -603,6 +603,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(a, b) => Res::SelfTy(a, b),
+            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`"))

From 680523560f9f2906a4fbab61c9254ee40d27e8e9 Mon Sep 17 00:00:00 2001
From: Camille GILLOT <gillot.camille@gmail.com>
Date: Sat, 29 Jan 2022 19:06:27 +0100
Subject: [PATCH 3/3] Only the query once for `find_by_def_id`.

---
 compiler/rustc_middle/src/hir/map/mod.rs | 19 ++++++++++++++-----
 1 file changed, 14 insertions(+), 5 deletions(-)

diff --git a/compiler/rustc_middle/src/hir/map/mod.rs b/compiler/rustc_middle/src/hir/map/mod.rs
index e75c16337ee4..87c45b9eca33 100644
--- a/compiler/rustc_middle/src/hir/map/mod.rs
+++ b/compiler/rustc_middle/src/hir/map/mod.rs
@@ -231,8 +231,7 @@ impl<'hir> Map<'hir> {
     }
 
     pub fn opt_def_kind(&self, local_def_id: LocalDefId) -> Option<DefKind> {
-        let hir_id = self.local_def_id_to_hir_id(local_def_id);
-        let def_kind = match self.find(hir_id)? {
+        let def_kind = match self.find_by_def_id(local_def_id)? {
             Node::Item(item) => match item.kind {
                 ItemKind::Static(..) => DefKind::Static,
                 ItemKind::Const(..) => DefKind::Const,
@@ -272,6 +271,7 @@ impl<'hir> Map<'hir> {
                 // FIXME(eddyb) is this even possible, if we have a `Node::Ctor`?
                 assert_ne!(variant_data.ctor_hir_id(), None);
 
+                let hir_id = self.local_def_id_to_hir_id(local_def_id);
                 let ctor_of = match self.find(self.get_parent_node(hir_id)) {
                     Some(Node::Item(..)) => def::CtorOf::Struct,
                     Some(Node::Variant(..)) => def::CtorOf::Variant,
@@ -280,6 +280,7 @@ impl<'hir> Map<'hir> {
                 DefKind::Ctor(ctor_of, def::CtorKind::from_hir(variant_data))
             }
             Node::AnonConst(_) => {
+                let hir_id = self.local_def_id_to_hir_id(local_def_id);
                 let inline = match self.find(self.get_parent_node(hir_id)) {
                     Some(Node::Expr(&Expr {
                         kind: ExprKind::ConstBlock(ref anon_const), ..
@@ -292,7 +293,10 @@ impl<'hir> Map<'hir> {
             Node::Expr(expr) => match expr.kind {
                 ExprKind::Closure(.., None) => DefKind::Closure,
                 ExprKind::Closure(.., Some(_)) => DefKind::Generator,
-                _ => bug!("def_kind: unsupported node: {}", self.node_to_string(hir_id)),
+                _ => {
+                    let hir_id = self.local_def_id_to_hir_id(local_def_id);
+                    bug!("def_kind: unsupported node: {}", self.node_to_string(hir_id))
+                }
             },
             Node::GenericParam(param) => match param.kind {
                 GenericParamKind::Lifetime { .. } => DefKind::LifetimeParam,
@@ -352,7 +356,12 @@ impl<'hir> Map<'hir> {
     /// Retrieves the `Node` corresponding to `id`, returning `None` if cannot be found.
     #[inline]
     pub fn find_by_def_id(&self, id: LocalDefId) -> Option<Node<'hir>> {
-        self.find(self.local_def_id_to_hir_id(id))
+        let owner = self.tcx.hir_owner(id);
+        match owner {
+            MaybeOwner::Owner(o) => Some(o.node.into()),
+            MaybeOwner::NonOwner(hir_id) => self.find(hir_id),
+            MaybeOwner::Phantom => bug!("No HirId for {:?}", id),
+        }
     }
 
     /// Retrieves the `Node` corresponding to `id`, panicking if it cannot be found.
@@ -367,7 +376,7 @@ impl<'hir> Map<'hir> {
     }
 
     pub fn get_if_local(&self, id: DefId) -> Option<Node<'hir>> {
-        id.as_local().and_then(|id| self.find(self.local_def_id_to_hir_id(id)))
+        id.as_local().and_then(|id| self.find_by_def_id(id))
     }
 
     pub fn get_generics(&self, id: LocalDefId) -> Option<&'hir Generics<'hir>> {