diff --git a/Cargo.lock b/Cargo.lock index 9a6aad0b..fe2a2f5e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -365,7 +365,6 @@ dependencies = [ "iddqd-test-utils", "proptest", "ref-cast", - "rustc-hash", "schemars", "serde", "serde_core", @@ -710,12 +709,6 @@ dependencies = [ "memchr", ] -[[package]] -name = "rustc-hash" -version = "2.1.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "357703d41365b4b27c590e3ed91eabb1b663f07c4c084095e60cbed4362dff0d" - [[package]] name = "rustix" version = "0.38.44" diff --git a/Cargo.toml b/Cargo.toml index d490a5e1..75292176 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -32,7 +32,6 @@ iddqd = { path = "crates/iddqd", default-features = false } iddqd-test-utils = { path = "crates/iddqd-test-utils" } proptest = { version = "1.7.0", default-features = false, features = ["std"] } ref-cast = "1.0.24" -rustc-hash = { version = "2.1.1", default-features = false } schemars = "0.8.22" serde = "1.0.223" serde_core = "1.0.223" diff --git a/crates/iddqd/Cargo.toml b/crates/iddqd/Cargo.toml index eeb096a0..d822b9e3 100644 --- a/crates/iddqd/Cargo.toml +++ b/crates/iddqd/Cargo.toml @@ -29,7 +29,6 @@ equivalent.workspace = true foldhash = { workspace = true, optional = true } hashbrown.workspace = true ref-cast = { workspace = true, optional = true } -rustc-hash.workspace = true schemars = { workspace = true, optional = true } serde_core = { workspace = true, optional = true } serde_json = { workspace = true, optional = true } @@ -50,7 +49,7 @@ default-hasher = ["dep:foldhash", "iddqd-test-utils/default-hasher"] proptest = ["dep:proptest"] schemars08 = ["dep:schemars", "dep:serde_json", "serde"] serde = ["dep:serde_core", "iddqd-test-utils/serde"] -std = ["dep:foldhash", "iddqd-test-utils/std", "rustc-hash/std"] +std = ["dep:foldhash", "iddqd-test-utils/std"] # Internal-only feature for testing that schemars/preserve_order works. internal-schemars08-preserve-order = ["schemars08", "schemars/preserve_order"] diff --git a/crates/iddqd/src/bi_hash_map/imp.rs b/crates/iddqd/src/bi_hash_map/imp.rs index ca8e5782..86e8c7b6 100644 --- a/crates/iddqd/src/bi_hash_map/imp.rs +++ b/crates/iddqd/src/bi_hash_map/imp.rs @@ -819,9 +819,7 @@ impl BiHashMap { &mut self, additional: usize, ) -> Result<(), crate::errors::TryReserveError> { - self.items - .try_reserve(additional) - .map_err(crate::errors::TryReserveError::from_hashbrown)?; + self.items.try_reserve(additional)?; let items = &self.items; let state = &self.tables.state; self.tables @@ -872,7 +870,11 @@ impl BiHashMap { /// # } /// ``` pub fn shrink_to_fit(&mut self) { - self.items.shrink_to_fit(); + let remap = self.items.shrink_to_fit(); + if !remap.is_identity() { + self.tables.k1_to_item.remap_indexes(&remap); + self.tables.k2_to_item.remap_indexes(&remap); + } let items = &self.items; let state = &self.tables.state; self.tables @@ -925,7 +927,11 @@ impl BiHashMap { /// # } /// ``` pub fn shrink_to(&mut self, min_capacity: usize) { - self.items.shrink_to(min_capacity); + let remap = self.items.shrink_to(min_capacity); + if !remap.is_identity() { + self.tables.k1_to_item.remap_indexes(&remap); + self.tables.k2_to_item.remap_indexes(&remap); + } let items = &self.items; let state = &self.tables.state; self.tables @@ -1052,7 +1058,7 @@ impl BiHashMap { self.tables.validate(self.len(), compactness)?; // Check that the indexes are all correct. - for (&ix, item) in self.items.iter() { + for (ix, item) in self.items.iter() { let key1 = item.key1(); let key2 = item.key2(); diff --git a/crates/iddqd/src/bi_hash_map/iter.rs b/crates/iddqd/src/bi_hash_map/iter.rs index f2e4c182..d077c6ed 100644 --- a/crates/iddqd/src/bi_hash_map/iter.rs +++ b/crates/iddqd/src/bi_hash_map/iter.rs @@ -2,13 +2,11 @@ use super::{RefMut, tables::BiHashMapTables}; use crate::{ BiHashItem, DefaultHashBuilder, support::{ - ItemIndex, - alloc::{AllocWrapper, Allocator, Global}, - item_set::ItemSet, + alloc::{Allocator, Global}, + item_set::{self, ItemSet}, }, }; use core::{hash::BuildHasher, iter::FusedIterator}; -use hashbrown::hash_map; /// An iterator over the elements of a [`BiHashMap`] by shared reference. /// Created by [`BiHashMap::iter`]. @@ -21,7 +19,7 @@ use hashbrown::hash_map; /// [`HashMap`]: std::collections::HashMap #[derive(Clone, Debug, Default)] pub struct Iter<'a, T: BiHashItem> { - inner: hash_map::Values<'a, ItemIndex, T>, + inner: item_set::Values<'a, T>, } impl<'a, T: BiHashItem> Iter<'a, T> { @@ -46,7 +44,6 @@ impl ExactSizeIterator for Iter<'_, T> { } } -// hash_map::Iter is a FusedIterator, so Iter is as well. impl FusedIterator for Iter<'_, T> {} /// An iterator over the elements of a [`BiHashMap`] by mutable reference. @@ -68,7 +65,7 @@ pub struct IterMut< A: Allocator = Global, > { tables: &'a BiHashMapTables, - inner: hash_map::ValuesMut<'a, ItemIndex, T>, + inner: item_set::ValuesMut<'a, T>, } impl<'a, T: BiHashItem, S: Clone + BuildHasher, A: Allocator> @@ -104,7 +101,6 @@ impl ExactSizeIterator } } -// hash_map::IterMut is a FusedIterator, so IterMut is as well. impl FusedIterator for IterMut<'_, T, S, A> { @@ -121,7 +117,7 @@ impl FusedIterator /// [`HashMap`]: std::collections::HashMap #[derive(Debug)] pub struct IntoIter { - inner: hash_map::IntoValues>, + inner: item_set::IntoValues, } impl IntoIter { @@ -146,5 +142,4 @@ impl ExactSizeIterator for IntoIter { } } -// hash_map::IterMut is a FusedIterator, so IterMut is as well. impl FusedIterator for IntoIter {} diff --git a/crates/iddqd/src/errors.rs b/crates/iddqd/src/errors.rs index aa254702..e611fac6 100644 --- a/crates/iddqd/src/errors.rs +++ b/crates/iddqd/src/errors.rs @@ -77,9 +77,9 @@ enum TryReserveErrorKind { /// (usually `isize::MAX` bytes). CapacityOverflow, - /// The memory allocator returned an error + /// The memory allocator returned an error. AllocError { - /// The layout of the allocation request that failed + /// The layout of the allocation request that failed. layout: core::alloc::Layout, }, } @@ -97,6 +97,20 @@ impl TryReserveError { }; Self { kind } } + + /// Converts from an `allocator_api2` `TryReserveError`. + pub(crate) fn from_allocator_api2( + error: allocator_api2::collections::TryReserveError, + ) -> Self { + use allocator_api2::collections::TryReserveErrorKind as Kind; + let kind = match error.kind() { + Kind::CapacityOverflow => TryReserveErrorKind::CapacityOverflow, + Kind::AllocError { layout, .. } => { + TryReserveErrorKind::AllocError { layout } + } + }; + Self { kind } + } } impl fmt::Display for TryReserveError { diff --git a/crates/iddqd/src/id_hash_map/imp.rs b/crates/iddqd/src/id_hash_map/imp.rs index eecde817..79fa35dc 100644 --- a/crates/iddqd/src/id_hash_map/imp.rs +++ b/crates/iddqd/src/id_hash_map/imp.rs @@ -695,9 +695,7 @@ impl IdHashMap { &mut self, additional: usize, ) -> Result<(), crate::errors::TryReserveError> { - self.items - .try_reserve(additional) - .map_err(crate::errors::TryReserveError::from_hashbrown)?; + self.items.try_reserve(additional)?; let items = &self.items; let state = &self.tables.state; self.tables @@ -740,7 +738,10 @@ impl IdHashMap { /// # } /// ``` pub fn shrink_to_fit(&mut self) { - self.items.shrink_to_fit(); + let remap = self.items.shrink_to_fit(); + if !remap.is_identity() { + self.tables.key_to_item.remap_indexes(&remap); + } let items = &self.items; let state = &self.tables.state; self.tables @@ -786,7 +787,10 @@ impl IdHashMap { /// # } /// ``` pub fn shrink_to(&mut self, min_capacity: usize) { - self.items.shrink_to(min_capacity); + let remap = self.items.shrink_to(min_capacity); + if !remap.is_identity() { + self.tables.key_to_item.remap_indexes(&remap); + } let items = &self.items; let state = &self.tables.state; self.tables @@ -895,7 +899,7 @@ impl IdHashMap { self.tables.validate(self.len(), compactness)?; // Check that the indexes are all correct. - for (&ix, item) in self.items.iter() { + for (ix, item) in self.items.iter() { let key = item.key(); let Some(ix1) = self.find_index(&key) else { return Err(ValidationError::general(format!( diff --git a/crates/iddqd/src/id_hash_map/iter.rs b/crates/iddqd/src/id_hash_map/iter.rs index 1e293f50..d77b513d 100644 --- a/crates/iddqd/src/id_hash_map/iter.rs +++ b/crates/iddqd/src/id_hash_map/iter.rs @@ -2,13 +2,11 @@ use super::{RefMut, tables::IdHashMapTables}; use crate::{ DefaultHashBuilder, IdHashItem, support::{ - ItemIndex, - alloc::{AllocWrapper, Allocator, Global}, - item_set::ItemSet, + alloc::{Allocator, Global}, + item_set::{self, ItemSet}, }, }; use core::{hash::BuildHasher, iter::FusedIterator}; -use hashbrown::hash_map; /// An iterator over the elements of a [`IdHashMap`] by shared reference. /// Created by [`IdHashMap::iter`]. @@ -21,7 +19,7 @@ use hashbrown::hash_map; /// [`HashMap`]: std::collections::HashMap #[derive(Clone, Debug, Default)] pub struct Iter<'a, T: IdHashItem> { - inner: hash_map::Values<'a, ItemIndex, T>, + inner: item_set::Values<'a, T>, } impl<'a, T: IdHashItem> Iter<'a, T> { @@ -46,7 +44,6 @@ impl ExactSizeIterator for Iter<'_, T> { } } -// hash_map::Iter is a FusedIterator, so Iter is as well. impl FusedIterator for Iter<'_, T> {} /// An iterator over the elements of a [`IdHashMap`] by mutable reference. @@ -68,7 +65,7 @@ pub struct IterMut< A: Allocator = Global, > { tables: &'a IdHashMapTables, - inner: hash_map::ValuesMut<'a, ItemIndex, T>, + inner: item_set::ValuesMut<'a, T>, } impl<'a, T: IdHashItem, S: BuildHasher, A: Allocator> IterMut<'a, T, S, A> { @@ -102,7 +99,6 @@ impl ExactSizeIterator } } -// hash_map::IterMut is a FusedIterator, so IterMut is as well. impl FusedIterator for IterMut<'_, T, S, A> { @@ -119,7 +115,7 @@ impl FusedIterator /// [`HashMap`]: std::collections::HashMap #[derive(Debug)] pub struct IntoIter { - inner: hash_map::IntoValues>, + inner: item_set::IntoValues, } impl IntoIter { @@ -144,5 +140,4 @@ impl ExactSizeIterator for IntoIter { } } -// hash_map::IterMut is a FusedIterator, so IterMut is as well. impl FusedIterator for IntoIter {} diff --git a/crates/iddqd/src/id_ord_map/imp.rs b/crates/iddqd/src/id_ord_map/imp.rs index e13a7d7c..ed726ec4 100644 --- a/crates/iddqd/src/id_ord_map/imp.rs +++ b/crates/iddqd/src/id_ord_map/imp.rs @@ -432,7 +432,10 @@ impl IdOrdMap { /// assert!(map.capacity() >= 2); /// ``` pub fn shrink_to_fit(&mut self) { - self.items.shrink_to_fit(); + let remap = self.items.shrink_to_fit(); + if !remap.is_identity() { + self.tables.key_to_item.remap_indexes(&remap); + } } /// Shrinks the capacity of the map with a lower limit. It will drop @@ -475,7 +478,10 @@ impl IdOrdMap { /// assert!(map.capacity() >= 2); /// ``` pub fn shrink_to(&mut self, min_capacity: usize) { - self.items.shrink_to(min_capacity); + let remap = self.items.shrink_to(min_capacity); + if !remap.is_identity() { + self.tables.key_to_item.remap_indexes(&remap); + } } /// Iterates over the items in the map. @@ -592,7 +598,7 @@ impl IdOrdMap { // Check that the indexes are all correct. - for (&ix, item) in self.items.iter() { + for (ix, item) in self.items.iter() { let key = item.key(); let ix1 = match chaos { ValidateChaos::Yes => { @@ -1363,7 +1369,7 @@ impl IdOrdMap { Q: ?Sized + Ord + Equivalent>, { self.items.iter().find_map(|(index, item)| { - (k.equivalent(&item.key())).then_some(*index) + (k.equivalent(&item.key())).then_some(index) }) } diff --git a/crates/iddqd/src/id_ord_map/iter.rs b/crates/iddqd/src/id_ord_map/iter.rs index 3524e749..77aa6932 100644 --- a/crates/iddqd/src/id_ord_map/iter.rs +++ b/crates/iddqd/src/id_ord_map/iter.rs @@ -1,8 +1,11 @@ use super::{IdOrdItem, RefMut, tables::IdOrdMapTables}; use crate::support::{ - alloc::Global, borrow::DormantMutRef, btree_table, item_set::ItemSet, + alloc::Global, + borrow::DormantMutRef, + btree_table, + item_set::{ConsumingItemSet, ItemSet, ItemSlot}, }; -use core::{hash::Hash, iter::FusedIterator}; +use core::{hash::Hash, iter::FusedIterator, marker::PhantomData}; /// An iterator over the elements of an [`IdOrdMap`] by shared reference. /// @@ -45,6 +48,40 @@ impl ExactSizeIterator for Iter<'_, T> { // btree_set::Iter is a FusedIterator, so Iter is as well. impl FusedIterator for Iter<'_, T> {} +/// A raw pointer into the start of an `ItemSet`'s slot buffer, with the same +/// thread-safety properties as an `&'a mut ItemSet`. +/// +/// We use a raw pointer rather than lifetime extension as done by the hash map +/// iterators to avoid reborrow invalidation under Stacked Borrows. Due to the +/// way `Vec::index_mut` works, each iteration reborrowing `&mut self.items` +/// would invalidate previously yielded `&mut T` children. +struct ItemSetPtr<'a, T: IdOrdItem> { + // The pointer to the start of the slot buffer. + start_ptr: *mut ItemSlot, + // Number of slots in the backing buffer at construction time. + slot_count: usize, + // Borrow the ItemSet for `'a` so the raw pointer stays live, and so that + // variance and drop-check work the same as `&'a mut ItemSet`. + _marker: PhantomData<&'a mut ItemSet>, +} + +impl core::fmt::Debug for ItemSetPtr<'_, T> { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + f.debug_struct("ItemSetPtr") + .field("start_ptr", &self.start_ptr) + .field("slot_count", &self.slot_count) + .finish() + } +} + +// SAFETY: `ItemSetPtr<'a, T>` has the same thread-safety semantics as `&'a mut +// ItemSet`, which is `Send`/`Sync` iff `ItemSet` is +// either of those, respectively. This reduces to `T: Send` / `T: Sync`, since +// the global allocator `Global` is always `Send` + `Sync`. +unsafe impl<'a, T: IdOrdItem + Send> Send for ItemSetPtr<'a, T> {} +// SAFETY: see the `Send` impl above. +unsafe impl<'a, T: IdOrdItem + Sync> Sync for ItemSetPtr<'a, T> {} + /// An iterator over the elements of a [`IdOrdMap`] by mutable reference. /// /// This iterator returns [`RefMut`] instances. @@ -58,7 +95,7 @@ pub struct IterMut<'a, T: IdOrdItem> where T::Key<'a>: Hash, { - items: &'a mut ItemSet, + items: ItemSetPtr<'a, T>, tables: &'a IdOrdMapTables, iter: btree_table::Iter<'a>, } @@ -71,7 +108,13 @@ where items: &'a mut ItemSet, tables: &'a IdOrdMapTables, ) -> Self { - Self { items, tables, iter: tables.key_to_item.iter() } + let slot_count = items.slot_count(); + let start_ptr = items.start_ptr(); + Self { + items: ItemSetPtr { start_ptr, slot_count, _marker: PhantomData }, + tables, + iter: tables.key_to_item.iter(), + } } } @@ -84,35 +127,43 @@ where #[inline] fn next(&mut self) -> Option { let index = self.iter.next()?; - - let item = &mut self.items[index]; - - // SAFETY: This lifetime extension from self to 'a is safe based on two - // things: + let raw_index = index.as_u32() as usize; + + // This is a belt-and-suspenders bounds check. As of 2026-04-28, we've + // carefully analyzed all the code paths (including for panic safety) to + // ensure that indexes stored in the B-tree are always in bounds. But a + // future change might inadvertently break things. Handle this kind of + // programmer error as a panic rather than UB. + assert!( + raw_index < self.items.slot_count, + "btree index {raw_index} out of bounds for slot count {}", + self.items.slot_count, + ); + + // SAFETY: We need to show: // - // 1. We never repeat indexes, i.e. for an index i, once we've handed - // out an item at i, creating `&mut T`, we'll never get the index i - // again. (This is guaranteed from the set-based nature of the - // iterator.) This means that we don't ever create a mutable alias to - // the same memory. + // * `self.items.ptr.add(raw_index)` points at valid memory. + // * There are no overlapping mutable borrows of the same memory. // - // In particular, unlike all the other places we look up data from a - // btree table, we don't pass a lookup function into - // self.iter.next(). If we did, then it is possible the lookup - // function would have been called with an old index i. But we don't - // need to do that. + // This is shown by the following observations: // - // 2. All mutable references to data within self.items are derived from - // self.items. So, the rule described at [1] is upheld: - // - // > When creating a mutable reference, then while this reference - // > exists, the memory it points to must not get accessed (read or - // > written) through any other pointer or reference not derived from - // > this reference. - // - // [1]: - // https://doc.rust-lang.org/std/ptr/index.html#pointer-to-reference-conversion - let item = unsafe { core::mem::transmute::<&mut T, &'a mut T>(item) }; + // * We construct `ItemSetPtr` by mutably borrowing the item set, + // which means that while this iterator is alive, no other code + // can access the item set. + // * The bounds check above shows that `raw_index` is in bounds. + // * The B-tree only stores indexes that currently point at `Some` + // slots in the backing `ItemSet`, so the slot is initialized. + // (Again, as of 2026-04-28 we've verified this invariant, but + // a future change might break things, so we use `expect` and not + // `unwrap_unchecked`.) + // * The B-tree is a set, so each call to `self.iter.next()` yields a + // distinct `index`. This means that the handed-out `&mut T`s + // never point to the same memory. + let item: &'a mut T = unsafe { + (*self.items.start_ptr.add(raw_index)) + .as_mut() + .expect("btree index points at an Occupied slot in ItemSet") + }; let (hash, dormant) = { let (item, dormant) = DormantMutRef::new(item); @@ -120,8 +171,8 @@ where (hash, dormant) }; - // SAFETY: item is dropped above, and self is no longer used after this - // point. + // SAFETY: item is dropped above, and self is no longer used + // after this point. let item = unsafe { dormant.awaken() }; Some(RefMut::new(self.tables.state().clone(), hash, item)) @@ -138,7 +189,6 @@ where } } -// hash_map::IterMut is a FusedIterator, so IterMut is as well. impl<'a, T: IdOrdItem + 'a> FusedIterator for IterMut<'a, T> where T::Key<'a>: Hash { @@ -152,7 +202,7 @@ impl<'a, T: IdOrdItem + 'a> FusedIterator for IterMut<'a, T> where /// [`IdOrdMap::into_iter`]: crate::IdOrdMap::into_iter #[derive(Debug)] pub struct IntoIter { - items: ItemSet, + items: ConsumingItemSet, iter: btree_table::IntoIter, } @@ -161,7 +211,10 @@ impl IntoIter { items: ItemSet, tables: IdOrdMapTables, ) -> Self { - Self { items, iter: tables.key_to_item.into_iter() } + Self { + items: items.into_consuming(), + iter: tables.key_to_item.into_iter(), + } } } @@ -171,9 +224,11 @@ impl Iterator for IntoIter { #[inline] fn next(&mut self) -> Option { let index = self.iter.next()?; + // We own `self.items` and the B-tree's indexes are never revisited, so + // we can take directly from the consuming view. let next = self .items - .remove(index) + .take(index) .unwrap_or_else(|| panic!("index {index} not found in items")); Some(next) } diff --git a/crates/iddqd/src/support/btree_table.rs b/crates/iddqd/src/support/btree_table.rs index 9ef8eb94..30b4c147 100644 --- a/crates/iddqd/src/support/btree_table.rs +++ b/crates/iddqd/src/support/btree_table.rs @@ -4,7 +4,7 @@ //! integers (that are indexes corresponding to items), but use an external //! comparator. -use super::{ItemIndex, map_hash::MapHash}; +use super::{ItemIndex, item_set::IndexRemap, map_hash::MapHash}; use crate::internal::{TableValidationError, ValidateCompact}; use alloc::{ collections::{BTreeSet, btree_set}, @@ -34,6 +34,10 @@ thread_local! { /// * When the CmpDropGuard is dropped (including due to a panic), we reset /// the comparator to None. /// + /// Comparators take `&Index` rather than `Index` by value because `Index` + /// wraps `IndexCell` (an `AtomicU32` newtype) for in-place mutation in + /// `remap_indexes`, and `AtomicU32` isn't `Copy`. + /// /// This is not great! (For one, thread-locals and no-std don't really mix.) /// Some alternatives: /// @@ -61,10 +65,13 @@ thread_local! { /// default choice to balance cache locality, but other options are worth /// benchmarking. We do need to provide a comparator, though, so radix /// trees and such are out of the question. - static CMP: Cell Ordering>> + static CMP: Cell>> = const { Cell::new(None) }; } +/// External comparator type used via `CMP`'s dynamic scoping. +type IndexCmp<'a> = dyn Fn(&Index, &Index) -> Ordering + 'a; + /// A B-tree-based table with an external comparator. #[derive(Clone, Debug, Default)] pub(crate) struct MapBTreeTable { @@ -108,18 +115,16 @@ impl MapBTreeTable { // All items between 0 (inclusive) and self.len() (exclusive) // are present, and there are no duplicates. Also, the sentinel // value should not be stored. - let mut indexes: Vec<_> = Vec::with_capacity(expected_len); + let mut indexes: Vec = + Vec::with_capacity(expected_len); for index in &self.items { - match index.0 { - Index::SENTINEL_VALUE => { - return Err(TableValidationError::new( - "sentinel value should not be stored in map", - )); - } - v => { - indexes.push(v); - } + let v = index.value(); + if v == Index::SENTINEL_VALUE { + return Err(TableValidationError::new( + "sentinel value should not be stored in map", + )); } + indexes.push(v); } indexes.sort_unstable(); for (i, index) in indexes.iter().enumerate() { @@ -133,9 +138,10 @@ impl MapBTreeTable { ValidateCompact::NonCompact => { // There should be no duplicates, and the sentinel value // should not be stored. - let indexes: Vec<_> = self.items.iter().copied().collect(); + let indexes: Vec = + self.items.iter().map(|ix| ix.value()).collect(); let index_set: BTreeSet = - indexes.iter().map(|ix| ix.0).collect(); + indexes.iter().copied().collect(); if index_set.len() != indexes.len() { return Err(TableValidationError::new(format!( "expected no duplicates, but found {} duplicates \ @@ -157,12 +163,12 @@ impl MapBTreeTable { #[inline] pub(crate) fn first(&self) -> Option { - self.items.first().map(|ix| ix.0) + self.items.first().map(|ix| ix.value()) } #[inline] pub(crate) fn last(&self) -> Option { - self.items.last().map(|ix| ix.0) + self.items.last().map(|ix| ix.value()) } pub(crate) fn find_index( @@ -179,11 +185,11 @@ impl MapBTreeTable { let guard = CmpDropGuard::new(&f); - let ret = match self.items.get(&Index::SENTINEL) { - Some(Index(v)) if *v == Index::SENTINEL_VALUE => { + let ret = match self.items.get(&Index::sentinel()) { + Some(ix) if ix.value() == Index::SENTINEL_VALUE => { panic!("internal map shouldn't store sentinel value") } - Some(Index(v)) => Some(*v), + Some(ix) => Some(ix.value()), None => { // The key is not in the table. None @@ -234,7 +240,32 @@ impl MapBTreeTable { { // We don't need to set up a comparator in the environment because // `retain` doesn't do any comparisons as part of its operation. - self.items.retain(|index| f(index.0)); + self.items.retain(|index| f(index.value())); + } + + /// 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. + /// + /// We do not rebuild the tree. [`IndexRemap`] preserves relative + /// order, so the tree's iteration order — which is the user's + /// `Ord` over items — matches before and after the rewrite. Only + /// the stored index values change; node structure, pointers, and + /// the user-visible total order are all preserved. The walk is + /// O(N) with no comparator calls and no allocations. + /// + /// 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 + pub(crate) fn remap_indexes(&mut self, remap: &IndexRemap) { + for idx in self.items.iter() { + let new = remap.remap(idx.value()); + idx.set_value(new); + } } /// Clears the B-tree table, removing all items. @@ -279,7 +310,7 @@ impl<'a> Iterator for Iter<'a> { type Item = ItemIndex; fn next(&mut self) -> Option { - self.inner.next().map(|index| index.0) + self.inner.next().map(|index| index.value()) } } @@ -298,21 +329,22 @@ impl Iterator for IntoIter { type Item = ItemIndex; fn next(&mut self) -> Option { - self.inner.next().map(|index| index.0) + self.inner.next().map(|index| index.value()) } } fn find_cmp<'a, K, Q, F>( key: &'a Q, lookup: F, -) -> impl Fn(Index, Index) -> Ordering + 'a +) -> impl Fn(&Index, &Index) -> Ordering + 'a where Q: ?Sized + Comparable, F: 'a + Fn(ItemIndex) -> K, K: Ord, { - move |a: Index, b: Index| { - if a.0 == b.0 { + move |a: &Index, b: &Index| { + let (a, b) = (a.value(), b.value()); + if a == b { // This is potentially load-bearing! It means that even if the Eq // implementation on map items is wrong, we treat items at the same // index as equal. @@ -321,7 +353,7 @@ where // multiple mutable references to the same index. return Ordering::Equal; } - match (a.0, b.0) { + match (a, b) { (Index::SENTINEL_VALUE, v) => key.compare(&lookup(v)), (v, Index::SENTINEL_VALUE) => key.compare(&lookup(v)).reverse(), (a, b) => lookup(a).cmp(&lookup(b)), @@ -333,14 +365,15 @@ fn insert_cmp<'a, K, Q, F>( index: ItemIndex, key: &'a Q, lookup: F, -) -> impl Fn(Index, Index) -> Ordering + 'a +) -> impl Fn(&Index, &Index) -> Ordering + 'a where Q: ?Sized + Comparable, F: 'a + Fn(ItemIndex) -> K, K: Ord, { - move |a: Index, b: Index| { - if a.0 == b.0 { + move |a: &Index, b: &Index| { + let (a, b) = (a.value(), b.value()); + if a == b { // This is potentially load-bearing! It means that even if the Eq // implementation on map items is wrong, we treat items at the same // index as equal. @@ -349,7 +382,7 @@ where // multiple mutable references to the same index. return Ordering::Equal; } - match (a.0, b.0) { + match (a, b) { // The sentinel value should not be invoked at all, because it's not // passed in during insert and not stored in the table. (Index::SENTINEL_VALUE, _) | (_, Index::SENTINEL_VALUE) => { @@ -367,7 +400,7 @@ struct CmpDropGuard<'a> { } impl<'a> CmpDropGuard<'a> { - fn new(f: &'a dyn Fn(Index, Index) -> Ordering) -> Self { + fn new(f: &'a IndexCmp<'a>) -> Self { // CMP lasts only as long as this function and is immediately reset to // None once this scope is left. let ret = Self { _marker: PhantomData }; @@ -375,10 +408,9 @@ impl<'a> CmpDropGuard<'a> { // SAFETY: This is safe because we are not storing the reference // anywhere, and it is only used for the lifetime of this CmpDropGuard. let as_static = unsafe { - std::mem::transmute::< - &'a dyn Fn(Index, Index) -> Ordering, - &'static dyn Fn(Index, Index) -> Ordering, - >(f) + std::mem::transmute::<&'a IndexCmp<'a>, &'static IndexCmp<'static>>( + f, + ) }; CMP.set(Some(as_static)); @@ -392,19 +424,92 @@ impl Drop for CmpDropGuard<'_> { } } -#[derive(Clone, Copy, Debug)] -struct Index(ItemIndex); +/// An [`ItemIndex`] (= `u32`) with interior mutability, layout-identical +/// to `u32`. +/// +/// Backed by `AtomicU32`. We use `Relaxed` ordering everywhere because +/// the only caller of `set` holds `&mut MapBTreeTable`, which excludes +/// every other reference — so there is never a race between a reader +/// and a writer. `Relaxed` loads compile to a plain `mov` on x86-64 +/// and similar instructions on other architectures, so this gives us +/// interior mutability at the cost of a normal load. +/// +/// Going through `AtomicU32` rather than `Cell` keeps us +/// naturally `Sync` without an `unsafe impl Sync` — `AtomicU32` is +/// designed to be accessed from multiple threads. +#[repr(transparent)] +#[derive(Debug, Default)] +struct IndexCell(core::sync::atomic::AtomicU32); + +impl Clone for IndexCell { + fn clone(&self) -> Self { + Self(core::sync::atomic::AtomicU32::new(self.get().as_u32())) + } +} + +impl IndexCell { + #[inline] + const fn new(value: ItemIndex) -> Self { + Self(core::sync::atomic::AtomicU32::new(value.as_u32())) + } + + #[inline] + fn get(&self) -> ItemIndex { + ItemIndex::new(self.0.load(core::sync::atomic::Ordering::Relaxed)) + } + + /// Overwrite the stored value. The atomic store makes this safe to + /// call through `&self`, though in practice callers only invoke it + /// while holding `&mut` on the enclosing `MapBTreeTable`. + #[inline] + fn set(&self, value: ItemIndex) { + debug_assert_ne!( + value, + ItemIndex::SENTINEL, + "IndexCell::set: sentinel must never be stored in the table", + ); + self.0.store(value.as_u32(), core::sync::atomic::Ordering::Relaxed); + } +} + +#[derive(Clone, Debug)] +struct Index(IndexCell); impl Index { const SENTINEL_VALUE: ItemIndex = ItemIndex::SENTINEL; - const SENTINEL: Self = Self(Self::SENTINEL_VALUE); + + /// Returns a fresh sentinel `Index`. + /// + /// A function rather than an associated `const` because `IndexCell` + /// wraps an `AtomicU32` (interior mutability), and a `const Self` + /// would trigger `clippy::declare_interior_mutable_const` at every + /// borrow site. + #[inline] + const fn sentinel() -> Self { + Self(IndexCell::new(Self::SENTINEL_VALUE)) + } #[inline] fn new(value: ItemIndex) -> Self { if value == Self::SENTINEL_VALUE { panic!("btree map overflow, index with value {value:?} was added") } - Self(value) + Self(IndexCell::new(value)) + } + + #[inline] + fn value(&self) -> ItemIndex { + self.0.get() + } + + /// Overwrite the stored index value in place. + /// + /// Safe thanks to the atomic store inside [`IndexCell::set`]. In + /// practice we only call this from `remap_indexes`, which holds + /// `&mut MapBTreeTable`. + #[inline] + fn set_value(&self, value: ItemIndex) { + self.0.set(value) } } @@ -413,15 +518,16 @@ impl PartialEq for Index { // For non-sentinel indexes, two values are the same iff their indexes // are the same. This is ensured by the fact that our key types // implement Eq (as part of implementing Ord). - if self.0 != Self::SENTINEL_VALUE && other.0 != Self::SENTINEL_VALUE { - return self.0 == other.0; + let (a, b) = (self.value(), other.value()); + if a != Self::SENTINEL_VALUE && b != Self::SENTINEL_VALUE { + return a == b; } // If any of the two indexes is the sentinel, we're required to perform // a lookup. CMP.with(|cmp| { let cmp = cmp.get().expect("cmp should be set"); - cmp(*self, *other) == Ordering::Equal + cmp(self, other) == Ordering::Equal }) } } @@ -435,7 +541,7 @@ impl Ord for Index { // which should have set the thread local. CMP.with(|cmp| { let cmp = cmp.get().expect("cmp should be set"); - cmp(*self, *other) + cmp(self, other) }) } } @@ -446,3 +552,95 @@ impl PartialOrd for Index { Some(self.cmp(other)) } } + +#[cfg(all(test, feature = "std"))] +mod tests { + use super::*; + use crate::support::{alloc::Global, item_set::ItemSet}; + use core::cell::Cell; + + thread_local! { + /// When set, `PanickingKey::cmp` panics on invocation. Scoped by + /// individual tests. + static PANIC_TRIGGER: Cell = const { Cell::new(false) }; + } + + /// A key type whose `Ord` impl can be made to panic on demand. + #[derive(Clone, Debug, PartialEq, Eq)] + struct PanickingKey(u32); + + impl PartialOrd for PanickingKey { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } + } + + impl Ord for PanickingKey { + fn cmp(&self, other: &Self) -> Ordering { + if PANIC_TRIGGER.with(|c| c.get()) { + panic!("simulated Ord panic"); + } + self.0.cmp(&other.0) + } + } + + /// `remap_indexes` must not invoke the user-supplied `Ord` impl. We arm + /// `PANIC_TRIGGER` for the duration of the call and verify the rebuild + /// succeeds — any stray user-Ord invocation would panic the test. + #[test] + fn remap_indexes_does_not_call_user_ord() { + // Build an IndexRemap with holes. This yields holes = [1, 3]. + let mut set: ItemSet = ItemSet::new(); + for i in 0..5u32 { + set.assert_can_grow().insert(PanickingKey(i * 10)); + } + set.remove(ItemIndex::new(1)); + set.remove(ItemIndex::new(3)); + let remap = set.shrink_to_fit(); + assert!(!remap.is_identity(), "remap should carry two holes"); + + // A MapBTreeTable populated to match the pre-compaction live indexes + // 0, 2, 4 — these are the indexes the outer map would have stored + // before shrink. Setup uses the user `Ord`, so the trigger is off. + let mut table = MapBTreeTable::new(); + let pre_lookup = |ix: ItemIndex| -> PanickingKey { + match ix.as_u32() { + 0 => PanickingKey(0), + 2 => PanickingKey(20), + 4 => PanickingKey(40), + _ => panic!("unexpected index in pre-compaction lookup: {ix}"), + } + }; + for ix in [0u32, 2, 4] { + let ix = ItemIndex::new(ix); + let key = pre_lookup(ix); + table.insert(ix, &key, pre_lookup); + } + assert_eq!(table.len(), 3); + assert_eq!( + table + .items + .iter() + .map(|i| i.value().as_u32()) + .collect::>(), + [0u32, 2, 4], + ); + + // Arm the trigger: any call into `PanickingKey::cmp` during the + // rebuild below will panic this test. + PANIC_TRIGGER.with(|c| c.set(true)); + table.remap_indexes(&remap); + PANIC_TRIGGER.with(|c| c.set(false)); + + // Remap 0 -> 0, 2 -> 1, 4 -> 2, and key order is preserved, so the + // final contents must be [0, 1, 2]. + assert_eq!( + table + .items + .iter() + .map(|i| i.value().as_u32()) + .collect::>(), + [0u32, 1, 2], + ); + } +} diff --git a/crates/iddqd/src/support/hash_table.rs b/crates/iddqd/src/support/hash_table.rs index 9d5a3ad5..ee2822b9 100644 --- a/crates/iddqd/src/support/hash_table.rs +++ b/crates/iddqd/src/support/hash_table.rs @@ -3,6 +3,7 @@ use super::{ ItemIndex, alloc::{AllocWrapper, Allocator}, + item_set::IndexRemap, map_hash::MapHash, }; use crate::internal::{TableValidationError, ValidateCompact}; @@ -189,6 +190,21 @@ impl MapHashTable { self.items.reserve(additional, hasher); } + /// 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. + /// + /// [`ItemSet::shrink_to_fit`]: super::item_set::ItemSet::shrink_to_fit + /// [`ItemSet::shrink_to`]: super::item_set::ItemSet::shrink_to + pub(crate) fn remap_indexes(&mut self, remap: &IndexRemap) { + for slot in self.items.iter_mut() { + *slot = remap.remap(*slot); + } + } + /// Shrinks the capacity of the hash table as much as possible. /// /// See [`Self::reserve`] for the contract `hasher` must satisfy. diff --git a/crates/iddqd/src/support/item_index.rs b/crates/iddqd/src/support/item_index.rs index 8e507306..e6d0e214 100644 --- a/crates/iddqd/src/support/item_index.rs +++ b/crates/iddqd/src/support/item_index.rs @@ -7,30 +7,23 @@ use core::fmt; /// /// We use a `u32` and not a `usize` for these indexes because the increased /// density leads to meaningful performance improvements on 64-bit targets. This -/// does mean that the maximum number of items is limited to 2^32 - 1 (we -/// reserve `u32::MAX` for the sentinel), but we consider that to be a -/// reasonable tradeoff. The limit is enforced within -/// [`ItemSet::assert_can_grow`]; see [`Self::MAX_VALID`]. +/// does mean that the maximum number of concurrently live items is limited to +/// `u32::MAX - 1` slots (the -1 is because `u32::MAX` is reserved for +/// [`Self::SENTINEL`]). This limit is enforced within +/// [`ItemSet::assert_can_grow`]. /// /// [`ItemSet::assert_can_grow`]: super::item_set::ItemSet::assert_can_grow #[derive(Clone, Copy, Default, PartialEq, Eq, PartialOrd, Ord, Hash)] pub(crate) struct ItemIndex(u32); impl ItemIndex { - /// The smallest possible index. - pub(crate) const ZERO: Self = Self(0); - /// The largest index that may be assigned to an item. /// /// One below `u32::MAX`, which is reserved as [`Self::SENTINEL`]. - /// Equivalently, the maximum number of items that may ever be inserted - /// into a single map (across the map's lifetime, since indexes are never - /// reused other than through `last_index`) is `u32::MAX`. pub(crate) const MAX_VALID: Self = Self(u32::MAX - 1); /// Reserved sentinel value marking the root/empty slot. Never assigned to /// an item. - #[cfg_attr(not(feature = "std"), expect(dead_code))] pub(crate) const SENTINEL: Self = Self(u32::MAX); /// Wraps a raw `u32`. @@ -44,24 +37,6 @@ impl ItemIndex { pub(crate) const fn as_u32(self) -> u32 { self.0 } - - /// Returns this index plus one, panicking on overflow. - /// - /// Used by [`ItemSet`](super::item_set::ItemSet) to advance - /// `next_index` after an insert. - #[inline] - pub(crate) fn next(self) -> Self { - Self(self.0.checked_add(1).expect("ItemIndex did not overflow")) - } - - /// Returns this index minus one, panicking on underflow. - /// - /// Used by [`ItemSet`](super::item_set::ItemSet) to roll back - /// `next_index` when removing the highest-index item. - #[inline] - pub(crate) fn prev(self) -> Self { - Self(self.0.checked_sub(1).expect("ItemIndex did not underflow")) - } } impl fmt::Debug for ItemIndex { diff --git a/crates/iddqd/src/support/item_set.rs b/crates/iddqd/src/support/item_set.rs index eda24794..dc0d393e 100644 --- a/crates/iddqd/src/support/item_set.rs +++ b/crates/iddqd/src/support/item_set.rs @@ -1,41 +1,115 @@ -use super::{ItemIndex, alloc::AllocWrapper}; +//! A dense, index-keyed container for items. +//! +//! # Design +//! +//! Each slot is an `ItemSlot` that is either `Occupied(T)` or `Vacant { +//! next }`. The free chain consists of vacant slots that are linked together +//! via `next` pointers, with `free_head` as its LIFO top and +//! [`ItemIndex::SENTINEL`] as the end-of-list sentinel. +//! +//! Removed slots are recycled by the next [`GrowHandle::insert`], so a churn +//! workload stabilizes at the high-water mark of simultaneously-live items. +//! +//! The container maintains a single allocation (`items`) and uses two `u32`s +//! of stack footprint beyond it (the `free_head` and the current `len`). +//! +//! # Why slot-based +//! +//! We also tried a `Vec>` plus a separately allocated free list for +//! vacant indexes. That was optimal storage for any `T` with a niche +//! (`size_of::>() == size_of::()`). But this came at the cost of a +//! hand-rolled unsafe allocator to manage the secondary allocation (i.e., north +//! of 350 lines of layout-math, lifetime, and `Send`/`Sync` reasoning). The +//! slot-based layout eliminates that module entirely: the only unsafe in this +//! file is the disjoint-indexes trick in `get_disjoint_mut`, which any +//! slot-based container needs regardless of backend. +//! +//! The tradeoff is that `ItemSlot` carries a discriminant, so slots are at +//! least `max(size_of::(), size_of::()) + align_of::>()`. +//! For types with a niche (including structs where a field has a niche), this +//! is one word larger per slot than `Option` would be. Benchmarking +//! indicates that overall this is a wash. Based on that, we choose the +//! implementation with less unsafe code. +//! +//! # Invariants +//! +//! 1. For every `i < items.len()`: `items[i]` is `Occupied` iff `i` is +//! not reachable from `free_head` via the `Vacant::next` chain. +//! 2. The `Vacant::next` chain starting at `free_head` terminates at +//! `SENTINEL`, visits every vacant slot exactly once, and +//! contains no in-bounds index that refers to an `Occupied` slot. +//! 3. `len == items.iter().filter(|e| matches!(e, Occupied(_))).count()`. +//! +//! Under these invariants: +//! +//! * The number of occupied slots is `self.len`. +//! * The number of vacant slots is `items.len() - self.len`. + +use super::{ + ItemIndex, + alloc::{AllocWrapper, Allocator, Global, global_alloc}, +}; use crate::{ + errors::TryReserveError, internal::{ValidateCompact, ValidationError}, - support::alloc::{Allocator, Global, global_alloc}, }; +use allocator_api2::vec::Vec; use core::{ fmt, + iter::FusedIterator, ops::{Index, IndexMut}, }; -use hashbrown::{HashMap, hash_map}; -use rustc_hash::FxBuildHasher; -/// A map of items stored by integer index. -#[derive(Clone)] -pub(crate) struct ItemSet { - // rustc-hash's FxHashMap is custom-designed for compact-ish integer keys. - items: HashMap>, - // The next index to use. This only ever goes up, not down (modulo the - // small `remove`-of-tail optimization below). - // - // An alternative might be to use a free list of indexes, but that's - // unnecessarily complex. - next_index: ItemIndex, +/// 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. +/// +/// Two cases: +/// +/// - [`IndexRemap::Identity`]: compaction was a no-op (no holes were +/// filled), so every old index is still valid as-is. +/// - [`IndexRemap::Permuted`]: holes were filled. The contained +/// `Vec` is a direct position array — `new_pos[old]` is +/// the new index, or [`ItemIndex::SENTINEL`] for slots that were vacated. +#[derive(Clone, Debug)] +pub(crate) enum IndexRemap { + /// Compaction was a no-op: every old slot index is still valid. + Identity, + /// Slots moved during compaction. `new_pos[old]` is either the new index + /// for the slot that used to live at `old`, or [`ItemIndex::SENTINEL`] if + /// `old` was vacant at compaction time. + Permuted(alloc::vec::Vec), } -impl fmt::Debug for ItemSet { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_struct("ItemSet") - .field("items", &self.items) - .field("next_index", &self.next_index) - .finish() +impl IndexRemap { + #[inline] + pub(crate) fn is_identity(&self) -> bool { + matches!(self, Self::Identity) } -} -impl ItemSet { + /// Looks up the post-compaction index for `old`. + /// + /// 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. #[inline] - pub(crate) const fn new() -> Self { - Self::new_in(global_alloc()) + pub(crate) fn remap(&self, old: ItemIndex) -> ItemIndex { + match self { + Self::Identity => old, + Self::Permuted(new_pos) => { + let new = new_pos[old.as_u32() as usize]; + if new == ItemIndex::SENTINEL { + panic!( + "IndexRemap::remap called on a compacted-away \ + index {old}" + ) + } + new + } + } } } @@ -54,12 +128,12 @@ pub(crate) struct GrowHandle<'a, T, A: Allocator> { items: &'a mut ItemSet, } -impl Index for GrowHandle<'_, T, A> { - type Output = T; +impl core::ops::Deref for GrowHandle<'_, T, A> { + type Target = ItemSet; #[inline] - fn index(&self, index: ItemIndex) -> &T { - &self.items[index] + fn deref(&self) -> &ItemSet { + self.items } } @@ -68,7 +142,14 @@ impl<'a, T, A: Allocator> GrowHandle<'a, T, A> { #[cfg_attr(not(feature = "std"), expect(dead_code))] #[inline] pub(crate) fn next_index(&self) -> ItemIndex { - self.items.next_index + if self.free_head == ItemIndex::SENTINEL { + // `assert_can_grow` enforces `items.len() <= ItemIndex::MAX_VALID`, + // so this conversion is lossless. + ItemIndex::new(self.items.len() as u32) + } else { + // Use the LIFO slot. + self.free_head + } } /// Inserts `value` at [`Self::next_index`] and returns the chosen @@ -77,13 +158,120 @@ impl<'a, T, A: Allocator> GrowHandle<'a, T, A> { /// This is the only way to grow an [`ItemSet`]. #[inline] pub(crate) fn insert(self, value: T) -> ItemIndex { - let index = self.items.next_index; - self.items.items.insert(index, value); - // `assert_can_grow` guarantees `next_index <= ItemIndex::MAX_VALID` - // here, so the post-insert bump is at most `ItemIndex::SENTINEL` — - // no overflow. - self.items.next_index = self.items.next_index.next(); - index + if self.items.free_head == ItemIndex::SENTINEL { + // `assert_can_grow` guarantees `items.len() <= ItemIndex::MAX_VALID`, + // so this u32 conversion cannot lose precision. + let idx = ItemIndex::new(self.items.items.len() as u32); + self.items.items.push(ItemSlot::Occupied(value)); + self.items.len += 1; + idx + } else { + let idx = self.items.free_head; + // Replace the `Vacant { next }` at `idx` with `Occupied`, + // and advance `free_head` to `next`. + let slot = &mut self.items.items[idx.as_u32() as usize]; + let next = match slot { + ItemSlot::Occupied(_) => { + panic!("ItemSet free chain points at occupied slot {idx}") + } + ItemSlot::Vacant { next } => *next, + }; + *slot = ItemSlot::Occupied(value); + self.items.free_head = next; + self.items.len += 1; + idx + } + } +} + +/// A single slot in an [`ItemSet`]. +/// +/// Exposed at `pub(crate)` because `ItemSet::as_mut_ptr` hands out a +/// raw pointer into the `items` buffer; callers need to name the element +/// type. All other interaction with slots goes through `ItemSet`'s safe +/// methods. +#[derive(Clone, Debug)] +pub(crate) enum ItemSlot { + /// The slot holds a live value. + Occupied(T), + /// The slot is free. + /// + /// `next` is the index of the next free slot, or [`ItemIndex::SENTINEL`] if + /// this is the end of the chain. + Vacant { next: ItemIndex }, +} + +impl ItemSlot { + /// Returns a reference to the contained value, if occupied. + #[inline] + fn as_ref(&self) -> Option<&T> { + match self { + ItemSlot::Occupied(v) => Some(v), + ItemSlot::Vacant { .. } => None, + } + } + + /// Returns a mutable reference to the contained value, if occupied. + #[inline] + pub(crate) fn as_mut(&mut self) -> Option<&mut T> { + match self { + ItemSlot::Occupied(v) => Some(v), + ItemSlot::Vacant { .. } => None, + } + } + + #[inline] + fn is_occupied(&self) -> bool { + match self { + ItemSlot::Occupied(_) => true, + ItemSlot::Vacant { .. } => false, + } + } +} + +/// A dense, index-keyed container for items. +/// +/// See the [module-level docs](self) for the design and tradeoffs. +pub(crate) struct ItemSet { + items: Vec, AllocWrapper>, + /// LIFO head of the embedded free chain, or [`ItemIndex::SENTINEL`] when no + /// slots are free. + free_head: ItemIndex, + /// Count of `Occupied` slots, maintained by insert/remove. + /// + /// (ItemIndex is a u32, as is len, so the struct can be more tightly packed + /// than if both were usizes.) + len: u32, +} + +impl Clone for ItemSet { + fn clone(&self) -> Self { + Self { + items: self.items.clone(), + free_head: self.free_head, + len: self.len, + } + } +} + +impl fmt::Debug for ItemSet { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("ItemSet") + .field("len", &self.len) + .field("slots", &self.items.len()) + .field("free_head", &self.free_head) + .finish() + } +} + +impl ItemSet { + #[inline] + pub(crate) const fn new() -> Self { + Self { + items: Vec::new_in(AllocWrapper(global_alloc())), + free_head: ItemIndex::SENTINEL, + len: 0, + } } } @@ -91,19 +279,17 @@ impl ItemSet { #[inline] pub(crate) const fn new_in(alloc: A) -> Self { Self { - items: HashMap::with_hasher_in(FxBuildHasher, AllocWrapper(alloc)), - next_index: ItemIndex::ZERO, + items: Vec::new_in(AllocWrapper(alloc)), + free_head: ItemIndex::SENTINEL, + len: 0, } } pub(crate) fn with_capacity_in(capacity: usize, alloc: A) -> Self { Self { - items: HashMap::with_capacity_and_hasher_in( - capacity, - Default::default(), - AllocWrapper(alloc), - ), - next_index: ItemIndex::ZERO, + items: Vec::with_capacity_in(capacity, AllocWrapper(alloc)), + free_head: ItemIndex::SENTINEL, + len: 0, } } @@ -111,27 +297,96 @@ impl ItemSet { &self.items.allocator().0 } - /// Validates the item set. + /// Returns a raw pointer to the start of the backing slot buffer. + #[inline] + #[cfg_attr(not(feature = "std"), expect(dead_code))] + pub(crate) fn start_ptr(&mut self) -> *mut ItemSlot { + self.items.as_mut_ptr() + } + + /// Returns the number of slots in the backing buffer. + #[inline] + #[cfg_attr(not(feature = "std"), expect(dead_code))] + pub(crate) fn slot_count(&self) -> usize { + self.items.len() + } + pub(crate) fn validate( &self, compactness: ValidateCompact, ) -> Result<(), ValidationError> { - // If the map is expected to be compact, then ensure that all keys - // between 0 and next_index are present. - match compactness { - ValidateCompact::Compact => { - for i in 0..self.next_index.as_u32() { - let idx = ItemIndex::new(i); - if !self.items.contains_key(&idx) { + let occupied_count = + self.items.iter().filter(|e| e.is_occupied()).count(); + if occupied_count != self.len as usize { + return Err(ValidationError::General(format!( + "ItemSet len ({}) disagrees with occupied-slot count ({})", + self.len, occupied_count, + ))); + } + + // Walk the free chain and verify the following properties: + // + // * Every visited index is in bounds. + // * Every visited slot is `Vacant`. + // * We visit exactly `items.len() - len` slots (i.e. each + // vacant slot exactly once); this detects cycles and missing + // links at the same time. + let Some(expected_vacant) = + self.items.len().checked_sub(self.len as usize) + else { + return Err(ValidationError::General(format!( + "ItemSet len ({}) exceeds items.len() ({})", + self.len, + self.items.len(), + ))); + }; + + let mut walked = 0usize; + let mut cursor = self.free_head; + while cursor != ItemIndex::SENTINEL { + let cursor_idx = cursor.as_u32() as usize; + if cursor_idx >= self.items.len() { + return Err(ValidationError::General(format!( + "ItemSet free chain has out-of-range index {cursor}" + ))); + } + match &self.items[cursor_idx] { + ItemSlot::Occupied(_) => { + return Err(ValidationError::General(format!( + "ItemSet free chain points at occupied slot \ + {cursor}" + ))); + } + ItemSlot::Vacant { next } => { + walked += 1; + if walked > expected_vacant { return Err(ValidationError::General(format!( - "ItemSet is not compact: missing index {idx}" + "ItemSet free chain cycles or overshoots: \ + walked {walked} vacant slots, expected \ + {expected_vacant}" ))); } + cursor = *next; } } - ValidateCompact::NonCompact => { - // No real checks can be done in this case. + } + if walked != expected_vacant { + return Err(ValidationError::General(format!( + "ItemSet free chain length {walked} disagrees with \ + vacant-slot count {expected_vacant}" + ))); + } + + match compactness { + ValidateCompact::Compact => { + if expected_vacant != 0 { + return Err(ValidationError::General(format!( + "ItemSet is not compact: {expected_vacant} \ + vacant slots", + ))); + } } + ValidateCompact::NonCompact => {} } Ok(()) @@ -143,60 +398,93 @@ impl ItemSet { #[inline] pub(crate) fn is_empty(&self) -> bool { - self.items.is_empty() + self.len == 0 } #[inline] pub(crate) fn len(&self) -> usize { - self.items.len() + self.len as usize } #[inline] - pub(crate) fn iter(&self) -> hash_map::Iter<'_, ItemIndex, T> { - self.items.iter() + pub(crate) fn iter(&self) -> Iter<'_, T> { + Iter::new(self) } #[inline] #[expect(dead_code)] - pub(crate) fn iter_mut(&mut self) -> hash_map::IterMut<'_, ItemIndex, T> { - self.items.iter_mut() + pub(crate) fn iter_mut(&mut self) -> IterMut<'_, T> { + IterMut::new(self) } #[inline] - pub(crate) fn values(&self) -> hash_map::Values<'_, ItemIndex, T> { - self.items.values() + pub(crate) fn values(&self) -> Values<'_, T> { + Values::new(self) } #[inline] - pub(crate) fn values_mut( - &mut self, - ) -> hash_map::ValuesMut<'_, ItemIndex, T> { - self.items.values_mut() + pub(crate) fn values_mut(&mut self) -> ValuesMut<'_, T> { + ValuesMut::new(self) } #[inline] - pub(crate) fn into_values( - self, - ) -> hash_map::IntoValues> { - self.items.into_values() + pub(crate) fn into_values(self) -> IntoValues { + IntoValues::new(self) } #[inline] pub(crate) fn get(&self, index: ItemIndex) -> Option<&T> { - self.items.get(&index) + self.items.get(index.as_u32() as usize).and_then(ItemSlot::as_ref) } #[inline] pub(crate) fn get_mut(&mut self, index: ItemIndex) -> Option<&mut T> { - self.items.get_mut(&index) + self.items.get_mut(index.as_u32() as usize).and_then(ItemSlot::as_mut) } - #[inline] + /// Returns mutable references to up to `N` distinct indexes. + /// + /// Returns `None` for any index that is out of bounds, vacant, or + /// that duplicates an earlier index in the array. pub(crate) fn get_disjoint_mut( &mut self, indexes: [&ItemIndex; N], ) -> [Option<&mut T>; N] { - self.items.get_many_mut(indexes) + let len = self.items.len(); + let mut valid = [false; N]; + for i in 0..N { + let idx = indexes[i].as_u32() as usize; + if idx >= len { + continue; + } + // SAFETY: idx < len, so `items[idx]` is in bounds. + if !unsafe { self.items.get_unchecked(idx) }.is_occupied() { + continue; + } + let mut dup = false; + for j in 0..i { + if valid[j] && indexes[j].as_u32() == indexes[i].as_u32() { + dup = true; + break; + } + } + if !dup { + valid[i] = true; + } + } + + let base = self.items.as_mut_ptr(); + core::array::from_fn(|i| { + if valid[i] { + let idx = indexes[i].as_u32() as usize; + // SAFETY: we verified idx is in bounds, the slot is + // `Occupied`, and no earlier valid entry shares this + // index. Therefore the `&mut` references are disjoint. + unsafe { (*base.add(idx)).as_mut() } + } else { + None + } + }) } /// Returns a [`GrowHandle`] that grants exclusive access to grow the set by @@ -208,84 +496,157 @@ impl ItemSet { #[inline] #[must_use = "GrowHandle must be passed to GrowHandle::insert"] pub(crate) fn assert_can_grow(&mut self) -> GrowHandle<'_, T, A> { - // `MAX_VALID` is the largest assignable index. After that final insert, - // `next_index` is bumped to `SENTINEL` and the next `assert_can_grow` - // will panic. - assert!( - self.next_index <= ItemIndex::MAX_VALID, - "ItemSet index exceeds maximum index {}", - ItemIndex::MAX_VALID, - ); + if self.free_head == ItemIndex::SENTINEL { + assert!( + self.items.len() <= ItemIndex::MAX_VALID.as_u32() as usize, + "ItemSet index exceeds maximum index {}", + ItemIndex::MAX_VALID, + ); + } else { + // At least one vacant slot is available in self.items. + } GrowHandle { items: self } } + /// Removes the item at `index`, if any. + /// + /// This does not allocate: the freed index threads onto the embedded chain + /// in place. + /// + /// `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). #[inline] pub(crate) fn remove(&mut self, index: ItemIndex) -> Option { - let entry = self.items.remove(&index); - if entry.is_some() && index == self.next_index.prev() { - // If we removed the last entry, decrement next_index. Not strictly - // necessary but a nice optimization. - // - // This does not guarantee compactness, since it's possible for the - // following set of operations to occur: - // - // 0. start at next_index = 0 - // 1. insert 0, next_index = 1 - // 2. insert 1, next_index = 2 - // 3. remove 0, next_index = 2 - // 4. remove 1, next_index = 1 (not 0, even though the map is empty) - // - // Compactness would require a heap acting as a free list. But that - // seems generally unnecessary. - // - // `entry.is_some()` implies an item lived at `index`, so - // `next_index > index >= 0` and `prev()` cannot underflow. - self.next_index = self.next_index.prev(); + let slot = self.items.get_mut(index.as_u32() as usize)?; + if !slot.is_occupied() { + return None; } - entry + let ItemSlot::Occupied(v) = + core::mem::replace(slot, ItemSlot::Vacant { next: self.free_head }) + else { + unreachable!("is_occupied was just checked") + }; + self.free_head = index; + self.len = self.len.checked_sub(1).expect("ItemSet len should be > 0"); + Some(v) + } + + /// Consumes this set into an owned, invariant-free + /// [`ConsumingItemSet`]. + pub(crate) fn into_consuming(self) -> ConsumingItemSet { + ConsumingItemSet { items: self.items } } /// Clears the item set, removing all items. - #[inline] + /// + /// Preserves `items.capacity()`, matching the behavior of + /// [`Vec::clear`]. Any prior [`try_reserve`](Self::try_reserve) + /// reservation survives a `clear`. pub(crate) fn clear(&mut self) { self.items.clear(); - self.next_index = ItemIndex::ZERO; + self.free_head = ItemIndex::SENTINEL; + self.len = 0; } - // This method assumes that value has the same ID. It also asserts that - // `index` is valid (and panics if it isn't). + /// This method assumes that value has the same ID. It also asserts + /// that `index` is valid (and panics if it isn't). #[inline] pub(crate) fn replace(&mut self, index: ItemIndex, value: T) -> T { - self.items - .insert(index, value) - .unwrap_or_else(|| panic!("ItemSet index not found: {index}")) + let Some(slot) = self + .items + .get_mut(index.as_u32() as usize) + .filter(|s| s.is_occupied()) + else { + panic!("ItemSet index not found: {index}") + }; + let ItemSlot::Occupied(old) = + core::mem::replace(slot, ItemSlot::Occupied(value)) + else { + unreachable!("slot was just matched as Occupied") + }; + old } - /// Reserves capacity for at least `additional` more items. #[inline] pub(crate) fn reserve(&mut self, additional: usize) { self.items.reserve(additional); } - /// Shrinks the capacity of the item set as much as possible. #[inline] - pub(crate) fn shrink_to_fit(&mut self) { + pub(crate) fn shrink_to_fit(&mut self) -> IndexRemap { + let remap = self.compact(); self.items.shrink_to_fit(); + remap } - /// Shrinks the capacity of the item set with a lower limit. #[inline] - pub(crate) fn shrink_to(&mut self, min_capacity: usize) { + pub(crate) fn shrink_to(&mut self, min_capacity: usize) -> IndexRemap { + let remap = self.compact(); 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 { + let pre_len = self.items.len(); + if pre_len == self.len as usize { + // Already compact, so there's nothing to remap. + debug_assert_eq!( + self.free_head, + ItemIndex::SENTINEL, + "compact: items full but free_head not SENTINEL ({})", + self.free_head, + ); + return IndexRemap::Identity; + } + + // Do a forward scan, writing each `Occupied` into the next write + // position. As we go, build a `new_pos[old] = new` index so callers can + // rewrite their stored indexes. + assert!( + pre_len <= ItemIndex::MAX_VALID.as_u32() as usize, + "compact: items.len() {pre_len} exceeds MAX_VALID {}", + ItemIndex::MAX_VALID, + ); + let mut new_pos: alloc::vec::Vec = + alloc::vec::Vec::with_capacity(pre_len); + let mut write: u32 = 0; + for read in 0..pre_len { + match &self.items[read] { + ItemSlot::Occupied(_) => { + new_pos.push(ItemIndex::new(write)); + if write as usize != read { + self.items.swap(write as usize, read); + } + write += 1; + } + ItemSlot::Vacant { .. } => { + new_pos.push(ItemIndex::SENTINEL); + } + } + } + self.items.truncate(write as usize); + self.free_head = ItemIndex::SENTINEL; + // `len` is unchanged: we truncated only `Vacant` entries. + + IndexRemap::Permuted(new_pos) } /// Tries to reserve capacity for at least `additional` more items. + /// + /// After this call returns `Ok(())`, the next `additional` calls + /// to [`GrowHandle::insert`] are OOM-free. `remove` is always + /// OOM-free regardless. #[inline] pub(crate) fn try_reserve( &mut self, additional: usize, - ) -> Result<(), hashbrown::TryReserveError> { - self.items.try_reserve(additional) + ) -> Result<(), TryReserveError> { + self.items + .try_reserve(additional) + .map_err(TryReserveError::from_allocator_api2) } } @@ -300,9 +661,9 @@ mod serde_impls { where S: Serializer, { - // Serialize just the items -- don't serialize the map keys. We'll - // rebuild the map keys on deserialization. - serializer.collect_seq(self.items.values()) + // Serialize just the items -- don't serialize the map keys. + // We'll rebuild the map keys on deserialization. + serializer.collect_seq(self.values()) } } } @@ -312,8 +673,7 @@ impl Index for ItemSet { #[inline] fn index(&self, index: ItemIndex) -> &Self::Output { - self.items - .get(&index) + self.get(index) .unwrap_or_else(|| panic!("ItemSet index not found: {index}")) } } @@ -321,8 +681,424 @@ impl Index for ItemSet { impl IndexMut for ItemSet { #[inline] fn index_mut(&mut self, index: ItemIndex) -> &mut Self::Output { - self.items - .get_mut(&index) + self.get_mut(index) .unwrap_or_else(|| panic!("ItemSet index not found: {index}")) } } + +// --- Iterators ---------------------------------------------------------- + +/// An iterator over `(index, &item)` pairs in an [`ItemSet`]. +pub(crate) struct Iter<'a, T> { + inner: core::iter::Enumerate>>, + remaining: usize, +} + +impl<'a, T> Iter<'a, T> { + fn new(set: &'a ItemSet) -> Self { + Self { inner: set.items.iter().enumerate(), remaining: set.len() } + } +} + +impl Clone for Iter<'_, T> { + fn clone(&self) -> Self { + Self { inner: self.inner.clone(), remaining: self.remaining } + } +} + +impl fmt::Debug for Iter<'_, T> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Iter").field("remaining", &self.remaining).finish() + } +} + +impl Default for Iter<'_, T> { + fn default() -> Self { + let empty: &[ItemSlot] = &[]; + Self { inner: empty.iter().enumerate(), remaining: 0 } + } +} + +impl<'a, T> Iterator for Iter<'a, T> { + type Item = (ItemIndex, &'a T); + + #[inline] + fn next(&mut self) -> Option { + for (i, slot) in self.inner.by_ref() { + if let ItemSlot::Occupied(v) = slot { + debug_assert!( + self.remaining > 0, + "iterator yielded more items than ItemSet::len()", + ); + self.remaining -= 1; + return Some((ItemIndex::new(i as u32), v)); + } + } + None + } + + #[inline] + fn size_hint(&self) -> (usize, Option) { + (self.remaining, Some(self.remaining)) + } +} + +impl ExactSizeIterator for Iter<'_, T> { + #[inline] + fn len(&self) -> usize { + self.remaining + } +} + +impl FusedIterator for Iter<'_, T> {} + +/// An iterator over `(index, &mut item)` pairs in an [`ItemSet`]. +pub(crate) struct IterMut<'a, T> { + inner: core::iter::Enumerate>>, + remaining: usize, +} + +impl<'a, T> IterMut<'a, T> { + fn new(set: &'a mut ItemSet) -> Self { + let remaining = set.len(); + Self { inner: set.items.iter_mut().enumerate(), remaining } + } +} + +impl fmt::Debug for IterMut<'_, T> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("IterMut").field("remaining", &self.remaining).finish() + } +} + +impl<'a, T> Iterator for IterMut<'a, T> { + type Item = (ItemIndex, &'a mut T); + + #[inline] + fn next(&mut self) -> Option { + for (i, slot) in self.inner.by_ref() { + if let ItemSlot::Occupied(v) = slot { + debug_assert!( + self.remaining > 0, + "iterator yielded more items than ItemSet::len()", + ); + self.remaining -= 1; + return Some((ItemIndex::new(i as u32), v)); + } + } + None + } + + #[inline] + fn size_hint(&self) -> (usize, Option) { + (self.remaining, Some(self.remaining)) + } +} + +impl ExactSizeIterator for IterMut<'_, T> { + #[inline] + fn len(&self) -> usize { + self.remaining + } +} + +impl FusedIterator for IterMut<'_, T> {} + +/// An iterator over `&item` references in an [`ItemSet`]. +pub(crate) struct Values<'a, T> { + inner: core::slice::Iter<'a, ItemSlot>, + remaining: usize, +} + +impl<'a, T> Values<'a, T> { + fn new(set: &'a ItemSet) -> Self { + Self { inner: set.items.iter(), remaining: set.len() } + } +} + +impl Clone for Values<'_, T> { + fn clone(&self) -> Self { + Self { inner: self.inner.clone(), remaining: self.remaining } + } +} + +impl fmt::Debug for Values<'_, T> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Values").field("remaining", &self.remaining).finish() + } +} + +impl Default for Values<'_, T> { + fn default() -> Self { + let empty: &[ItemSlot] = &[]; + Self { inner: empty.iter(), remaining: 0 } + } +} + +impl<'a, T> Iterator for Values<'a, T> { + type Item = &'a T; + + #[inline] + fn next(&mut self) -> Option { + for slot in self.inner.by_ref() { + if let ItemSlot::Occupied(v) = slot { + debug_assert!( + self.remaining > 0, + "iterator yielded more items than ItemSet::len()", + ); + self.remaining -= 1; + return Some(v); + } + } + None + } + + #[inline] + fn size_hint(&self) -> (usize, Option) { + (self.remaining, Some(self.remaining)) + } +} + +impl ExactSizeIterator for Values<'_, T> { + #[inline] + fn len(&self) -> usize { + self.remaining + } +} + +impl FusedIterator for Values<'_, T> {} + +/// An iterator over `&mut item` references in an [`ItemSet`]. +pub(crate) struct ValuesMut<'a, T> { + inner: core::slice::IterMut<'a, ItemSlot>, + remaining: usize, +} + +impl<'a, T> ValuesMut<'a, T> { + fn new(set: &'a mut ItemSet) -> Self { + let remaining = set.len(); + Self { inner: set.items.iter_mut(), remaining } + } +} + +impl fmt::Debug for ValuesMut<'_, T> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("ValuesMut").field("remaining", &self.remaining).finish() + } +} + +impl<'a, T> Iterator for ValuesMut<'a, T> { + type Item = &'a mut T; + + #[inline] + fn next(&mut self) -> Option { + for slot in self.inner.by_ref() { + if let ItemSlot::Occupied(v) = slot { + debug_assert!( + self.remaining > 0, + "iterator yielded more items than ItemSet::len()", + ); + self.remaining -= 1; + return Some(v); + } + } + None + } + + #[inline] + fn size_hint(&self) -> (usize, Option) { + (self.remaining, Some(self.remaining)) + } +} + +impl ExactSizeIterator for ValuesMut<'_, T> { + #[inline] + fn len(&self) -> usize { + self.remaining + } +} + +impl FusedIterator for ValuesMut<'_, T> {} + +/// An owning iterator over the items in an [`ItemSet`]. +pub(crate) struct IntoValues { + inner: allocator_api2::vec::IntoIter, AllocWrapper>, + remaining: usize, +} + +impl IntoValues { + fn new(set: ItemSet) -> Self { + let remaining = set.len(); + let consuming = set.into_consuming(); + Self { inner: consuming.items.into_iter(), remaining } + } +} + +impl fmt::Debug for IntoValues { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("IntoValues") + .field("remaining", &self.remaining) + .finish() + } +} + +impl Iterator for IntoValues { + type Item = T; + + #[inline] + fn next(&mut self) -> Option { + for slot in self.inner.by_ref() { + if let ItemSlot::Occupied(v) = slot { + debug_assert!( + self.remaining > 0, + "iterator yielded more items than ItemSet::len()", + ); + self.remaining -= 1; + return Some(v); + } + } + None + } + + #[inline] + fn size_hint(&self) -> (usize, Option) { + (self.remaining, Some(self.remaining)) + } +} + +impl ExactSizeIterator for IntoValues { + #[inline] + fn len(&self) -> usize { + self.remaining + } +} + +impl FusedIterator for IntoValues {} + +/// An [`ItemSet`] consumed into an owned, by-index take-only version. +/// +/// Produced by [`ItemSet::into_consuming`]. The free chain is no longer +/// maintained from here on. +pub(crate) struct ConsumingItemSet { + items: Vec, AllocWrapper>, +} + +impl fmt::Debug for ConsumingItemSet { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("ConsumingItemSet") + .field("slots", &self.items.len()) + .finish() + } +} + +impl ConsumingItemSet { + /// Takes the item at `index`, leaving a `Vacant { next: SENTINEL }` + /// slot behind. + /// + /// Returns `None` if `index` is out of bounds or the slot has + /// already been taken. O(1) regardless of position. + #[cfg_attr(not(feature = "std"), expect(dead_code))] + #[inline] + pub(crate) fn take(&mut self, index: ItemIndex) -> Option { + let slot = self.items.get_mut(index.as_u32() as usize)?; + if !slot.is_occupied() { + return None; + } + // The free chain is no longer maintained in this view, so any + // `next` value is fine. `SENTINEL` is a natural choice. + let ItemSlot::Occupied(v) = core::mem::replace( + slot, + ItemSlot::Vacant { next: ItemIndex::SENTINEL }, + ) else { + unreachable!("is_occupied was just checked") + }; + Some(v) + } +} + +#[cfg(all(test, feature = "std"))] +mod tests { + use super::*; + use crate::internal::ValidateCompact; + + fn ix(value: u32) -> ItemIndex { + ItemIndex::new(value) + } + + #[test] + fn shrink_to_fit_compacts_middle_holes() { + let mut set = ItemSet::::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(); + + assert_eq!(set.len(), 3); + set.validate(ValidateCompact::Compact).unwrap(); + assert_eq!(&[set[ix(0)], set[ix(1)], set[ix(2)]], &[0, 20, 40]); + + assert!(!remap.is_identity()); + assert_eq!(remap.remap(ix(0)), ix(0)); + assert_eq!(remap.remap(ix(2)), ix(1)); + assert_eq!(remap.remap(ix(4)), ix(2)); + } + + #[test] + fn shrink_to_fit_without_holes_returns_empty_remap() { + let mut set = ItemSet::::new(); + for i in 0..4 { + set.assert_can_grow().insert(i); + } + let remap = set.shrink_to_fit(); + assert!(remap.is_identity()); + set.validate(ValidateCompact::Compact) + .expect("a hole-free set is trivially compact after shrink_to_fit"); + } + + #[test] + fn free_chain_is_lifo_and_well_formed() { + let mut set = ItemSet::::new(); + for i in 0..6 { + set.assert_can_grow().insert(i * 10); + } + // Remove 1, then 3, then 5 — free chain after is [5 -> 3 -> 1]. + assert_eq!(set.remove(ix(1)), Some(10)); + assert_eq!(set.remove(ix(3)), Some(30)); + assert_eq!(set.remove(ix(5)), Some(50)); + set.validate(ValidateCompact::NonCompact).unwrap(); + assert_eq!(set.len(), 3); + + // LIFO reuse: next three inserts go into 5, 3, 1. + assert_eq!(set.assert_can_grow().insert(100), ix(5)); + assert_eq!(set.assert_can_grow().insert(200), ix(3)); + assert_eq!(set.assert_can_grow().insert(300), ix(1)); + set.validate(ValidateCompact::Compact).unwrap(); + assert_eq!(set[ix(1)], 300); + assert_eq!(set[ix(3)], 200); + assert_eq!(set[ix(5)], 100); + + // Fourth insert allocates a new slot. + assert_eq!(set.assert_can_grow().insert(400), ix(6)); + } + + #[test] + fn clone_preserves_free_chain_and_values() { + let mut set = ItemSet::::new(); + for i in 0..4 { + set.assert_can_grow().insert(i); + } + set.remove(ix(1)); + set.remove(ix(2)); + + let cloned = set.clone(); + cloned.validate(ValidateCompact::NonCompact).unwrap(); + assert_eq!(cloned.len(), set.len()); + assert_eq!(cloned.get(ix(0)), Some(&0)); + assert_eq!(cloned.get(ix(1)), None); + assert_eq!(cloned.get(ix(2)), None); + assert_eq!(cloned.get(ix(3)), Some(&3)); + } +} diff --git a/crates/iddqd/src/tri_hash_map/imp.rs b/crates/iddqd/src/tri_hash_map/imp.rs index 9ce1171f..a56c6af1 100644 --- a/crates/iddqd/src/tri_hash_map/imp.rs +++ b/crates/iddqd/src/tri_hash_map/imp.rs @@ -923,9 +923,7 @@ impl TriHashMap { &mut self, additional: usize, ) -> Result<(), crate::errors::TryReserveError> { - self.items - .try_reserve(additional) - .map_err(crate::errors::TryReserveError::from_hashbrown)?; + self.items.try_reserve(additional)?; let items = &self.items; let state = &self.tables.state; self.tables @@ -995,7 +993,12 @@ impl TriHashMap { /// # } /// ``` pub fn shrink_to_fit(&mut self) { - self.items.shrink_to_fit(); + let remap = self.items.shrink_to_fit(); + 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); + } let items = &self.items; let state = &self.tables.state; self.tables @@ -1066,7 +1069,12 @@ impl TriHashMap { /// # } /// ``` pub fn shrink_to(&mut self, min_capacity: usize) { - self.items.shrink_to(min_capacity); + let remap = self.items.shrink_to(min_capacity); + 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); + } let items = &self.items; let state = &self.tables.state; self.tables @@ -1223,7 +1231,7 @@ impl TriHashMap { self.tables.validate(self.len(), compactness)?; // Check that the indexes are all correct. - for (&ix, item) in self.items.iter() { + for (ix, item) in self.items.iter() { let key1 = item.key1(); let key2 = item.key2(); let key3 = item.key3(); diff --git a/crates/iddqd/src/tri_hash_map/iter.rs b/crates/iddqd/src/tri_hash_map/iter.rs index 8c472d2a..da2cbe4e 100644 --- a/crates/iddqd/src/tri_hash_map/iter.rs +++ b/crates/iddqd/src/tri_hash_map/iter.rs @@ -2,13 +2,11 @@ use super::{RefMut, tables::TriHashMapTables}; use crate::{ DefaultHashBuilder, TriHashItem, support::{ - ItemIndex, - alloc::{AllocWrapper, Allocator, Global}, - item_set::ItemSet, + alloc::{Allocator, Global}, + item_set::{self, ItemSet}, }, }; use core::{hash::BuildHasher, iter::FusedIterator}; -use hashbrown::hash_map; /// An iterator over the elements of a [`TriHashMap`] by shared reference. /// Created by [`TriHashMap::iter`]. @@ -21,7 +19,7 @@ use hashbrown::hash_map; /// [`HashMap`]: std::collections::HashMap #[derive(Clone, Debug, Default)] pub struct Iter<'a, T: TriHashItem> { - inner: hash_map::Values<'a, ItemIndex, T>, + inner: item_set::Values<'a, T>, } impl<'a, T: TriHashItem> Iter<'a, T> { @@ -46,7 +44,6 @@ impl ExactSizeIterator for Iter<'_, T> { } } -// hash_map::Iter is a FusedIterator, so Iter is as well. impl FusedIterator for Iter<'_, T> {} /// An iterator over the elements of a [`TriHashMap`] by mutable reference. @@ -68,7 +65,7 @@ pub struct IterMut< A: Allocator = Global, > { tables: &'a TriHashMapTables, - inner: hash_map::ValuesMut<'a, ItemIndex, T>, + inner: item_set::ValuesMut<'a, T>, } impl<'a, T: TriHashItem, S: Clone + BuildHasher, A: Allocator> @@ -104,7 +101,6 @@ impl ExactSizeIterator } } -// hash_map::IterMut is a FusedIterator, so IterMut is as well. impl FusedIterator for IterMut<'_, T, S, A> { @@ -121,7 +117,7 @@ impl FusedIterator /// [`HashMap`]: std::collections::HashMap #[derive(Debug)] pub struct IntoIter { - inner: hash_map::IntoValues>, + inner: item_set::IntoValues, } impl IntoIter { @@ -146,5 +142,4 @@ impl ExactSizeIterator for IntoIter { } } -// hash_map::IterMut is a FusedIterator, so IterMut is as well. impl FusedIterator for IntoIter {} diff --git a/crates/iddqd/tests/integration/bi_hash_map.rs b/crates/iddqd/tests/integration/bi_hash_map.rs index 4ac8a0c9..2e33d21f 100644 --- a/crates/iddqd/tests/integration/bi_hash_map.rs +++ b/crates/iddqd/tests/integration/bi_hash_map.rs @@ -44,11 +44,10 @@ fn debug_impls() { assert_eq!( format!("{map:?}"), - // This is a small-enough map that the order of iteration is - // deterministic. + // Iteration is in insertion order. "{{k1: 1, k2: 'a'}: SimpleItem { key1: 1, key2: 'a' }, \ - {k1: 10, k2: 'c'}: SimpleItem { key1: 10, key2: 'c' }, \ - {k1: 20, k2: 'b'}: SimpleItem { key1: 20, key2: 'b' }}", + {k1: 20, k2: 'b'}: SimpleItem { key1: 20, key2: 'b' }, \ + {k1: 10, k2: 'c'}: SimpleItem { key1: 10, key2: 'c' }}", ); assert_eq!( format!("{:?}", map.get1_mut(&1).unwrap()), @@ -67,7 +66,7 @@ fn debug_impls_borrowed() { assert_eq!( format!("{before:?}"), - r#"{{k1: "a", k2: [98, 48]}: BorrowedItem { key1: "a", key2: [98, 48], key3: "path0" }, {k1: "c", k2: [98, 50]}: BorrowedItem { key1: "c", key2: [98, 50], key3: "path2" }, {k1: "b", k2: [98, 49]}: BorrowedItem { key1: "b", key2: [98, 49], key3: "path1" }}"# + r#"{{k1: "a", k2: [98, 48]}: BorrowedItem { key1: "a", key2: [98, 48], key3: "path0" }, {k1: "b", k2: [98, 49]}: BorrowedItem { key1: "b", key2: [98, 49], key3: "path1" }, {k1: "c", k2: [98, 50]}: BorrowedItem { key1: "c", key2: [98, 50], key3: "path2" }}"# ); #[cfg(feature = "daft")] @@ -84,7 +83,7 @@ fn debug_impls_borrowed() { let diff = before.diff(&after).by_unique(); assert_eq!( format!("{diff:?}"), - r#"Diff { common: {{k1: "a", k2: [98, 48]}: IdLeaf { before: BorrowedItem { key1: "a", key2: [98, 48], key3: "path0" }, after: BorrowedItem { key1: "a", key2: [98, 48], key3: "path0" } }}, added: {{k1: "d", k2: [98, 52]}: BorrowedItem { key1: "d", key2: [98, 52], key3: "path4" }, {k1: "c", k2: [98, 51]}: BorrowedItem { key1: "c", key2: [98, 51], key3: "path3" }}, removed: {{k1: "c", k2: [98, 50]}: BorrowedItem { key1: "c", key2: [98, 50], key3: "path2" }, {k1: "b", k2: [98, 49]}: BorrowedItem { key1: "b", key2: [98, 49], key3: "path1" }} }"# + r#"Diff { common: {{k1: "a", k2: [98, 48]}: IdLeaf { before: BorrowedItem { key1: "a", key2: [98, 48], key3: "path0" }, after: BorrowedItem { key1: "a", key2: [98, 48], key3: "path0" } }}, added: {{k1: "c", k2: [98, 51]}: BorrowedItem { key1: "c", key2: [98, 51], key3: "path3" }, {k1: "d", k2: [98, 52]}: BorrowedItem { key1: "d", key2: [98, 52], key3: "path4" }}, removed: {{k1: "b", k2: [98, 49]}: BorrowedItem { key1: "b", key2: [98, 49], key3: "path1" }, {k1: "c", k2: [98, 50]}: BorrowedItem { key1: "c", key2: [98, 50], key3: "path2" }} }"# ); } } @@ -280,17 +279,11 @@ enum Operation { impl Operation { fn compactness_change(&self) -> CompactnessChange { match self { - // `shrink_to_fit` / `shrink_to` flow through hashbrown's - // rehash path the same way `reserve` does; like `reserve` - // they touch the allocation, not the item set's index - // space, so compactness is unchanged. Operation::InsertUnique(_) | Operation::Get1(_) | Operation::Get2(_) | Operation::Reserve(_) - | Operation::TryReserve(_) - | Operation::ShrinkToFit - | Operation::ShrinkTo(_) => CompactnessChange::NoChange, + | Operation::TryReserve(_) => CompactnessChange::NoChange, // The act of removing items, including calls to insert_overwrite, // can make the map non-compact. Operation::InsertOverwrite(_) @@ -299,8 +292,12 @@ impl Operation { | Operation::RetainValueContains(_, _) | Operation::RetainModulo(_, _, _) | Operation::Extend(_) => CompactnessChange::NoLongerCompact, - // Clear always makes the map compact (empty). - Operation::Clear => CompactnessChange::BecomesCompact, + // Clear always makes the map compact (empty). Shrink + // fully compacts the backing store, restoring the + // `Compact` invariant. + Operation::Clear + | Operation::ShrinkToFit + | Operation::ShrinkTo(_) => CompactnessChange::BecomesCompact, } } } diff --git a/crates/iddqd/tests/integration/id_hash_map.rs b/crates/iddqd/tests/integration/id_hash_map.rs index 5df020db..275caa41 100644 --- a/crates/iddqd/tests/integration/id_hash_map.rs +++ b/crates/iddqd/tests/integration/id_hash_map.rs @@ -38,9 +38,8 @@ fn debug_impls() { assert_eq!( format!("{map:?}"), - // This is a small-enough map that the order of iteration is - // deterministic. - r#"{1: SimpleItem { key: 1 }, 10: SimpleItem { key: 10 }, 20: SimpleItem { key: 20 }}"# + // Iteration is in insertion order. + r#"{1: SimpleItem { key: 1 }, 20: SimpleItem { key: 20 }, 10: SimpleItem { key: 10 }}"# ); assert_eq!( format!("{:?}", map.get_mut(&1).unwrap()), @@ -59,7 +58,7 @@ fn debug_impls_borrowed() { assert_eq!( format!("{before:?}"), - r#"{"a": BorrowedItem { key1: "a", key2: [98, 48], key3: "path0" }, "c": BorrowedItem { key1: "c", key2: [98, 50], key3: "path2" }, "b": BorrowedItem { key1: "b", key2: [98, 49], key3: "path1" }}"# + r#"{"a": BorrowedItem { key1: "a", key2: [98, 48], key3: "path0" }, "b": BorrowedItem { key1: "b", key2: [98, 49], key3: "path1" }, "c": BorrowedItem { key1: "c", key2: [98, 50], key3: "path2" }}"# ); #[cfg(feature = "daft")] @@ -232,16 +231,10 @@ enum Operation { impl Operation { fn compactness_change(&self) -> CompactnessChange { match self { - // `shrink_to_fit` / `shrink_to` flow through hashbrown's - // rehash path the same way `reserve` does; like `reserve` - // they touch the allocation, not the item set's index - // space, so compactness is unchanged. Operation::InsertUnique(_) | Operation::Get(_) | Operation::Reserve(_) - | Operation::TryReserve(_) - | Operation::ShrinkToFit - | Operation::ShrinkTo(_) => CompactnessChange::NoChange, + | Operation::TryReserve(_) => CompactnessChange::NoChange, // The act of removing items, including calls to insert_overwrite, // can make the map non-compact. Operation::InsertOverwrite(_) @@ -249,8 +242,12 @@ impl Operation { | Operation::RetainValueContains(_, _) | Operation::RetainModulo(_, _, _) | Operation::Extend(_) => CompactnessChange::NoLongerCompact, - // Clear always makes the map compact (empty). - Operation::Clear => CompactnessChange::BecomesCompact, + // Clear empties the map, so it is de-facto compact. Shrink + // operations fully compact the backing store, restoring the + // `Compact` invariant. + Operation::Clear + | Operation::ShrinkToFit + | Operation::ShrinkTo(_) => CompactnessChange::BecomesCompact, } } } @@ -363,10 +360,12 @@ fn proptest_ops( } Operation::ShrinkToFit => { map.shrink_to_fit(); + // The naive map has no shrink operation. map.validate(compactness).expect("map should be valid"); } Operation::ShrinkTo(min_capacity) => { map.shrink_to(min_capacity); + // The naive map has no shrink operation. map.validate(compactness).expect("map should be valid"); } } diff --git a/crates/iddqd/tests/integration/id_ord_map.rs b/crates/iddqd/tests/integration/id_ord_map.rs index dc53e64e..aaff92b9 100644 --- a/crates/iddqd/tests/integration/id_ord_map.rs +++ b/crates/iddqd/tests/integration/id_ord_map.rs @@ -268,6 +268,8 @@ enum Operation { ), // clear is set to a lower weight since it makes the map empty. Clear, + ShrinkToFit, + ShrinkTo(#[strategy(0..256_usize)] usize), } impl Operation { @@ -288,8 +290,12 @@ impl Operation { | Operation::RetainValueContains(_, _) | Operation::RetainModulo(_, _, _) | Operation::Extend(_) => CompactnessChange::NoLongerCompact, - // Clear always makes the map compact (empty). - Operation::Clear => CompactnessChange::BecomesCompact, + // Clear always makes the map compact (empty). Shrink + // fully compacts the backing store, restoring the + // `Compact` invariant. + Operation::Clear + | Operation::ShrinkToFit + | Operation::ShrinkTo(_) => CompactnessChange::BecomesCompact, } } } @@ -464,6 +470,16 @@ fn proptest_ops( map.validate(compactness, ValidateChaos::No) .expect("map should be valid"); } + Operation::ShrinkToFit => { + map.shrink_to_fit(); + map.validate(compactness, ValidateChaos::No) + .expect("map should be valid"); + } + Operation::ShrinkTo(min_capacity) => { + map.shrink_to(min_capacity); + map.validate(compactness, ValidateChaos::No) + .expect("map should be valid"); + } } // Check that the iterators work correctly. diff --git a/crates/iddqd/tests/integration/tri_hash_map.rs b/crates/iddqd/tests/integration/tri_hash_map.rs index 5dc4c952..fedc4bdc 100644 --- a/crates/iddqd/tests/integration/tri_hash_map.rs +++ b/crates/iddqd/tests/integration/tri_hash_map.rs @@ -51,11 +51,10 @@ fn debug_impls() { assert_eq!( format!("{map:?}"), - // This is a small-enough map that the order of iteration is - // deterministic. + // Iteration is in insertion order. "{{k1: 1, k2: 'a', k3: 0}: SimpleItem { key1: 1, key2: 'a', key3: 0 }, \ - {k1: 10, k2: 'c', k3: 2}: SimpleItem { key1: 10, key2: 'c', key3: 2 }, \ - {k1: 20, k2: 'b', k3: 1}: SimpleItem { key1: 20, key2: 'b', key3: 1 }}", + {k1: 20, k2: 'b', k3: 1}: SimpleItem { key1: 20, key2: 'b', key3: 1 }, \ + {k1: 10, k2: 'c', k3: 2}: SimpleItem { key1: 10, key2: 'c', key3: 2 }}", ); assert_eq!( format!("{:?}", map.get1_mut(&1).unwrap()), @@ -74,7 +73,7 @@ fn debug_impls_borrowed() { assert_eq!( format!("{before:?}"), - r#"{{k1: "a", k2: [98, 48], k3: "path0"}: BorrowedItem { key1: "a", key2: [98, 48], key3: "path0" }, {k1: "c", k2: [98, 50], k3: "path2"}: BorrowedItem { key1: "c", key2: [98, 50], key3: "path2" }, {k1: "b", k2: [98, 49], k3: "path1"}: BorrowedItem { key1: "b", key2: [98, 49], key3: "path1" }}"# + r#"{{k1: "a", k2: [98, 48], k3: "path0"}: BorrowedItem { key1: "a", key2: [98, 48], key3: "path0" }, {k1: "b", k2: [98, 49], k3: "path1"}: BorrowedItem { key1: "b", key2: [98, 49], key3: "path1" }, {k1: "c", k2: [98, 50], k3: "path2"}: BorrowedItem { key1: "c", key2: [98, 50], key3: "path2" }}"# ); #[cfg(feature = "daft")] @@ -91,7 +90,7 @@ fn debug_impls_borrowed() { let diff = before.diff(&after).by_unique(); assert_eq!( format!("{diff:?}"), - r#"Diff { common: {{k1: "a", k2: [98, 48], k3: "path0"}: IdLeaf { before: BorrowedItem { key1: "a", key2: [98, 48], key3: "path0" }, after: BorrowedItem { key1: "a", key2: [98, 48], key3: "path0" } }}, added: {{k1: "d", k2: [98, 52], k3: "path4"}: BorrowedItem { key1: "d", key2: [98, 52], key3: "path4" }, {k1: "c", k2: [98, 51], k3: "path3"}: BorrowedItem { key1: "c", key2: [98, 51], key3: "path3" }}, removed: {{k1: "c", k2: [98, 50], k3: "path2"}: BorrowedItem { key1: "c", key2: [98, 50], key3: "path2" }, {k1: "b", k2: [98, 49], k3: "path1"}: BorrowedItem { key1: "b", key2: [98, 49], key3: "path1" }} }"# + r#"Diff { common: {{k1: "a", k2: [98, 48], k3: "path0"}: IdLeaf { before: BorrowedItem { key1: "a", key2: [98, 48], key3: "path0" }, after: BorrowedItem { key1: "a", key2: [98, 48], key3: "path0" } }}, added: {{k1: "c", k2: [98, 51], k3: "path3"}: BorrowedItem { key1: "c", key2: [98, 51], key3: "path3" }, {k1: "d", k2: [98, 52], k3: "path4"}: BorrowedItem { key1: "d", key2: [98, 52], key3: "path4" }}, removed: {{k1: "b", k2: [98, 49], k3: "path1"}: BorrowedItem { key1: "b", key2: [98, 49], key3: "path1" }, {k1: "c", k2: [98, 50], k3: "path2"}: BorrowedItem { key1: "c", key2: [98, 50], key3: "path2" }} }"# ); } } @@ -310,18 +309,12 @@ enum Operation { impl Operation { fn compactness_change(&self) -> CompactnessChange { match self { - // `shrink_to_fit` / `shrink_to` flow through hashbrown's - // rehash path the same way `reserve` does; like `reserve` - // they touch the allocation, not the item set's index - // space, so compactness is unchanged. Operation::InsertUnique(_) | Operation::Get1(_) | Operation::Get2(_) | Operation::Get3(_) | Operation::Reserve(_) - | Operation::TryReserve(_) - | Operation::ShrinkToFit - | Operation::ShrinkTo(_) => CompactnessChange::NoChange, + | Operation::TryReserve(_) => CompactnessChange::NoChange, // The act of removing items, including calls to insert_overwrite, // can make the map non-compact. Operation::InsertOverwrite(_) @@ -331,8 +324,12 @@ impl Operation { | Operation::RetainValueContains(_, _) | Operation::RetainModulo(_, _, _) | Operation::Extend(_) => CompactnessChange::NoLongerCompact, - // Clear always makes the map compact (empty). - Operation::Clear => CompactnessChange::BecomesCompact, + // Clear always makes the map compact (empty). Shrink + // fully compacts the backing store, restoring the + // `Compact` invariant. + Operation::Clear + | Operation::ShrinkToFit + | Operation::ShrinkTo(_) => CompactnessChange::BecomesCompact, } } } diff --git a/crates/iddqd/tests/snapshots/map_sizes.txt b/crates/iddqd/tests/snapshots/map_sizes.txt index c9077c6a..af1d0932 100644 --- a/crates/iddqd/tests/snapshots/map_sizes.txt +++ b/crates/iddqd/tests/snapshots/map_sizes.txt @@ -1,10 +1,10 @@ -IdHashMap: 80 -IdHashMap: 88 +IdHashMap: 72 +IdHashMap: 80 -BiHashMap: 112 -BiHashMap: 120 +BiHashMap: 104 +BiHashMap: 112 -TriHashMap: 144 -TriHashMap: 152 +TriHashMap: 136 +TriHashMap: 144 -IdOrdMap: 72 +IdOrdMap: 64