panic safety PBTs now cover allocator panics

crates/iddqd/src/bi_hash_map/imp.rs

@@ -865,11 +865,22 @@ impl<T: BiHashItem, S: Clone + BuildHasher, A: Allocator> BiHashMap<T, S, A> {
     /// # }
     /// ```
     pub fn shrink_to_fit(&mut self) {
-        let remap = self.items.shrink_to_fit();
+        // Sequence this carefully.
+        //
+        // * First, compact the item set. This does not allocate through A
+        //   (it allocates a small remap buffer through the global allocator),
+        //   and returns a remapper.
+        // * Then, remap the tables using the remapper.
+        // * Finally, shrink the capacities of the tables and items.
+        //
+        // An allocator panic during either capacity shrink leaves the tables
+        // and items already in sync, because remap has already been committed.
+        let remap = self.items.compact();
         if !remap.is_identity() {
             self.tables.k1_to_item.remap_indexes(&remap);
             self.tables.k2_to_item.remap_indexes(&remap);
         }
+        self.items.shrink_capacity_to_fit();
         self.tables.k1_to_item.shrink_to_fit();
         self.tables.k2_to_item.shrink_to_fit();
     }
@@ -916,11 +927,13 @@ impl<T: BiHashItem, S: Clone + BuildHasher, A: Allocator> BiHashMap<T, S, A> {
     /// # }
     /// ```
     pub fn shrink_to(&mut self, min_capacity: usize) {
-        let remap = self.items.shrink_to(min_capacity);
+        // See `shrink_to_fit` for the rationale behind the sequence.
+        let remap = self.items.compact();
         if !remap.is_identity() {
             self.tables.k1_to_item.remap_indexes(&remap);
             self.tables.k2_to_item.remap_indexes(&remap);
         }
+        self.items.shrink_capacity_to(min_capacity);
         self.tables.k1_to_item.shrink_to(min_capacity);
         self.tables.k2_to_item.shrink_to(min_capacity);
     }

crates/iddqd/src/id_hash_map/imp.rs

@@ -735,10 +735,21 @@ impl<T: IdHashItem, S: Clone + BuildHasher, A: Allocator> IdHashMap<T, S, A> {
     /// # }
     /// ```
     pub fn shrink_to_fit(&mut self) {
-        let remap = self.items.shrink_to_fit();
+        // Sequence this carefully.
+        //
+        // * First, compact the item set. This does not allocate through A
+        //   (it allocates a small remap buffer through the global allocator),
+        //   and returns a remapper.
+        // * Then, remap the tables using the remapper.
+        // * Finally, shrink the capacities of the tables and items.
+        //
+        // An allocator panic during either capacity shrink leaves the tables
+        // and items already in sync, because remap has already been committed.
+        let remap = self.items.compact();
         if !remap.is_identity() {
             self.tables.key_to_item.remap_indexes(&remap);
         }
+        self.items.shrink_capacity_to_fit();
         self.tables.key_to_item.shrink_to_fit();
     }
 
@@ -780,10 +791,12 @@ impl<T: IdHashItem, S: Clone + BuildHasher, A: Allocator> IdHashMap<T, S, A> {
     /// # }
     /// ```
     pub fn shrink_to(&mut self, min_capacity: usize) {
-        let remap = self.items.shrink_to(min_capacity);
+        // See `shrink_to_fit` for the rationale behind the sequence.
+        let remap = self.items.compact();
         if !remap.is_identity() {
             self.tables.key_to_item.remap_indexes(&remap);
         }
+        self.items.shrink_capacity_to(min_capacity);
         self.tables.key_to_item.shrink_to(min_capacity);
     }
 

crates/iddqd/src/id_ord_map/imp.rs

@@ -435,10 +435,23 @@ impl<T: IdOrdItem> IdOrdMap<T> {
     /// assert!(map.capacity() >= 2);
     /// ```
     pub fn shrink_to_fit(&mut self) {
-        let remap = self.items.shrink_to_fit();
+        // Sequence this carefully.
+        //
+        // * First, compact the item set. This does not allocate through A
+        //   (it allocates a small remap buffer through the global allocator),
+        //   and returns a remapper.
+        // * Then, remap the table using the remapper.
+        // * Finally, shrink the capacity of the items. (BTreeMap has no
+        //   capacity to shrink.)
+        //
+        // An allocator panic during the capacity shrink leaves the table
+        // and items already in sync, because remap has already been
+        // committed.
+        let remap = self.items.compact();
         if !remap.is_identity() {
             self.tables.key_to_item.remap_indexes(&remap);
         }
+        self.items.shrink_capacity_to_fit();
     }
 
     /// Shrinks the capacity of the map with a lower limit. It will drop
@@ -481,10 +494,12 @@ impl<T: IdOrdItem> IdOrdMap<T> {
     /// assert!(map.capacity() >= 2);
     /// ```
     pub fn shrink_to(&mut self, min_capacity: usize) {
-        let remap = self.items.shrink_to(min_capacity);
+        // See `shrink_to_fit` for the rationale behind the sequence.
+        let remap = self.items.compact();
         if !remap.is_identity() {
             self.tables.key_to_item.remap_indexes(&remap);
         }
+        self.items.shrink_capacity_to(min_capacity);
     }
 
     /// Iterates over the items in the map.

crates/iddqd/src/support/btree_table.rs

@@ -286,9 +286,9 @@ impl MapBTreeTable {
 
     /// Rewrites every stored index via `remap`.
     ///
-    /// Called after [`ItemSet::shrink_to_fit`] or [`ItemSet::shrink_to`]
-    /// compacts the backing items buffer. Each stored `Index` needs to be
-    /// rewritten to point at the item's new position.
+    /// Called after [`ItemSet::compact`] compacts the backing items buffer.
+    /// Each stored `Index` needs to be rewritten to point at the item's new
+    /// position.
     ///
     /// We do not rebuild the tree. [`IndexRemap`] preserves relative
     /// order, so the tree's iteration order — which is the user's
@@ -300,8 +300,7 @@ impl MapBTreeTable {
     /// In-place mutation through `&Index` is provided by [`IndexCell`],
     /// which uses an `AtomicU32` for `&self`-based stores.
     ///
-    /// [`ItemSet::shrink_to_fit`]: super::item_set::ItemSet::shrink_to_fit
-    /// [`ItemSet::shrink_to`]: super::item_set::ItemSet::shrink_to
+    /// [`ItemSet::compact`]: super::item_set::ItemSet::compact
     pub(crate) fn remap_indexes(&mut self, remap: &IndexRemap) {
         for idx in self.items.keys() {
             let new = remap.remap(idx.value());
@@ -700,7 +699,8 @@ mod tests {
         }
         set.remove(ItemIndex::new(1));
         set.remove(ItemIndex::new(3));
-        let remap = set.shrink_to_fit();
+        let remap = set.compact();
+        set.shrink_capacity_to_fit();
         assert!(!remap.is_identity(), "remap should carry two holes");
 
         // A MapBTreeTable populated to match the pre-compaction live indexes

crates/iddqd/src/support/hash_table.rs

@@ -252,13 +252,11 @@ impl<A: Allocator> MapHashTable<A> {
 
     /// Rewrites every stored index via `remap`.
     ///
-    /// Called after [`ItemSet::shrink_to_fit`] / [`ItemSet::shrink_to`]
-    /// compacts the backing items buffer. We store hashes of *keys* (not of
-    /// indexes), so rewriting an index does not invalidate its hash and no
-    /// rehash is needed.
+    /// Called after [`ItemSet::compact`] compacts the backing items buffer.
+    /// We store hashes of *keys* (not of indexes), so rewriting an index does
+    /// not invalidate its hash and no rehash is needed.
     ///
-    /// [`ItemSet::shrink_to_fit`]: super::item_set::ItemSet::shrink_to_fit
-    /// [`ItemSet::shrink_to`]: super::item_set::ItemSet::shrink_to
+    /// [`ItemSet::compact`]: super::item_set::ItemSet::compact
     pub(crate) fn remap_indexes(&mut self, remap: &IndexRemap) {
         for stored in self.items.iter_mut() {
             stored.ix = remap.remap(stored.ix);

crates/iddqd/src/support/item_set.rs

@@ -63,10 +63,11 @@ use core::{
 
 /// A remap from old (pre-compaction) to new (post-compaction) indexes.
 ///
-/// Produced by [`ItemSet::shrink_to_fit`] and [`ItemSet::shrink_to`],
-/// consumed by the outer tables (hash / btree index tables) so they
-/// can rewrite their stored indexes to point at the compacted `items`
-/// buffer.
+/// Produced by [`ItemSet::compact`], and consumed by the outer tables (hash /
+/// btree index tables) so they can rewrite their stored indexes to point at
+/// the compacted `items` buffer. The capacity shrink that typically follows
+/// is a separate step so that an allocator panic during the
+/// shrink cannot leave the tables pointing at pre-compaction indexes.
 ///
 /// Two cases:
 ///
@@ -95,7 +96,7 @@ impl IndexRemap {
     ///
     /// Panics if `old` was a slot that compaction vacated. This indicates a
     /// caller bug: those indexes should already have been removed from the
-    /// outer index before `shrink_to_fit` was called.
+    /// outer index before [`ItemSet::compact`] was called.
     #[inline]
     pub(crate) fn remap(&self, old: ItemIndex) -> ItemIndex {
         match self {
@@ -512,7 +513,7 @@ impl<T, A: Allocator> ItemSet<T, A> {
     ///
     /// `items` is not truncated here, even for a trailing remove. The vacated
     /// slot stays in place until reused by the next insert or reclaimed by
-    /// [`shrink_to_fit`](Self::shrink_to_fit).
+    /// [`compact`](Self::compact).
     #[inline]
     pub(crate) fn remove(&mut self, index: ItemIndex) -> Option<T> {
         let slot = self.items.get_mut(index.as_u32() as usize)?;
@@ -572,23 +573,37 @@ impl<T, A: Allocator> ItemSet<T, A> {
         self.items.reserve(additional);
     }
 
+    /// Shrinks the backing buffer's capacity to fit the current length.
+    ///
+    /// Must be called *after* [`compact`](Self::compact) has been run and
+    /// the outer tables have been remapped via the returned [`IndexRemap`].
+    /// Splitting capacity-shrinking out of compaction means that an
+    /// allocator panic in this call cannot leave the tables pointing at
+    /// pre-compaction indexes (by the time we get here, the tables and
+    /// `items` are already in sync).
     #[inline]
-    pub(crate) fn shrink_to_fit(&mut self) -> IndexRemap {
-        let remap = self.compact();
+    pub(crate) fn shrink_capacity_to_fit(&mut self) {
         self.items.shrink_to_fit();
-        remap
     }
 
+    /// Shrinks the backing buffer's capacity, leaving at least `min_capacity`
+    /// slots reserved.
+    ///
+    /// See [`shrink_capacity_to_fit`](Self::shrink_capacity_to_fit) for the
+    /// panic-safety rationale for splitting this from compaction.
     #[inline]
-    pub(crate) fn shrink_to(&mut self, min_capacity: usize) -> IndexRemap {
-        let remap = self.compact();
+    pub(crate) fn shrink_capacity_to(&mut self, min_capacity: usize) {
         self.items.shrink_to(min_capacity);
-        remap
     }
 
     /// Moves every live slot down to fill `Vacant` holes, truncates
     /// `items` to its new length, and clears the free chain.
-    fn compact(&mut self) -> IndexRemap {
+    ///
+    /// Does *not* shrink the underlying buffer's capacity. Callers should
+    /// remap any externally-stored indexes via the returned [`IndexRemap`]
+    /// before calling [`shrink_capacity_to_fit`](Self::shrink_capacity_to_fit)
+    /// or [`shrink_capacity_to`](Self::shrink_capacity_to).
+    pub(crate) fn compact(&mut self) -> IndexRemap {
         let pre_len = self.items.len();
         if pre_len == self.len as usize {
             // Already compact, so there's nothing to remap.
@@ -1140,15 +1155,16 @@ mod tests {
     }
 
     #[test]
-    fn shrink_to_fit_compacts_middle_holes() {
+    fn compact_fills_middle_holes() {
         let mut set = ItemSet::<u32, Global>::new();
         for i in 0..5 {
             set.assert_can_grow().insert(i * 10);
         }
         set.remove(ix(1)).expect("slot was occupied");
         set.remove(ix(3)).expect("slot was occupied");
 
-        let remap = set.shrink_to_fit();
+        let remap = set.compact();
+        set.shrink_capacity_to_fit();
 
         assert_eq!(set.len(), 3);
         set.validate(ValidateCompact::Compact).unwrap();
@@ -1161,15 +1177,16 @@ mod tests {
     }
 
     #[test]
-    fn shrink_to_fit_without_holes_returns_empty_remap() {
+    fn compact_without_holes_returns_identity_remap() {
         let mut set = ItemSet::<u32, Global>::new();
         for i in 0..4 {
             set.assert_can_grow().insert(i);
         }
-        let remap = set.shrink_to_fit();
+        let remap = set.compact();
+        set.shrink_capacity_to_fit();
         assert!(remap.is_identity());
         set.validate(ValidateCompact::Compact)
-            .expect("a hole-free set is trivially compact after shrink_to_fit");
+            .expect("a hole-free set is trivially compact after compact");
     }
 
     #[test]

crates/iddqd/src/tri_hash_map/imp.rs

@@ -986,12 +986,23 @@ impl<T: TriHashItem, S: Clone + BuildHasher, A: Allocator> TriHashMap<T, S, A> {
     /// # }
     /// ```
     pub fn shrink_to_fit(&mut self) {
-        let remap = self.items.shrink_to_fit();
+        // Sequence this carefully.
+        //
+        // * First, compact the item set. This does not allocate through A
+        //   (it allocates a small remap buffer through the global allocator),
+        //   and returns a remapper.
+        // * Then, remap the tables using the remapper.
+        // * Finally, shrink the capacities of the tables and items.
+        //
+        // An allocator panic during either capacity shrink leaves the tables
+        // and items already in sync, because remap has already been committed.
+        let remap = self.items.compact();
         if !remap.is_identity() {
             self.tables.k1_to_item.remap_indexes(&remap);
             self.tables.k2_to_item.remap_indexes(&remap);
             self.tables.k3_to_item.remap_indexes(&remap);
         }
+        self.items.shrink_capacity_to_fit();
         self.tables.k1_to_item.shrink_to_fit();
         self.tables.k2_to_item.shrink_to_fit();
         self.tables.k3_to_item.shrink_to_fit();
@@ -1054,12 +1065,14 @@ impl<T: TriHashItem, S: Clone + BuildHasher, A: Allocator> TriHashMap<T, S, A> {
     /// # }
     /// ```
     pub fn shrink_to(&mut self, min_capacity: usize) {
-        let remap = self.items.shrink_to(min_capacity);
+        // See `shrink_to_fit` for the rationale behind the sequence.
+        let remap = self.items.compact();
         if !remap.is_identity() {
             self.tables.k1_to_item.remap_indexes(&remap);
             self.tables.k2_to_item.remap_indexes(&remap);
             self.tables.k3_to_item.remap_indexes(&remap);
         }
+        self.items.shrink_capacity_to(min_capacity);
         self.tables.k1_to_item.shrink_to(min_capacity);
         self.tables.k2_to_item.shrink_to(min_capacity);
         self.tables.k3_to_item.shrink_to(min_capacity);