Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
93 changes: 14 additions & 79 deletions crates/iddqd/src/id_ord_map/iter.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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.
///
Expand Down Expand Up @@ -48,40 +48,6 @@ impl<T: IdOrdItem> ExactSizeIterator for Iter<'_, T> {
// btree_set::Iter is a FusedIterator, so Iter is as well.
impl<T: IdOrdItem> 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<T, Global>`.
///
/// 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<T>,
// 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<T, Global>`.
_marker: PhantomData<&'a mut ItemSet<T, Global>>,
}

impl<T: IdOrdItem> 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<T, Global>`, which is `Send`/`Sync` iff `ItemSet<T, Global>` 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.
Expand All @@ -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>,
}
Expand All @@ -108,10 +74,8 @@ where
items: &'a mut ItemSet<T, Global>,
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(),
}
Expand All @@ -127,52 +91,23 @@ where
#[inline]
fn next(&mut self) -> Option<Self::Item> {
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);
let hash = self.tables.make_hash(item);
(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))
Expand Down
140 changes: 126 additions & 14 deletions crates/iddqd/src/support/item_set.rs
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,7 @@ use allocator_api2::vec::Vec;
use core::{
fmt,
iter::FusedIterator,
marker::PhantomData,
ops::{Index, IndexMut},
};

Expand Down Expand Up @@ -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<T> {
/// The slot holds a live value.
Expand Down Expand Up @@ -297,18 +298,14 @@ impl<T, A: Allocator> ItemSet<T, A> {
&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<T> {
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<T>] {
&mut self.items
}

pub(crate) fn validate(
Expand Down Expand Up @@ -686,6 +683,121 @@ impl<T, A: Allocator> IndexMut<ItemIndex> for ItemSet<T, A> {
}
}

// --- ItemSlotsPtr ---------------------------------------------------------

/// A raw pointer into the start of an [`ItemSet`]'s slot buffer, with the same
/// thread-safety properties as `&'a mut [ItemSlot<T>]`.
///
/// 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<T>,
/// 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<T>]`.
/// This deliberately does not mention `ItemSet<T, A>` so the iterator's
/// public surface stays allocator-agnostic.
_marker: PhantomData<&'a mut [ItemSlot<T>]>,
}

impl<T> 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<T>` as the `&'a mut [ItemSlot<T>]` 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<T: Send> Send for ItemSlotsPtr<'_, T> {}
// SAFETY: see the `Send` impl above.
unsafe impl<T: Sync> 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<T>]) -> 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`].
Expand Down
Loading