diff --git a/crates/iddqd/src/id_ord_map/iter.rs b/crates/iddqd/src/id_ord_map/iter.rs index 77aa693..ef50d55 100644 --- a/crates/iddqd/src/id_ord_map/iter.rs +++ b/crates/iddqd/src/id_ord_map/iter.rs @@ -3,9 +3,9 @@ use crate::support::{ alloc::Global, borrow::DormantMutRef, btree_table, - item_set::{ConsumingItemSet, ItemSet, ItemSlot}, + item_set::{ConsumingItemSet, ItemSet, ItemSlotsPtr}, }; -use core::{hash::Hash, iter::FusedIterator, marker::PhantomData}; +use core::{hash::Hash, iter::FusedIterator}; /// An iterator over the elements of an [`IdOrdMap`] by shared reference. /// @@ -48,40 +48,6 @@ 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. @@ -95,7 +61,7 @@ pub struct IterMut<'a, T: IdOrdItem> where T::Key<'a>: Hash, { - items: ItemSetPtr<'a, T>, + items: ItemSlotsPtr<'a, T>, tables: &'a IdOrdMapTables, iter: btree_table::Iter<'a>, } @@ -108,10 +74,8 @@ where items: &'a mut ItemSet, tables: &'a IdOrdMapTables, ) -> Self { - let slot_count = items.slot_count(); - let start_ptr = items.start_ptr(); Self { - items: ItemSetPtr { start_ptr, slot_count, _marker: PhantomData }, + items: ItemSlotsPtr::new(items.slots_mut()), tables, iter: tables.key_to_item.iter(), } @@ -127,43 +91,11 @@ where #[inline] fn next(&mut self) -> Option { let index = self.iter.next()?; - 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: - // - // * `self.items.ptr.add(raw_index)` points at valid memory. - // * There are no overlapping mutable borrows of the same memory. - // - // This is shown by the following observations: - // - // * 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") - }; + + // SAFETY: The B-tree is a set, so each call to `self.iter.next()` + // yields a distinct `index`. Therefore the `&mut T` references that + // `get_mut` hands out across iterations never alias. + let item: &'a mut T = unsafe { self.items.get_mut(index) }; let (hash, dormant) = { let (item, dormant) = DormantMutRef::new(item); @@ -171,8 +103,11 @@ where (hash, dormant) }; - // SAFETY: item is dropped above, and self is no longer used - // after this point. + // SAFETY: The `&mut T` that `DormantMutRef::new` produced inside + // the block above (and used for hashing) was dropped when the + // block closed, so the dormant ref is now the unique borrow of + // the slot. The `self.tables.state()` access below touches a + // different allocation and does not alias. let item = unsafe { dormant.awaken() }; Some(RefMut::new(self.tables.state().clone(), hash, item)) diff --git a/crates/iddqd/src/support/item_set.rs b/crates/iddqd/src/support/item_set.rs index dc0d393..0be3c34 100644 --- a/crates/iddqd/src/support/item_set.rs +++ b/crates/iddqd/src/support/item_set.rs @@ -57,6 +57,7 @@ use allocator_api2::vec::Vec; use core::{ fmt, iter::FusedIterator, + marker::PhantomData, ops::{Index, IndexMut}, }; @@ -186,10 +187,10 @@ impl<'a, T, A: Allocator> GrowHandle<'a, T, A> { /// 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. +/// Exposed at `pub(crate)` because [`ItemSet::slots_mut`] hands out a slot +/// slice for `ItemSlotsPtr` to build an allocator-agnostic raw pointer over. +/// 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. @@ -297,18 +298,14 @@ impl ItemSet { &self.items.allocator().0 } - /// 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. + /// Returns the backing slot buffer as a mutable slice. + /// + /// Used by [`ItemSlotsPtr::new`] to build an allocator-agnostic raw + /// pointer over the slot buffer for the per-map `IterMut` iterators. #[inline] #[cfg_attr(not(feature = "std"), expect(dead_code))] - pub(crate) fn slot_count(&self) -> usize { - self.items.len() + pub(crate) fn slots_mut(&mut self) -> &mut [ItemSlot] { + &mut self.items } pub(crate) fn validate( @@ -686,6 +683,121 @@ impl IndexMut for ItemSet { } } +// --- ItemSlotsPtr --------------------------------------------------------- + +/// A raw pointer into the start of an [`ItemSet`]'s slot buffer, with the same +/// thread-safety properties as `&'a mut [ItemSlot]`. +/// +/// This is used by iterators that yield `&mut T` references one at a time and +/// need to keep handing out distinct references for the duration of the borrow. +/// +/// This is equivalent to reborrowing the slot slice each iteration and using +/// unsafe code for lifetime extension, but that would invalidate previously +/// yielded `&mut T` children under Stacked Borrows. By using a raw pointer, we +/// keep the original mutable borrow live for the full iteration while still +/// being able to hand out element references that outlive `&mut self`. (Note +/// that the lifetime extension approach is not rejected by Tree Borrows, which +/// indicates that it's probably sound. But it's nice for iddqd to pass both +/// Stacked and Tree Borrows.) +/// +/// The only way to read a slot through this pointer is the `unsafe` +/// [`Self::get_mut`] method, which the caller is responsible for invoking +/// with each `index` at most once across the lifetime of `'a`. +pub(crate) struct ItemSlotsPtr<'a, T> { + /// 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 slot slice for `'a` so the raw pointer stays live, and so + /// that variance and drop-check work the same as `&'a mut [ItemSlot]`. + /// This deliberately does not mention `ItemSet` so the iterator's + /// public surface stays allocator-agnostic. + _marker: PhantomData<&'a mut [ItemSlot]>, +} + +impl fmt::Debug for ItemSlotsPtr<'_, T> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("ItemSlotsPtr") + .field("start_ptr", &self.start_ptr) + .field("slot_count", &self.slot_count) + .finish() + } +} + +// SAFETY: We treat the `*mut ItemSlot` as the `&'a mut [ItemSlot]` that +// the `PhantomData` already encodes. Auto-trait inference would give us +// `Send`/`Sync` for that slice under `T: Send` (resp. `T: Sync`); raw +// pointers don't carry the auto-traits, so we state the same bound here. +unsafe impl Send for ItemSlotsPtr<'_, T> {} +// SAFETY: see the `Send` impl above. +unsafe impl Sync for ItemSlotsPtr<'_, T> {} + +impl<'a, T> ItemSlotsPtr<'a, T> { + /// Captures a raw pointer into the slot buffer. + /// + /// The returned handle borrows `slots` for `'a`. + #[inline] + #[cfg_attr(not(feature = "std"), expect(dead_code))] + pub(crate) fn new(slots: &'a mut [ItemSlot]) -> Self { + Self { + start_ptr: slots.as_mut_ptr(), + slot_count: slots.len(), + _marker: PhantomData, + } + } + + /// Returns a mutable reference to the item at `index`. + /// + /// The lifetime of the returned reference is the borrow `'a` captured at + /// construction. This lets callers (typically iterators) hand the reference + /// out and then call `get_mut` again for a different index without + /// invalidating prior yields. + /// + /// # Panics + /// + /// Panics if `index` is out of bounds for the captured slot buffer or + /// if the slot is `Vacant`. Both indicate a stale or invalid index + /// reached us from an outer index table: this is a programmer error, + /// not undefined behavior. + /// + /// # Safety + /// + /// The caller must guarantee that, across the lifetime of this + /// [`ItemSlotsPtr`], each `index` value is passed to `get_mut` at most + /// once. That is the only thing that keeps the returned `&mut T` + /// references disjoint and aliasing-free. + #[inline] + #[cfg_attr(not(feature = "std"), expect(dead_code))] + pub(crate) unsafe fn get_mut(&mut self, index: ItemIndex) -> &'a mut T { + let raw_index = index.as_u32() as usize; + // Belt-and-suspenders bounds check. The outer index tables only ever + // store indexes that point at occupied slots, but we panic rather + // than risk UB if a future change inadvertently breaks that invariant. + assert!( + raw_index < self.slot_count, + "ItemSlotsPtr index {raw_index} should be in bounds \ + for slot count {}", + self.slot_count, + ); + // SAFETY: + // + // * `raw_index < self.slot_count`, so `self.start_ptr.add(raw_index)` + // is in-bounds for the original slot slice. + // * `ItemSlotsPtr::new` mutably borrowed the slot slice for `'a`, so + // no other code can touch it for the duration. + // * The caller's distinctness contract guarantees that we never hand + // out two `&mut T` references to the same slot. + // * The `expect` below verifies that the slot is `Occupied`; an outer + // index table that points at a `Vacant` slot is a programmer error + // that we surface as a panic rather than UB. + unsafe { + (*self.start_ptr.add(raw_index)) + .as_mut() + .expect("ItemSlotsPtr index points at an occupied slot") + } + } +} + // --- Iterators ---------------------------------------------------------- /// An iterator over `(index, &item)` pairs in an [`ItemSet`].