stable_set: make small index type generic for testing, write better tests for the set type

Author: aiju

Cargo.lock

@@ -14,4 +14,5 @@ hashbrown = { version = "0.14.5", default-features = false, features = ["inline-
 [dev-dependencies]
 zwohash = "0.1.2"
 rand = "0.8.5"
-rand_pcg = "0.3.1"
+rand_pcg = "0.3.1"
+indexmap = "2.6.0"

stable_set/Cargo.toml

@@ -1,71 +1,91 @@
 use hashbrown::hash_table::HashTable;
 
+/// Trait for a small index type. You can assume `T: SmallIndex` means `T = u32` most of the time.
+///
+/// The only exception is for testing this crate, where we use `T = u8`.
+///
+/// This is basically `TryInto<usize> + TryFrom<usize>` with some convenience methods
+pub trait SmallIndex: Copy + std::fmt::Debug + TryInto<usize> + TryFrom<usize> {
+    fn into_usize(self) -> usize {
+        // the unreachable! should be straightforwardly impossible bc the conversion is infallible
+        // it would be nice if we could just rely on Into<usize> existing
+        // but std doesn't define it for u32 because you might be on a 16-bit platform...
+        self.try_into().unwrap_or_else(|_| unreachable!())
+    }
+    fn from_usize_unwrap(index: usize) -> Self {
+        // unfortunately .try_into().unwrap() doesn't work in the generic case
+        // because TryFrom::Error might not be Debug
+        index
+            .try_into()
+            .unwrap_or_else(|_| panic!("shouldn't happen: small index overflow in IndexTable"))
+    }
+}
+
+impl SmallIndex for u8 {}
+impl SmallIndex for u16 {}
+impl SmallIndex for u32 {}
+impl SmallIndex for usize {}
+
+// We could just use u32 directly for the small index,
+// but this is slightly less error-prone and allows us to use a smaller type for testing.
 #[derive(Debug, Clone)]
-pub enum IndexTable {
-    Small(HashTable<u32>),
+pub enum IndexTable<S> {
+    Small(HashTable<S>),
     Large(HashTable<usize>),
 }
 
 #[derive(Debug)]
-pub enum OccupiedEntry<'a> {
-    Small(hashbrown::hash_table::OccupiedEntry<'a, u32>),
+pub enum OccupiedEntry<'a, S> {
+    Small(hashbrown::hash_table::OccupiedEntry<'a, S>),
     Large(hashbrown::hash_table::OccupiedEntry<'a, usize>),
 }
 
 #[derive(Debug)]
-pub enum VacantEntry<'a> {
-    Small(hashbrown::hash_table::VacantEntry<'a, u32>),
+pub enum VacantEntry<'a, S> {
+    Small(hashbrown::hash_table::VacantEntry<'a, S>),
     Large(hashbrown::hash_table::VacantEntry<'a, usize>),
 }
 
 #[derive(Debug)]
 #[allow(dead_code)]
-pub enum AbsentEntry<'a> {
-    Small(hashbrown::hash_table::AbsentEntry<'a, u32>),
+pub enum AbsentEntry<'a, S> {
+    Small(hashbrown::hash_table::AbsentEntry<'a, S>),
     Large(hashbrown::hash_table::AbsentEntry<'a, usize>),
 }
 
 #[derive(Debug)]
 #[allow(dead_code)]
-pub enum Entry<'a> {
-    Occupied(OccupiedEntry<'a>),
-    Vacant(VacantEntry<'a>),
+pub enum Entry<'a, S> {
+    Occupied(OccupiedEntry<'a, S>),
+    Vacant(VacantEntry<'a, S>),
 }
 
-impl Default for IndexTable {
+impl<S> Default for IndexTable<S> {
     fn default() -> Self {
         IndexTable::Small(HashTable::new())
     }
 }
 
-impl IndexTable {
+impl<S: SmallIndex> IndexTable<S> {
     pub fn with_capacity(capacity: usize) -> Self {
-        if Self::try_as_small(capacity).is_ok() {
+        if S::try_from(capacity).is_ok() {
             IndexTable::Small(HashTable::with_capacity(capacity))
         } else {
             IndexTable::Large(HashTable::with_capacity(capacity))
         }
     }
     #[inline(always)]
-    fn as_small(index: usize) -> u32 {
-        Self::try_as_small(index).unwrap()
-    }
-    #[inline(always)]
-    fn try_as_small(index: usize) -> Result<u32, std::num::TryFromIntError> {
-        u32::try_from(index)
-    }
-    #[inline(always)]
     pub fn entry(
         &mut self,
         hash: u64,
         mut eq: impl FnMut(usize) -> bool,
         hasher: impl Fn(usize) -> u64,
-    ) -> Entry<'_> {
+    ) -> Entry<'_, S> {
         match self {
             IndexTable::Small(table) => match table.entry(
                 hash,
-                |&index| eq(index as usize),
-                |&index| hasher(index as usize),
+                |&index| eq(index.into_usize()),
+                |&index| hasher(index.into_usize()),
             ) {
                 hashbrown::hash_table::Entry::Occupied(entry) => {
                     Entry::Occupied(OccupiedEntry::Small(entry))
@@ -91,12 +111,14 @@ impl IndexTable {
         &mut self,
         hash: u64,
         mut eq: impl FnMut(usize) -> bool,
-    ) -> Result<OccupiedEntry<'_>, AbsentEntry<'_>> {
+    ) -> Result<OccupiedEntry<'_, S>, AbsentEntry<'_, S>> {
         match self {
-            IndexTable::Small(table) => match table.find_entry(hash, |&index| eq(index as usize)) {
-                Ok(entry) => Ok(OccupiedEntry::Small(entry)),
-                Err(entry) => Err(AbsentEntry::Small(entry)),
-            },
+            IndexTable::Small(table) => {
+                match table.find_entry(hash, |&index| eq(index.into_usize())) {
+                    Ok(entry) => Ok(OccupiedEntry::Small(entry)),
+                    Err(entry) => Err(AbsentEntry::Small(entry)),
+                }
+            }
             IndexTable::Large(table) => match table.find_entry(hash, |&index| eq(index)) {
                 Ok(entry) => Ok(OccupiedEntry::Large(entry)),
                 Err(entry) => Err(AbsentEntry::Large(entry)),
@@ -107,8 +129,8 @@ impl IndexTable {
     pub fn find(&self, hash: u64, mut eq: impl FnMut(usize) -> bool) -> Option<usize> {
         match self {
             IndexTable::Small(table) => table
-                .find(hash, |&index| eq(index as usize))
-                .map(|&index| index as usize),
+                .find(hash, |&index| eq(index.into_usize()))
+                .map(|&index| index.into_usize()),
             IndexTable::Large(table) => table.find(hash, |&index| eq(index)).copied(),
         }
     }
@@ -121,9 +143,13 @@ impl IndexTable {
         value: usize,
     ) -> Option<usize> {
         match self {
-            IndexTable::Small(table) => table
-                .find_mut(hash, |&index| eq(index as usize))
-                .map(|index_ref| std::mem::replace(index_ref, Self::as_small(value)) as usize),
+            IndexTable::Small(table) => {
+                table
+                    .find_mut(hash, |&index| eq(index.into_usize()))
+                    .map(|index_ref| {
+                        std::mem::replace(index_ref, S::from_usize_unwrap(value)).into_usize()
+                    })
+            }
             IndexTable::Large(table) => table
                 .find_mut(hash, |&index| eq(index))
                 .map(|index_ref| std::mem::replace(index_ref, value)),
@@ -138,7 +164,7 @@ impl IndexTable {
     }
     #[inline(always)]
     pub fn grow_for(&mut self, index: usize, hasher: impl Fn(usize) -> u64) {
-        if Self::try_as_small(index).is_err() && self.is_small() {
+        if S::try_from(index).is_err() && self.is_small() {
             self.grow_cold(hasher)
         }
     }
@@ -155,7 +181,7 @@ impl IndexTable {
         };
         new_table.reserve(old_table.len(), |&j| hasher(j));
         for i in old_table {
-            new_table.insert_unique(hasher(i as usize), i as usize, |&j| hasher(j));
+            new_table.insert_unique(hasher(i.into_usize()), i.into_usize(), |&j| hasher(j));
         }
     }
     #[inline(always)]
@@ -175,9 +201,9 @@ impl IndexTable {
     #[inline(always)]
     pub fn retain(&mut self, mut f: impl FnMut(usize) -> Option<usize>) {
         match self {
-            IndexTable::Small(table) => table.retain(|index| match f(*index as usize) {
+            IndexTable::Small(table) => table.retain(|index| match f((*index).into_usize()) {
                 Some(new_index) => {
-                    *index = Self::as_small(new_index);
+                    *index = S::from_usize_unwrap(new_index);
                     true
                 }
                 None => false,
@@ -195,7 +221,9 @@ impl IndexTable {
     pub fn reserve(&mut self, additional: usize, hasher: impl Fn(usize) -> u64) {
         self.grow_for((self.len() + additional).saturating_sub(1), &hasher);
         match self {
-            IndexTable::Small(table) => table.reserve(additional, |&index| hasher(index as usize)),
+            IndexTable::Small(table) => {
+                table.reserve(additional, |&index| hasher(index.into_usize()))
+            }
             IndexTable::Large(table) => table.reserve(additional, |&index| hasher(index)),
         }
     }
@@ -217,12 +245,12 @@ impl IndexTable {
     }
 }
 
-impl<'a> VacantEntry<'a> {
+impl<'a, S: SmallIndex> VacantEntry<'a, S> {
     #[inline(always)]
-    pub fn insert(self, index: usize) -> OccupiedEntry<'a> {
+    pub fn insert(self, index: usize) -> OccupiedEntry<'a, S> {
         match self {
             VacantEntry::Small(entry) => {
-                OccupiedEntry::Small(entry.insert(IndexTable::as_small(index)))
+                OccupiedEntry::Small(entry.insert(S::from_usize_unwrap(index)))
             }
             VacantEntry::Large(entry) => OccupiedEntry::Large(entry.insert(index)),
         }
@@ -237,30 +265,34 @@ impl<'a> VacantEntry<'a> {
         value: usize,
     ) -> Option<usize> {
         match self {
-            VacantEntry::Small(entry) => entry.into_table()
-                .find_mut(hash, |&index| eq(index as usize))
-                .map(|index_ref| std::mem::replace(index_ref, IndexTable::as_small(value)) as usize),
-            VacantEntry::Large(entry) => entry.into_table()
+            VacantEntry::Small(entry) => entry
+                .into_table()
+                .find_mut(hash, |&index| eq(index.into_usize()))
+                .map(|index_ref| {
+                    std::mem::replace(index_ref, S::from_usize_unwrap(value)).into_usize()
+                }),
+            VacantEntry::Large(entry) => entry
+                .into_table()
                 .find_mut(hash, |&index| eq(index))
                 .map(|index_ref| std::mem::replace(index_ref, value)),
         }
     }
 }
 
-impl<'a> OccupiedEntry<'a> {
+impl<'a, S: SmallIndex> OccupiedEntry<'a, S> {
     #[inline(always)]
     pub fn get(&self) -> usize {
         match self {
-            OccupiedEntry::Small(entry) => *entry.get() as usize,
+            OccupiedEntry::Small(entry) => (*entry.get()).into_usize(),
             OccupiedEntry::Large(entry) => *entry.get(),
         }
     }
     #[inline(always)]
-    pub fn remove(self) -> (usize, VacantEntry<'a>) {
+    pub fn remove(self) -> (usize, VacantEntry<'a, S>) {
         match self {
             OccupiedEntry::Small(entry) => {
                 let (index, entry) = entry.remove();
-                (index as usize, VacantEntry::Small(entry))
+                (index.into_usize(), VacantEntry::Small(entry))
             }
             OccupiedEntry::Large(entry) => {
                 let (index, entry) = entry.remove();