Draft: TriHashMap Entry API

crates/iddqd-test-utils/src/naive_map.rs

@@ -190,6 +190,84 @@ impl NaiveMap {
         Some(self.items.remove(index))
     }
 
+    /// Removes and returns every item covered by a `TriHashMap` entry keyed on
+    /// `(key1, key2, key3)`, i.e. every distinct item matching `key1`, `key2`,
+    /// or `key3`.
+    ///
+    /// Mirrors `tri_hash_map::OccupiedEntry::remove`. The returned items are in
+    /// first-key-hit order.
+    pub fn entry_remove123(
+        &mut self,
+        key1: u8,
+        key2: char,
+        key3: &str,
+    ) -> Vec<TestItem> {
+        fn push_unique(indexes: &mut Vec<usize>, index: Option<usize>) {
+            if let Some(index) = index {
+                if !indexes.contains(&index) {
+                    indexes.push(index);
+                }
+            }
+        }
+
+        let mut key_order = Vec::new();
+
+        push_unique(
+            &mut key_order,
+            self.items.iter().position(|e| e.key1 == key1),
+        );
+        push_unique(
+            &mut key_order,
+            self.items.iter().position(|e| e.key2 == key2),
+        );
+        push_unique(
+            &mut key_order,
+            self.items.iter().position(|e| e.key3 == key3),
+        );
+
+        let mut remove_order = key_order.clone();
+        remove_order.sort_unstable_by(|a, b| b.cmp(a));
+
+        let mut removed = Vec::with_capacity(remove_order.len());
+        for index in remove_order {
+            removed.push((index, self.items.remove(index)));
+        }
+
+        key_order
+            .into_iter()
+            .map(|index| {
+                let removed_index = removed
+                    .iter()
+                    .position(|(removed_index, _)| *removed_index == index)
+                    .expect("index was removed");
+                removed.swap_remove(removed_index).1
+            })
+            .collect()
+    }
+
+    /// Mirrors the test harness behavior for
+    /// `TriHashMap::entry(...).Occupied(...).insert(...)`.
+    ///
+    /// Returns `None` for a vacant entry and does not insert. Returns
+    /// `Some(removed)` for an occupied entry, where `removed` is in
+    /// first-key-hit order.
+    pub fn entry_insert_overwrite123(
+        &mut self,
+        item: TestItem,
+    ) -> Option<Vec<TestItem>> {
+        let occupied = self.get1(item.key1).is_some()
+            || self.get2(item.key2).is_some()
+            || self.get3(&item.key3).is_some();
+
+        if !occupied {
+            return None;
+        }
+
+        let removed = self.entry_remove123(item.key1, item.key2, &item.key3);
+        self.items.push(item);
+        Some(removed)
+    }
+
     pub fn iter(&self) -> impl Iterator<Item = &TestItem> {
         self.items.iter()
     }

crates/iddqd-test-utils/src/panic_safety.rs

@@ -407,12 +407,10 @@ pub fn record_observation(
 #[derive(Clone, Copy, Debug, Default)]
 pub struct PanickyAlloc<A>(pub A);
 
-// SAFETY:
-//
-// * On the non-panic path, forwards to the wrapped allocator.
-// * On the armed path, panics before any inner allocation, so no pointer is
-//   observable to the caller.
 #[cfg(all(feature = "default-hasher", feature = "allocator-api2"))]
+// SAFETY: On the non-panic path, this forwards to the wrapped allocator. On
+// the armed path, it panics before any inner allocation, so no pointer is
+// observable to the caller.
 unsafe impl<A: Allocator> Allocator for PanickyAlloc<A> {
     fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
         observe_panicky_call("alloc");

crates/iddqd/src/bi_hash_map/entry_indexes.rs

@@ -1,17 +1,36 @@
-use crate::support::ItemIndex;
+use crate::support::{
+    ItemIndex,
+    entry::{
+        EntryIndexes as SupportEntryIndexes, EntryLookup, NonUniqueIndexes,
+    },
+};
 
 #[derive(Clone, Copy, Debug)]
 pub(super) enum EntryIndexes {
     Unique(ItemIndex),
     NonUnique {
-        // Invariant: at least one index is Some, and indexes are different from
-        // each other.
+        // Invariant: at least one index is Some, and indexes are not all the
+        // same Some value.
         index1: Option<ItemIndex>,
         index2: Option<ItemIndex>,
     },
 }
 
 impl EntryIndexes {
+    #[inline]
+    pub(super) fn classify(
+        index1: Option<ItemIndex>,
+        index2: Option<ItemIndex>,
+    ) -> EntryLookup<2> {
+        SupportEntryIndexes::new([index1, index2]).classify()
+    }
+
+    #[inline]
+    pub(super) fn from_non_unique(indexes: NonUniqueIndexes<2>) -> Self {
+        let [index1, index2] = *indexes.indexes();
+        EntryIndexes::NonUnique { index1, index2 }
+    }
+
     #[inline]
     pub(super) fn is_unique(&self) -> bool {
         matches!(self, EntryIndexes::Unique(_))

crates/iddqd/src/bi_hash_map/imp.rs

@@ -13,6 +13,7 @@ use crate::{
         ItemIndex,
         alloc::{Allocator, Global, global_alloc},
         borrow::DormantMutRef,
+        entry::EntryLookup,
         fmt_utils::StrDisplayAsDebug,
         hash_table,
         item_set::ItemSet,
@@ -1918,33 +1919,33 @@ impl<T: BiHashItem, S: Clone + BuildHasher, A: Allocator> BiHashMap<T, S, A> {
             (index1, index2)
         };
 
-        match (index1, index2) {
-            (Some(index1), Some(index2)) if index1 == index2 => {
+        match EntryIndexes::classify(index1, index2) {
+            EntryLookup::Unique(index) => {
                 // The item is already in the map.
                 drop(key1);
                 Entry::Occupied(
                     // SAFETY: `map` is not used after this point.
                     unsafe {
                         OccupiedEntry::new(
                             dormant_map,
-                            EntryIndexes::Unique(index1),
+                            EntryIndexes::Unique(index),
                         )
                     },
                 )
             }
-            (None, None) => {
+            EntryLookup::Vacant => {
                 let hashes = map.tables.make_hashes::<T>(&key1, &key2);
                 Entry::Vacant(
                     // SAFETY: `map` is not used after this point.
                     unsafe { VacantEntry::new(dormant_map, hashes) },
                 )
             }
-            (index1, index2) => Entry::Occupied(
+            EntryLookup::NonUnique(indexes) => Entry::Occupied(
                 // SAFETY: `map` is not used after this point.
                 unsafe {
                     OccupiedEntry::new(
                         dormant_map,
-                        EntryIndexes::NonUnique { index1, index2 },
+                        EntryIndexes::from_non_unique(indexes),
                     )
                 },
             ),