Skip to content

more resiliency to buggy user code in IdOrdMap #255

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
sunshowers merged 4 commits into from
May 22, 2026
Merged

more resiliency to buggy user code in IdOrdMap #255

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
c9fde77
[spr] initial version
sunshowers May 22, 2026
79882ff
minor fixes
sunshowers May 22, 2026
8a42cf6
more fixes:
sunshowers May 22, 2026
cc8487a
add Drop to panic safety PBTs
sunshowers May 22, 2026
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
47 changes: 30 additions & 17 deletions crates/iddqd/src/id_ord_map/imp.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -359,7 +359,7 @@ impl<T: IdOrdItem> IdOrdMap<T> {
/// Does nothing if capacity is already sufficient.
///
/// Note: This only reserves capacity in the item storage. The internal
/// [`BTreeSet`] used for key-to-item mapping does not support capacity
/// `BTreeMap` used for key-to-item mapping does not support capacity
/// reservation.
///
/// # Panics
Expand Down Expand Up @@ -402,7 +402,7 @@ impl<T: IdOrdItem> IdOrdMap<T> {
/// and possibly leaving some space in accordance with the resize policy.
///
/// Note: This only shrinks the item storage capacity. The internal
/// [`BTreeSet`] used for key-to-item mapping does not support capacity
/// `BTreeMap` used for key-to-item mapping does not support capacity
/// control.
///
/// # Examples
Expand Down Expand Up @@ -446,7 +446,7 @@ impl<T: IdOrdItem> IdOrdMap<T> {
/// If the current capacity is less than the lower limit, this is a no-op.
///
/// Note: This only shrinks the item storage capacity. The internal
/// [`BTreeSet`] used for key-to-item mapping does not support capacity
/// `BTreeMap` used for key-to-item mapping does not support capacity
/// control.
///
/// # Examples
Expand Down Expand Up @@ -1445,10 +1445,12 @@ impl<T: IdOrdItem> IdOrdMap<T> {
let grow_handle = self.items.assert_can_grow();
let next_index = grow_handle.next_index();
let key = value.key();
self.tables
.key_to_item
.insert(next_index, &key, |index| grow_handle[index].key());
let insert =
self.tables.key_to_item.prepare_insert(next_index, &key, |index| {
grow_handle[index].key()
});
drop(key);
insert.insert();
grow_handle.insert(value);

Ok(next_index)
Expand All @@ -1459,19 +1461,30 @@ impl<T: IdOrdItem> IdOrdMap<T> {
remove_index: ItemIndex,
) -> Option<T> {
// For panic safety, read the key while self.items still holds the slot,
// then remove from the B-tree *before* mutating self.items.
// then locate the B-tree entry before mutating self.items.
//
// `BTreeSet::remove` is panic-safe under user-`Ord` panics, since
// comparator panics during the internal binary search abort the
// operation without modifying the tree. (This is not a documented
// guarantee, but really the only reasonable way to implement a
// panic-safe B-tree map.) This means that a panic at this point leaves
// both items and the B-tree unmodified. `items.remove` afterwards
// cannot panic.
// `BTreeMap::entry` is panic-safe under user-`Ord` panics, since
// comparator panics during the internal binary search abort the lookup
// without modifying the tree. (This is not a documented guarantee, but
// really the only reasonable way to implement a panic-safe B-tree map.)
// This means that a panic at this point leaves both items and the
// B-tree unmodified. After the entry has been located, `drop(key)` can
// run user code, so it must happen before the B-tree or item slot is
// mutated.
//
// If BTreeMap::entry returns normally but misses due to already-broken
// tree ordering, the prepared remove falls back to exact-index cleanup
// before this item slot can be reused.
let key = self.items.get(remove_index)?.key();
self.tables
.key_to_item
.remove(remove_index, key, |index| self.items[index].key());
let remove = self.tables.key_to_item.prepare_remove(
remove_index,
&key,
|index| self.items[index].key(),
);
drop(key);
if !remove.remove() {
self.tables.key_to_item.remove_exact(remove_index);
}
Some(
self.items
.remove(remove_index)
Expand Down
132 changes: 102 additions & 30 deletions crates/iddqd/src/support/btree_table.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@
use super::{ItemIndex, item_set::IndexRemap, map_hash::MapHash};
use crate::internal::{TableValidationError, ValidateCompact};
use alloc::{
collections::{BTreeSet, btree_set},
collections::{BTreeMap, BTreeSet, btree_map},
vec::Vec,
};
use core::{
Expand All @@ -26,7 +26,7 @@ thread_local! {
///
/// This works by:
///
/// * We store an `Index` in the BTreeSet which knows how to call this
/// * We store an `Index` in the BTreeMap which knows how to call this
/// dynamic comparator.
/// * When we need to compare two `Index` values, we create a CmpDropGuard.
/// This struct is responsible for managing the lifetime of the
Expand All @@ -49,7 +49,7 @@ thread_local! {
/// If and when https://github.com/rust-lang/rust/issues/133549 lands,
/// this should become a viable option. Worth looking out for!
///
/// * Using a third-party BTreeSet implementation that allows passing in
/// * Using a third-party BTreeMap implementation that allows passing in
/// external comparators. As of 2025-05, there appear to be two options:
///
/// 1. copse (https://docs.rs/copse), which doesn't seem like a good fit
Expand All @@ -75,7 +75,7 @@ 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 {
items: BTreeSet<Index>,
items: BTreeMap<Index, ()>,
// We use foldhash directly here because we allow compiling with std but
// without the default-hasher. std turns on foldhash but not the default
// hasher.
Expand All @@ -85,7 +85,7 @@ pub(crate) struct MapBTreeTable {
impl MapBTreeTable {
pub(crate) const fn new() -> Self {
Self {
items: BTreeSet::new(),
items: BTreeMap::new(),
// FixedState::with_seed XORs the passed in seed with a fixed
// high-entropy value.
hash_state: foldhash::fast::FixedState::with_seed(0),
Expand Down Expand Up @@ -117,7 +117,7 @@ impl MapBTreeTable {
// value should not be stored.
let mut indexes: Vec<ItemIndex> =
Vec::with_capacity(expected_len);
for index in &self.items {
for index in self.items.keys() {
let v = index.value();
if v == Index::SENTINEL_VALUE {
return Err(TableValidationError::new(
Expand All @@ -139,7 +139,7 @@ impl MapBTreeTable {
// There should be no duplicates, and the sentinel value
// should not be stored.
let indexes: Vec<ItemIndex> =
self.items.iter().map(|ix| ix.value()).collect();
self.items.keys().map(|ix| ix.value()).collect();
let index_set: BTreeSet<ItemIndex> =
indexes.iter().copied().collect();
if index_set.len() != indexes.len() {
Expand All @@ -163,12 +163,12 @@ impl MapBTreeTable {

#[inline]
pub(crate) fn first(&self) -> Option<ItemIndex> {
self.items.first().map(|ix| ix.value())
self.items.first_key_value().map(|(ix, ())| ix.value())
}

#[inline]
pub(crate) fn last(&self) -> Option<ItemIndex> {
self.items.last().map(|ix| ix.value())
self.items.last_key_value().map(|(ix, ())| ix.value())
}

pub(crate) fn find_index<K, Q, F>(
Expand All @@ -185,11 +185,11 @@ impl MapBTreeTable {

let guard = CmpDropGuard::new(&f);

let ret = match self.items.get(&Index::sentinel()) {
Some(ix) if ix.value() == Index::SENTINEL_VALUE => {
let ret = match self.items.get_key_value(&Index::sentinel()) {
Some((ix, ())) if ix.value() == Index::SENTINEL_VALUE => {
panic!("internal map shouldn't store sentinel value")
}
Some(ix) => Some(ix.value()),
Some((ix, ())) => Some(ix.value()),
None => {
// The key is not in the table.
None
Expand All @@ -201,37 +201,78 @@ impl MapBTreeTable {
ret
}

pub(crate) fn insert<K, Q, F>(
pub(crate) fn prepare_insert<K, Q, F>(
&mut self,
index: ItemIndex,
key: &Q,
lookup: F,
) where
) -> PreparedBTreeInsert<'_>
where
K: Ord,
Q: ?Sized + Comparable<K>,
F: Fn(ItemIndex) -> K,
{
let f = insert_cmp(index, key, lookup);
let guard = CmpDropGuard::new(&f);

self.items.insert(Index::new(index));
let entry = match self.items.entry(Index::new(index)) {
btree_map::Entry::Vacant(entry) => entry,
btree_map::Entry::Occupied(_) => {
panic!("internal map already contains index {index}")
}
};

// drop(guard) isn't necessary, but we make it explicit
drop(guard);

PreparedBTreeInsert { entry }
}

pub(crate) fn remove<K, F>(&mut self, index: ItemIndex, key: K, lookup: F)
pub(crate) fn prepare_remove<K, F>(
&mut self,
index: ItemIndex,
key: &K,
lookup: F,
) -> PreparedBTreeRemove<'_>
where
F: Fn(ItemIndex) -> K,
K: Ord,
{
let f = insert_cmp(index, &key, lookup);
let f = insert_cmp(index, key, lookup);
let guard = CmpDropGuard::new(&f);

self.items.remove(&Index::new(index));
let entry = self.items.entry(Index::new(index));

// drop(guard) isn't necessary, but we make it explicit
drop(guard);

match entry {
btree_map::Entry::Vacant(_) => {
// The comparator-based search missed an entry that
// `remove_by_index` has just confirmed lives in the item set.
// The most likely cause is a hash-blind key mutation that
// `RefMut` could not detect: the tree's structural order is
// now wrong, so a binary search walks past the physical
// entry.
//
// This is a signal to the caller to fall back to the linear
// `remove_exact` fallback when it commits the change.
PreparedBTreeRemove { inner: PreparedBTreeRemoveInner::Missing }
}
btree_map::Entry::Occupied(entry) => PreparedBTreeRemove {
inner: PreparedBTreeRemoveInner::Occupied(entry),
},
}
}

pub(crate) fn remove_exact(&mut self, index: ItemIndex) {
// If an item key was changed without detection, the B-tree order can be
// wrong. A comparator-based search may then miss the physical entry for
// this index. Fall back to a linear exact-index cleanup before the
// item slot can be reused.
//
// (BTreeMap::retain does not re-sort the tree or recompute positions,
// so a comparator guard isn't necessary.)
self.items.retain(|stored_index, ()| stored_index.value() != index);
}

pub(crate) fn retain<F>(&mut self, mut f: F)
Expand All @@ -240,7 +281,7 @@ 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.value()));
self.items.retain(|index, ()| f(index.value()));
}

/// Rewrites every stored index via `remap`.
Expand All @@ -262,7 +303,7 @@ impl MapBTreeTable {
/// [`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() {
for idx in self.items.keys() {
let new = remap.remap(idx.value());
idx.set_value(new);
}
Expand All @@ -275,7 +316,7 @@ impl MapBTreeTable {
}

pub(crate) fn iter(&self) -> Iter<'_> {
Iter::new(self.items.iter())
Iter::new(self.items.keys())
}

pub(crate) fn into_iter(self) -> IntoIter {
Expand All @@ -293,11 +334,11 @@ impl MapBTreeTable {

#[derive(Clone, Debug)]
pub(crate) struct Iter<'a> {
inner: btree_set::Iter<'a, Index>,
inner: btree_map::Keys<'a, Index, ()>,
}

impl<'a> Iter<'a> {
fn new(inner: btree_set::Iter<'a, Index>) -> Self {
fn new(inner: btree_map::Keys<'a, Index, ()>) -> Self {
Self { inner }
}

Expand All @@ -316,11 +357,11 @@ impl<'a> Iterator for Iter<'a> {

#[derive(Debug)]
pub(crate) struct IntoIter {
inner: btree_set::IntoIter<Index>,
inner: btree_map::IntoIter<Index, ()>,
}

impl IntoIter {
fn new(inner: btree_set::IntoIter<Index>) -> Self {
fn new(inner: btree_map::IntoIter<Index, ()>) -> Self {
Self { inner }
}
}
Expand All @@ -329,7 +370,38 @@ impl Iterator for IntoIter {
type Item = ItemIndex;

fn next(&mut self) -> Option<Self::Item> {
self.inner.next().map(|index| index.value())
self.inner.next().map(|(index, ())| index.value())
}
}

pub(crate) struct PreparedBTreeInsert<'a> {
entry: btree_map::VacantEntry<'a, Index, ()>,
}

impl PreparedBTreeInsert<'_> {
pub(crate) fn insert(self) {
self.entry.insert(());
}
}

pub(crate) struct PreparedBTreeRemove<'a> {
inner: PreparedBTreeRemoveInner<'a>,
}

enum PreparedBTreeRemoveInner<'a> {
Occupied(btree_map::OccupiedEntry<'a, Index, ()>),
Missing,
}

impl PreparedBTreeRemove<'_> {
pub(crate) fn remove(self) -> bool {
match self.inner {
PreparedBTreeRemoveInner::Occupied(entry) => {
entry.remove_entry();
true
}
PreparedBTreeRemoveInner::Missing => false,
}
}
}

Expand Down Expand Up @@ -613,13 +685,13 @@ mod tests {
for ix in [0u32, 2, 4] {
let ix = ItemIndex::new(ix);
let key = pre_lookup(ix);
table.insert(ix, &key, pre_lookup);
table.prepare_insert(ix, &key, pre_lookup).insert();
}
assert_eq!(table.len(), 3);
assert_eq!(
table
.items
.iter()
.keys()
.map(|i| i.value().as_u32())
.collect::<alloc::vec::Vec<_>>(),
[0u32, 2, 4],
Expand All @@ -636,7 +708,7 @@ mod tests {
assert_eq!(
table
.items
.iter()
.keys()
.map(|i| i.value().as_u32())
.collect::<alloc::vec::Vec<_>>(),
[0u32, 1, 2],
Expand Down
Loading
Loading