From 85e86118f86b1e3d8a0bf703a488593824ec0add Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Thu, 5 Mar 2026 22:49:15 -0500 Subject: [PATCH 1/2] Add `MultiIndexed::iter_mut` --- ext/crates/once/src/multiindexed/iter.rs | 350 ++++++++++++++------- ext/crates/once/src/multiindexed/kdtrie.rs | 4 + ext/crates/once/src/multiindexed/mod.rs | 75 +++++ 3 files changed, 320 insertions(+), 109 deletions(-) diff --git a/ext/crates/once/src/multiindexed/iter.rs b/ext/crates/once/src/multiindexed/iter.rs index 1d95e81d85..198ae67966 100644 --- a/ext/crates/once/src/multiindexed/iter.rs +++ b/ext/crates/once/src/multiindexed/iter.rs @@ -1,103 +1,52 @@ -use std::ops::Range; +use std::{marker::PhantomData, ops::Range}; use super::{KdTrie, node::Node}; use crate::MultiIndexed; -/// A single frame in the iteration stack, representing the current state of traversal. -struct IterFrame<'a, V> { - /// The current depth in the multi-indexed structure - /// (0 for the root, 1 for the first level, etc.) - depth: usize, - - /// The current node being processed - current_node: &'a Node, +// --- Iterator --- - /// The range of indices left to iterate over in the current node +/// A stack frame in the depth-first traversal of a [`KdTrie`]. +/// +/// Each frame records the current node, its depth in the trie (i.e. which coordinate dimension it +/// indexes), and the remaining range of indices to visit at this node. +struct IterFrame { + depth: usize, + current_node: R, range: Range, } -impl<'a, V> IterFrame<'a, V> { - /// Creates the initial iteration frame. - fn new(dimensions: usize, root: &'a Node) -> Self { - // Safety: This function is only called by KdIterator::new, which is only called by the iter - // methods of KdTrie and MultiIndexed. Therefore, by definition, the number of dimensions - // can be trusted. There can not be any other caller because of the pub(self) visibility. - let root_range = if dimensions == 1 { - unsafe { root.leaf() }.range() - } else { - unsafe { root.inner() }.range() - }; - - Self { - depth: 0, - current_node: root, - range: root_range, - } - } -} - -/// Trait for managing coordinates during iteration -trait Coordinates { - fn set_coord(&mut self, depth: usize, value: i32); - fn truncate_to(&mut self, depth: usize); - fn get(&self) -> Self; -} - -/// Iterator implementation for multi-dimensional structures +/// A depth-first iterator over a [`KdTrie`], generic over: +/// +/// - `R: NodeRef` — the node handle type, either shared (`&Node`) or exclusive +/// (`NodePtrMut<'_, V>`), determining whether values are yielded as `&V` or `&mut V`. +/// - `C: Coordinates` — the coordinate accumulator, either `[i32; K]` (fixed-size, for +/// [`MultiIndexed`]) or `Vec` (dynamic, for [`KdTrie`]). /// -/// This abstracts over both dynamic and fixed-size coordinates, which allows us to iterate over -/// `KdTrie`s with vector coordinates and `MultiIndexed`s with fixed-size arrays. It's important to -/// allow fixed-size arrays to be used as coordinates, as they are `Copy` and can avoid the -/// expensive `clone`s. Empirically, this gives a 3x speedup. -struct KdIterator<'a, V, C> { +/// The iterator walks the trie in lexicographic order of coordinates, yielding `(C, R::Value)` for +/// each stored entry. +struct KdIterator { dimensions: usize, - stack: Vec>, + stack: Vec>, coordinates: C, } -impl<'a, V, C> KdIterator<'a, V, C> { - fn new(dimensions: usize, root: &'a Node, coordinates: C) -> Self { +impl KdIterator { + fn new(dimensions: usize, root: R, coordinates: C) -> Self { + let root_range = unsafe { root.range(dimensions == 1) }; Self { dimensions, - stack: vec![IterFrame::new(dimensions, root)], + stack: vec![IterFrame { + depth: 0, + current_node: root, + range: root_range, + }], coordinates, } } } -impl KdTrie { - pub fn iter(&self) -> impl Iterator, &V)> + '_ { - let dimensions = self.dimensions(); - KdIterator::new(dimensions, self.root(), Vec::with_capacity(dimensions)) - } -} - -impl MultiIndexed { - /// Returns an iterator over all coordinate-value pairs in the array. - /// - /// The iterator yields tuples of `([i32; K], &V)` where the first element is the coordinate - /// array and the second is a reference to the value. - /// - /// # Examples - /// - /// ``` - /// use once::MultiIndexed; - /// - /// let array = MultiIndexed::<2, i32>::new(); - /// array.insert([3, 4], 10); - /// array.insert([1, 2], 20); - /// - /// let mut items: Vec<_> = array.iter().collect(); - /// - /// assert_eq!(items, vec![([1, 2], &20), ([3, 4], &10)]); - /// ``` - pub fn iter(&self) -> impl Iterator { - KdIterator::new(K, self.0.root(), [0; K]) - } -} - -impl<'a, V, C: Coordinates> Iterator for KdIterator<'a, V, C> { - type Item = (C, &'a V); +impl Iterator for KdIterator { + type Item = (C, R::Value); fn next(&mut self) -> Option { while let Some(IterFrame { @@ -112,8 +61,7 @@ impl<'a, V, C: Coordinates> Iterator for KdIterator<'a, V, C> { while let Some(idx) = range.next() { if depth == self.dimensions - 1 { // This is a leaf node, check if there's a value at this index - let current_leaf = unsafe { current_node.leaf() }; - if let Some(value) = current_leaf.get(idx) { + if let Some(value) = unsafe { current_node.value(idx) } { // Push back the remaining range for this node if !range.is_empty() { self.stack.push(IterFrame { @@ -126,35 +74,29 @@ impl<'a, V, C: Coordinates> Iterator for KdIterator<'a, V, C> { self.coordinates.set_coord(depth, idx); return Some((self.coordinates.get(), value)); } - } else { + } else if let Some(child_node) = unsafe { current_node.child(idx) } { // This is an inner node, check if there's a child at this index - let current_inner = unsafe { current_node.inner() }; - if let Some(child_node) = current_inner.get(idx) { - // Push back the remaining range for this node - if !range.is_empty() { - self.stack.push(IterFrame { - depth, - current_node, - range, - }); - } - // Add the current index to coordinates and push the child - self.coordinates.set_coord(depth, idx); - let child_range = if depth + 1 == self.dimensions - 1 { - unsafe { child_node.leaf() }.range() - } else { - unsafe { child_node.inner() }.range() - }; + // Push back the remaining range for this node + if !range.is_empty() { self.stack.push(IterFrame { - depth: depth + 1, - current_node: child_node, - range: child_range, + depth, + current_node, + range, }); - - // Go to the next iteration of the outer loop, which will process the child - break; } + + // Add the current index to coordinates and push the child + self.coordinates.set_coord(depth, idx); + let child_range = unsafe { child_node.range(depth + 1 == self.dimensions - 1) }; + self.stack.push(IterFrame { + depth: depth + 1, + current_node: child_node, + range: child_range, + }); + + // Go to the next iteration of the outer loop, which will process the child + break; } } } @@ -163,20 +105,147 @@ impl<'a, V, C: Coordinates> Iterator for KdIterator<'a, V, C> { } } +// --- NodeRef --- + +/// Abstraction over shared (`&Node`) and exclusive (`*mut Node`) node access. +/// +/// This trait allows [`KdIterator`] to be generic over the borrowing mode, so a single iterator +/// implementation drives both `iter` (shared) and `iter_mut` (exclusive). +/// +/// # Safety +/// +/// Implementations must ensure that: +/// - `range`, `child`, and `value` uphold the safety preconditions of the underlying [`Node`] +/// methods (i.e. leaf methods are only called on leaf nodes, and inner methods on inner nodes). +/// - For mutable implementations, the returned value references do not alias. +unsafe trait NodeRef: Copy { + type Value; + + /// Returns the range of indices for this node. + /// + /// # Safety + /// + /// `is_leaf` must correctly indicate whether this is a leaf node. + unsafe fn range(self, is_leaf: bool) -> Range; + + /// Returns a handle to the child node at `idx`, or `None` if no child exists. + /// + /// # Safety + /// + /// Must only be called on inner nodes. + unsafe fn child(self, idx: i32) -> Option; + + /// Returns a reference to the value at `idx`, or `None` if the slot is empty. + /// + /// # Safety + /// + /// Must only be called on leaf nodes. + unsafe fn value(self, idx: i32) -> Option; +} + +/// Shared node reference. Yields `&V` values. +unsafe impl<'a, V> NodeRef for &'a Node { + type Value = &'a V; + + unsafe fn range(self, is_leaf: bool) -> Range { + if is_leaf { + unsafe { self.leaf() }.range() + } else { + unsafe { self.inner() }.range() + } + } + + unsafe fn child(self, idx: i32) -> Option { + unsafe { self.inner().get(idx) } + } + + unsafe fn value(self, idx: i32) -> Option { + unsafe { self.leaf().get(idx) } + } +} + +/// A `Copy` wrapper around `*mut Node` that serves as the exclusive counterpart to +/// `&Node` in the [`NodeRef`] trait. +/// +/// The phantom lifetime `'a` ties the yielded `&'a mut V` references back to the original +/// `&'a mut MultiIndexed` (or `&'a mut KdTrie`), ensuring soundness. +/// +/// This is safe to use because: +/// - It is only constructed from `&mut MultiIndexed` / `&mut KdTrie`, guaranteeing exclusive access +/// to the entire tree. +/// - The tree structure ensures that nodes at different positions are disjoint in memory. +/// - The iterator yields each value at most once. +struct NodePtrMut<'a, V>(*mut Node, PhantomData<&'a mut V>); + +impl Copy for NodePtrMut<'_, V> {} + +impl Clone for NodePtrMut<'_, V> { + fn clone(&self) -> Self { + *self + } +} + +/// Exclusive node reference. Yields `&mut V` values. +unsafe impl<'a, V> NodeRef for NodePtrMut<'a, V> { + type Value = &'a mut V; + + unsafe fn range(self, is_leaf: bool) -> Range { + if is_leaf { + unsafe { (*self.0).leaf() }.range() + } else { + unsafe { (*self.0).inner() }.range() + } + } + + unsafe fn child(self, idx: i32) -> Option { + let child = unsafe { (*self.0).get_child_mut(idx) }?; + Some(Self(child as *mut Node, PhantomData)) + } + + unsafe fn value(self, idx: i32) -> Option { + unsafe { (*self.0).get_value_mut(idx) } + } +} + +// --- Coordinates --- + +/// Trait for managing coordinates during iteration. +/// +/// The iterator accumulates coordinates dimension-by-dimension as it descends. When it backtracks, +/// it calls [`truncate_to`](Coordinates::truncate_to) to discard coordinates from deeper +/// dimensions. When it yields an entry, it calls [`get`](Coordinates::get) to snapshot the current +/// coordinates. +trait Coordinates { + /// Sets the coordinate at the given `depth` (dimension index) to `value`. + fn set_coord(&mut self, depth: usize, value: i32); + + /// Discards any coordinate data beyond `depth`, preparing for backtracking. + fn truncate_to(&mut self, depth: usize); + + /// Returns a snapshot of the current coordinates. + fn get(&self) -> Self; +} + +/// Fixed-size coordinate accumulator for [`MultiIndexed`]. +/// +/// `truncate_to` is a no-op since all dimensions are always present in the array; stale values at +/// deeper indices are simply overwritten by `set_coord` before they are ever read. impl Coordinates for [i32; K] { fn set_coord(&mut self, depth: usize, value: i32) { self[depth] = value; } - fn truncate_to(&mut self, _depth: usize) { - // Array doesn't need truncation - } + fn truncate_to(&mut self, _depth: usize) {} fn get(&self) -> Self { *self } } +/// Dynamic coordinate accumulator for [`KdTrie`]. +/// +/// `set_coord` pushes a new coordinate (asserting that `depth == len`, i.e. coordinates are always +/// built in order), and `truncate_to` pops coordinates back to the given depth. impl Coordinates for Vec { fn set_coord(&mut self, depth: usize, value: i32) { assert_eq!(self.len(), depth); @@ -191,3 +260,66 @@ impl Coordinates for Vec { self.clone() } } + +// --- Public API --- + +impl KdTrie { + pub fn iter(&self) -> impl Iterator, &V)> + '_ { + let dimensions = self.dimensions(); + KdIterator::new(dimensions, self.root(), Vec::with_capacity(dimensions)) + } + + pub fn iter_mut(&mut self) -> impl Iterator, &mut V)> + '_ { + let dimensions = self.dimensions(); + let root = NodePtrMut(self.root_mut() as *mut Node, PhantomData); + KdIterator::new(dimensions, root, Vec::with_capacity(dimensions)) + } +} + +impl MultiIndexed { + /// Returns an iterator over all coordinate-value pairs in the array. + /// + /// The iterator yields `([i32; K], &V)` tuples in lexicographic order of coordinates. + /// + /// # Examples + /// + /// ``` + /// use once::MultiIndexed; + /// + /// let array = MultiIndexed::<2, i32>::new(); + /// array.insert([3, 4], 10); + /// array.insert([1, 2], 20); + /// + /// let mut items: Vec<_> = array.iter().collect(); + /// + /// assert_eq!(items, vec![([1, 2], &20), ([3, 4], &10)]); + /// ``` + pub fn iter(&self) -> impl Iterator { + KdIterator::new(K, self.0.root(), [0; K]) + } + + /// Returns a mutable iterator over all coordinate-value pairs in the array. + /// + /// The iterator yields `([i32; K], &mut V)` tuples in lexicographic order of coordinates. + /// + /// # Examples + /// + /// ``` + /// use once::MultiIndexed; + /// + /// let mut array = MultiIndexed::<2, i32>::new(); + /// array.insert([1, 2], 10); + /// array.insert([3, 4], 20); + /// + /// for (_, v) in array.iter_mut() { + /// *v *= 2; + /// } + /// + /// assert_eq!(array.get([1, 2]), Some(&20)); + /// assert_eq!(array.get([3, 4]), Some(&40)); + /// ``` + pub fn iter_mut(&mut self) -> impl Iterator { + let root = NodePtrMut(self.0.root_mut() as *mut Node, PhantomData); + KdIterator::new(K, root, [0; K]) + } +} diff --git a/ext/crates/once/src/multiindexed/kdtrie.rs b/ext/crates/once/src/multiindexed/kdtrie.rs index 007f7c4e9a..af64be33e3 100644 --- a/ext/crates/once/src/multiindexed/kdtrie.rs +++ b/ext/crates/once/src/multiindexed/kdtrie.rs @@ -211,6 +211,10 @@ impl KdTrie { pub(super) fn root(&self) -> &Node { &self.root } + + pub(super) fn root_mut(&mut self) -> &mut Node { + &mut self.root + } } impl Drop for KdTrie { diff --git a/ext/crates/once/src/multiindexed/mod.rs b/ext/crates/once/src/multiindexed/mod.rs index 24c6eb638c..ee24688579 100644 --- a/ext/crates/once/src/multiindexed/mod.rs +++ b/ext/crates/once/src/multiindexed/mod.rs @@ -418,6 +418,44 @@ mod tests { assert_eq!(items1, items2); } + #[test] + fn test_iter_mut_empty() { + let mut arr = MultiIndexed::<2, i32>::new(); + let items: Vec<_> = arr.iter_mut().collect(); + assert_eq!(items, vec![]); + } + + #[test] + fn test_iter_mut_basic() { + let mut arr = MultiIndexed::<2, i32>::new(); + arr.insert([1, 2], 10); + arr.insert([3, 4], 20); + arr.insert([-5, 6], 30); + + // Mutate all values + for (_, v) in arr.iter_mut() { + *v *= 3; + } + + assert_eq!(arr.get([1, 2]), Some(&30)); + assert_eq!(arr.get([3, 4]), Some(&60)); + assert_eq!(arr.get([-5, 6]), Some(&90)); + } + + #[test] + fn test_iter_and_iter_mut_agree() { + let mut arr = MultiIndexed::<2, i32>::new(); + arr.insert([1, 2], 10); + arr.insert([3, -4], 20); + arr.insert([-5, 6], 30); + arr.insert([0, 0], 40); + + let immutable: Vec<_> = arr.iter().map(|(c, &v)| (c, v)).collect(); + let mutable: Vec<_> = arr.iter_mut().map(|(c, &mut v)| (c, v)).collect(); + + assert_eq!(immutable, mutable); + } + #[test] fn test_requires_drop() { use std::{ @@ -649,6 +687,33 @@ mod tests { assert_eq!(tagged_coords, items); } + fn proptest_multiindexed_iter_mut_kd(coords: Vec<[i32; K]>) { + let mut arr = MultiIndexed::::new(); + let mut reference = HashMap::new(); + for (i, coord) in coords.iter().enumerate() { + if arr.try_insert(*coord, i).is_ok() { + reference.insert(*coord, i); + } + } + + // Mutate via iter_mut: double every value + for (_, v) in arr.iter_mut() { + *v *= 2; + } + + // Verify against reference + for (coord, expected) in &reference { + assert_eq!(arr.get(*coord), Some(&(expected * 2))); + } + + // Verify iter_mut yields same coordinates as iter (after mutation) + let mut iter_items: Vec<_> = arr.iter().map(|(c, &v)| (c, v)).collect(); + iter_items.sort(); + let mut reference_items: Vec<_> = reference.iter().map(|(c, v)| (*c, v * 2)).collect(); + reference_items.sort(); + assert_eq!(iter_items, reference_items); + } + const MAX_LEN: usize = 10_000; proptest! { @@ -671,6 +736,16 @@ mod tests { fn proptest_multiindexed_iter_3d(coords in coords_vec_strategy::<3>(MAX_LEN)) { proptest_multiindexed_iter_kd::<3>(coords); } + + #[test] + fn proptest_multiindexed_iter_mut_2d(coords in coords_vec_strategy::<2>(MAX_LEN)) { + proptest_multiindexed_iter_mut_kd::<2>(coords); + } + + #[test] + fn proptest_multiindexed_iter_mut_3d(coords in coords_vec_strategy::<3>(MAX_LEN)) { + proptest_multiindexed_iter_mut_kd::<3>(coords); + } } } From 42a51ec53e9da6b9a3064de8aa030d08153d9e3d Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Sun, 15 Mar 2026 00:11:54 -0400 Subject: [PATCH 2/2] Detect and fix SB violation --- ext/crates/once/src/grove/block.rs | 10 ++++--- ext/crates/once/src/multiindexed/iter.rs | 33 ++++++++++++++++++++++++ 2 files changed, 39 insertions(+), 4 deletions(-) diff --git a/ext/crates/once/src/grove/block.rs b/ext/crates/once/src/grove/block.rs index 7c298007a8..b40545fe0b 100644 --- a/ext/crates/once/src/grove/block.rs +++ b/ext/crates/once/src/grove/block.rs @@ -145,13 +145,15 @@ impl Block { /// This method is safe to call even if the block is uninitialized. pub(super) fn get_mut(&mut self, index: usize) -> Option<&mut T> { let len = self.len.get_by_mut(); - if len == 0 { + if index >= len { return None; } let data_ptr = self.data.get_by_mut(); - // Safety: we just observed the length to be nonzero, so the pointer is not null - let data = unsafe { std::slice::from_raw_parts_mut(data_ptr, len) }; - data.get_mut(index).and_then(|w| w.get_mut()) + // Safety: index < len, so the pointer is in bounds. We reference a single element + // rather than creating a slice over the entire block, so that the resulting `&mut` + // does not alias other elements in the same allocation. + let elem = unsafe { &mut *data_ptr.add(index) }; + elem.get_mut() } /// Return the value at the given index. diff --git a/ext/crates/once/src/multiindexed/iter.rs b/ext/crates/once/src/multiindexed/iter.rs index 198ae67966..6ed9415e7d 100644 --- a/ext/crates/once/src/multiindexed/iter.rs +++ b/ext/crates/once/src/multiindexed/iter.rs @@ -323,3 +323,36 @@ impl MultiIndexed { KdIterator::new(K, root, [0; K]) } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_iter_mut_no_aliasing() { + let mut arr = MultiIndexed::<3, i32>::new(); + arr.insert([0, 0, 0], 10); + arr.insert([0, 0, 1], 20); + arr.insert([0, 1, 0], 30); + arr.insert([1, 0, 0], 40); + + let mut it = arr.iter_mut(); + let (_, a) = it.next().unwrap(); + let (_, b) = it.next().unwrap(); + let (_, c) = it.next().unwrap(); + let (_, d) = it.next().unwrap(); + + // Miri detects borrow-model violations if any of the references alias, even before the + // writes below. + *a += 1; + *b += 2; + *c += 3; + *d += 4; + drop(it); + + assert_eq!(arr.get([0, 0, 0]), Some(&11)); + assert_eq!(arr.get([0, 0, 1]), Some(&22)); + assert_eq!(arr.get([0, 1, 0]), Some(&33)); + assert_eq!(arr.get([1, 0, 0]), Some(&44)); + } +}