Make `IndexVec::ensure_contains_elem` return a reference to the element

6 files changed, 24 insertions, 26 deletions

diff --git a/compiler/rustc_ast_lowering/src/item.rs b/compiler/rustc_ast_lowering/src/item.rs
index f89e254a2f5..c061a244cf5 100644
--- a/compiler/rustc_ast_lowering/src/item.rs
+++ b/compiler/rustc_ast_lowering/src/item.rs
@@ -89,9 +89,9 @@ impl<'a, 'hir> ItemLowerer<'a, 'hir> {
         lctx.with_hir_id_owner(owner, |lctx| f(lctx));
 
         for (def_id, info) in lctx.children {
-            self.owners.ensure_contains_elem(def_id, || hir::MaybeOwner::Phantom);
-            debug_assert!(matches!(self.owners[def_id], hir::MaybeOwner::Phantom));
-            self.owners[def_id] = info;
+            let owner = self.owners.ensure_contains_elem(def_id, || hir::MaybeOwner::Phantom);
+            debug_assert!(matches!(owner, hir::MaybeOwner::Phantom));
+            *owner = info;
         }
     }
 
@@ -99,8 +99,8 @@ impl<'a, 'hir> ItemLowerer<'a, 'hir> {
         &mut self,
         def_id: LocalDefId,
     ) -> hir::MaybeOwner<&'hir hir::OwnerInfo<'hir>> {
-        self.owners.ensure_contains_elem(def_id, || hir::MaybeOwner::Phantom);
-        if let hir::MaybeOwner::Phantom = self.owners[def_id] {
+        let owner = self.owners.ensure_contains_elem(def_id, || hir::MaybeOwner::Phantom);
+        if let hir::MaybeOwner::Phantom = owner {
             let node = self.ast_index[def_id];
             match node {
                 AstOwner::NonOwner => {}
diff --git a/compiler/rustc_ast_lowering/src/lib.rs b/compiler/rustc_ast_lowering/src/lib.rs
index 2af47e11637..537c2a06183 100644
--- a/compiler/rustc_ast_lowering/src/lib.rs
+++ b/compiler/rustc_ast_lowering/src/lib.rs
@@ -368,8 +368,8 @@ fn index_crate<'a>(
     krate: &'a Crate,
 ) -> IndexVec<LocalDefId, AstOwner<'a>> {
     let mut indexer = Indexer { node_id_to_def_id, index: IndexVec::new() };
-    indexer.index.ensure_contains_elem(CRATE_DEF_ID, || AstOwner::NonOwner);
-    indexer.index[CRATE_DEF_ID] = AstOwner::Crate(krate);
+    *indexer.index.ensure_contains_elem(CRATE_DEF_ID, || AstOwner::NonOwner) =
+        AstOwner::Crate(krate);
     visit::walk_crate(&mut indexer, krate);
     return indexer.index;
 
@@ -386,22 +386,21 @@ fn index_crate<'a>(
 
         fn visit_item(&mut self, item: &'a ast::Item) {
             let def_id = self.node_id_to_def_id[&item.id];
-            self.index.ensure_contains_elem(def_id, || AstOwner::NonOwner);
-            self.index[def_id] = AstOwner::Item(item);
+            *self.index.ensure_contains_elem(def_id, || AstOwner::NonOwner) = AstOwner::Item(item);
             visit::walk_item(self, item)
         }
 
         fn visit_assoc_item(&mut self, item: &'a ast::AssocItem, ctxt: visit::AssocCtxt) {
             let def_id = self.node_id_to_def_id[&item.id];
-            self.index.ensure_contains_elem(def_id, || AstOwner::NonOwner);
-            self.index[def_id] = AstOwner::AssocItem(item, ctxt);
+            *self.index.ensure_contains_elem(def_id, || AstOwner::NonOwner) =
+                AstOwner::AssocItem(item, ctxt);
             visit::walk_assoc_item(self, item, ctxt);
         }
 
         fn visit_foreign_item(&mut self, item: &'a ast::ForeignItem) {
             let def_id = self.node_id_to_def_id[&item.id];
-            self.index.ensure_contains_elem(def_id, || AstOwner::NonOwner);
-            self.index[def_id] = AstOwner::ForeignItem(item);
+            *self.index.ensure_contains_elem(def_id, || AstOwner::NonOwner) =
+                AstOwner::ForeignItem(item);
             visit::walk_foreign_item(self, item);
         }
     }
diff --git a/compiler/rustc_hir_typeck/src/generator_interior/drop_ranges/mod.rs b/compiler/rustc_hir_typeck/src/generator_interior/drop_ranges/mod.rs
index f7b493bc224..d3685d21f25 100644
--- a/compiler/rustc_hir_typeck/src/generator_interior/drop_ranges/mod.rs
+++ b/compiler/rustc_hir_typeck/src/generator_interior/drop_ranges/mod.rs
@@ -268,8 +268,7 @@ impl DropRangesBuilder {
 
     fn node_mut(&mut self, id: PostOrderId) -> &mut NodeInfo {
         let size = self.num_values();
-        self.nodes.ensure_contains_elem(id, || NodeInfo::new(size));
-        &mut self.nodes[id]
+        self.nodes.ensure_contains_elem(id, || NodeInfo::new(size))
     }
 
     fn add_control_edge(&mut self, from: PostOrderId, to: PostOrderId) {
diff --git a/compiler/rustc_index/src/interval.rs b/compiler/rustc_index/src/interval.rs
index d809740c6ab..4605d42a15b 100644
--- a/compiler/rustc_index/src/interval.rs
+++ b/compiler/rustc_index/src/interval.rs
@@ -261,8 +261,7 @@ impl<R: Idx, C: Step + Idx> SparseIntervalMatrix<R, C> {
     }
 
     fn ensure_row(&mut self, row: R) -> &mut IntervalSet<C> {
-        self.rows.ensure_contains_elem(row, || IntervalSet::new(self.column_size));
-        &mut self.rows[row]
+        self.rows.ensure_contains_elem(row, || IntervalSet::new(self.column_size))
     }
 
     pub fn union_row(&mut self, row: R, from: &IntervalSet<C>) -> bool
diff --git a/compiler/rustc_index/src/vec.rs b/compiler/rustc_index/src/vec.rs
index ae2f52c513e..95600f7bcf1 100644
--- a/compiler/rustc_index/src/vec.rs
+++ b/compiler/rustc_index/src/vec.rs
@@ -236,12 +236,16 @@ impl<I: Idx, T> IndexVec<I, T> {
     /// `elem`; if that is already true, then has no
     /// effect. Otherwise, inserts new values as needed by invoking
     /// `fill_value`.
+    ///
+    /// Returns a reference to the `elem` entry.
     #[inline]
-    pub fn ensure_contains_elem(&mut self, elem: I, fill_value: impl FnMut() -> T) {
+    pub fn ensure_contains_elem(&mut self, elem: I, fill_value: impl FnMut() -> T) -> &mut T {
         let min_new_len = elem.index() + 1;
         if self.len() < min_new_len {
             self.raw.resize_with(min_new_len, fill_value);
         }
+
+        &mut self[elem]
     }
 
     #[inline]
@@ -446,20 +450,17 @@ impl<I: Idx, J: Idx> IndexSlice<I, J> {
 impl<I: Idx, T> IndexVec<I, Option<T>> {
     #[inline]
     pub fn insert(&mut self, index: I, value: T) -> Option<T> {
-        self.ensure_contains_elem(index, || None);
-        self[index].replace(value)
+        self.ensure_contains_elem(index, || None).replace(value)
     }
 
     #[inline]
     pub fn get_or_insert_with(&mut self, index: I, value: impl FnOnce() -> T) -> &mut T {
-        self.ensure_contains_elem(index, || None);
-        self[index].get_or_insert_with(value)
+        self.ensure_contains_elem(index, || None).get_or_insert_with(value)
     }
 
     #[inline]
     pub fn remove(&mut self, index: I) -> Option<T> {
-        self.ensure_contains_elem(index, || None);
-        self[index].take()
+        self.ensure_contains_elem(index, || None).take()
     }
 }
 
diff --git a/compiler/rustc_metadata/src/rmeta/table.rs b/compiler/rustc_metadata/src/rmeta/table.rs
index 364fa74ab7b..66e2518fa56 100644
--- a/compiler/rustc_metadata/src/rmeta/table.rs
+++ b/compiler/rustc_metadata/src/rmeta/table.rs
@@ -413,8 +413,8 @@ impl<I: Idx, const N: usize, T: FixedSizeEncoding<ByteArray = [u8; N]>> TableBui
             // > Space requirements could perhaps be optimized by using the HAMT `popcnt`
             // > trick (i.e. divide things into buckets of 32 or 64 items and then
             // > store bit-masks of which item in each bucket is actually serialized).
-            self.blocks.ensure_contains_elem(i, || [0; N]);
-            value.write_to_bytes(&mut self.blocks[i]);
+            let block = self.blocks.ensure_contains_elem(i, || [0; N]);
+            value.write_to_bytes(block);
         }
     }