Skip to main content

commonware_storage/index/partitioned/ordered/
mod.rs

1//! A partitioned index that stores each partition as sorted struct-of-arrays (see the
2//! `self::partition` module).
3//!
4//! The first `P` bytes of the (untranslated) key select a partition; the translator maps the
5//! remaining bytes to the partition-local key. Because the partitions are ordered by prefix and each
6//! partition's entries are sorted by translated key, this index is inherently ordered. It trades
7//! lookup/insert speed for memory density at scale; the unordered variant ([`super::unordered`])
8//! uses hash sub-indices instead and is faster when ordering is not required.
9//!
10//! # Spilling over-full partitions
11//!
12//! Each sorted-array insert is an O(occupancy) memmove, so a partition that grows large makes
13//! inserts expensive. When a partition's array reaches `SPILL_THRESHOLD` entries it converts to a
14//! `BTreeMap` (the `spilled` field) -- a supported alternate representation whose insert, lookup,
15//! and traversal are O(log occupancy). A partition reaches that size two ways:
16//!
17//! - *Adversarial grinding.* An order-preserving translator cannot randomize keys (that would break
18//!   the ordering), so an attacker can grind the key suffix to flood one partition with distinct
19//!   translated keys. Spilling bounds flooding M keys from O(M^2) to O(M log M).
20//! - *Honest high-occupancy growth at low `P`.* With few partitions a uniform workload fills them: a
21//!   `P=1` index (256 partitions) is guaranteed to spill once it holds more than 256*511 = 130,816
22//!   entries, `P=2` past ~33M, while `P=3`'s 16.8M partitions push this past ~8.5B (so P=3 is
23//!   effectively unreachable under honest load).
24//!
25//! A partition also fills when a single key collects many values -- keys that collide on the full
26//! prefix, or repeated inserts of one key. The spill covers this too: it triggers on the total
27//! value count, so a single over-full key still converts the partition and keeps inserts for the
28//! partition's other keys cheap. Values append to the end of a key's run. In the single-key case
29//! this makes inline inserts append-only, and after spilling they remain append-only in the run's
30//! `Vec`. Other inline inserts may shift later key runs, but the spill threshold bounds this cost.
31//! What spilling cannot bound is how many values one key holds, and a lookup must scan all of
32//! them: a key with `M` values costs O(M) per lookup. Every index that resolves collisions pays
33//! this scan (the flat `crate::index::ordered::Index` included); `M` stays near 1 only when the
34//! indexed `P + N`-byte prefix is well-distributed, so use enough prefix bytes and high-entropy
35//! keys.
36//!
37//! A caller-held cursor can temporarily grow an inline partition to or past the spill threshold.
38//! The next index mutation of that partition spills it before access. `insert_and_retain` performs
39//! the check after releasing its internal cursor.
40
41mod cursor;
42mod partition;
43
44pub use self::cursor::Cursor;
45use self::partition::Partition;
46#[commonware_macros::stability(ALPHA)]
47use crate::index::partitioned::{PartitionRange, Partitioned};
48use crate::{
49    index::{
50        Cursor as CursorTrait, Factory, Ordered, Unordered,
51        partitioned::partition_index_and_sub_key,
52    },
53    translator::Translator,
54};
55use commonware_runtime::{
56    Metrics,
57    telemetry::metrics::{Counter, Gauge, MetricsExt as _},
58};
59use std::{
60    collections::{BTreeMap, HashMap, btree_map, hash_map},
61    ops::Bound,
62};
63
64/// Sorted-array length at which a partition converts to a `BTreeMap`, bounding the O(occupancy)
65/// insert memmove to O(log occupancy). A partition reaches this from adversarial distinct-key
66/// grinding or from honest growth once partitions fill: a spill is guaranteed past 256*511 = 130,816
67/// entries at `P=1`, past ~33M at `P=2`, and only past ~8.5B at `P=3` (so P=3 effectively never
68/// spills under honest load). See the module docs.
69const SPILL_THRESHOLD: usize = 512;
70
71/// A partitioned index storing each partition as sorted struct-of-arrays, spilling an over-full
72/// partition to a `BTreeMap` to bound its O(occupancy) insert cost (see `spilled` and the module
73/// docs).
74pub struct Index<T: Translator, V: Send + Sync, const P: usize> {
75    /// Translates the prefix-stripped key bytes into a partition-local key.
76    translator: T,
77
78    /// The partitions, indexed by a key's partition index. A full index holds all `2^(8*P)`, while
79    /// the index inside a [RangeIndex] build worker holds only the worker's range (addressed by
80    /// local slot). Each stores its translated keys and values as sorted arrays (the inline
81    /// representation), though an emptied partition may instead have spilled (see `spilled`).
82    partitions: Box<[Partition<T::Key, V>]>,
83
84    /// Partitions that have spilled out of their sorted arrays (reached `SPILL_THRESHOLD` entries),
85    /// keyed by partition index; each maps translated keys to their value runs. Empty until a
86    /// partition fills, whether from honest growth at low `P` or adversarial grinding.
87    spilled: HashMap<usize, BTreeMap<T::Key, Vec<V>>>,
88
89    /// Sorted-array length at which a partition spills to `spilled`; [SPILL_THRESHOLD] in
90    /// production, lowered by tests to exercise spilling cheaply.
91    threshold: usize,
92
93    /// Metric: distinct translated keys currently held across all partitions.
94    keys: Gauge,
95
96    /// Metric: stored values currently held across all partitions.
97    items: Gauge,
98
99    /// Metric: cumulative values removed (via `remove`, cursor `delete`, or `retain`).
100    pruned: Counter,
101
102    /// Metric: cumulative partitions spilled to the side-table. Emptied partitions that de-spill
103    /// (rare) are not subtracted. Build workers hold clones of the handle, so their spills count
104    /// here live.
105    spills: Counter,
106}
107
108impl<T: Translator, V: Send + Sync, const P: usize> Index<T, V, P> {
109    /// Create a new [Index] with the given metrics context and translator.
110    pub fn new(ctx: impl Metrics, translator: T) -> Self {
111        const {
112            assert!(P > 0 && P <= 3, "P must be in 1..=3");
113        }
114        let count = 1usize << (P * 8);
115        let partitions = (0..count)
116            .map(|_| Partition::default())
117            .collect::<Vec<_>>()
118            .into_boxed_slice();
119        Self {
120            translator,
121            partitions,
122            spilled: HashMap::new(),
123            threshold: SPILL_THRESHOLD,
124            keys: ctx.gauge("keys", "Number of translated keys in the index"),
125            items: ctx.gauge("items", "Number of items in the index"),
126            pruned: ctx.counter("pruned", "Number of items pruned"),
127            spills: ctx.counter("spills", "Number of partitions spilled to the side-table"),
128        }
129    }
130
131    /// Create a new [Index] with an explicit spill threshold so tests can exercise spilling without
132    /// inserting [SPILL_THRESHOLD] keys. The threshold must be at least 1: `maybe_spill` relies on
133    /// an already-spilled partition's empty inline array staying strictly below the threshold.
134    #[cfg(test)]
135    pub(crate) fn with_threshold(ctx: impl Metrics, translator: T, threshold: usize) -> Self {
136        assert!(threshold > 0, "spill threshold must be at least 1");
137        let mut index = Self::new(ctx, translator);
138        index.threshold = threshold;
139        index
140    }
141
142    /// Visit every value held across all partitions (inline and spilled), in unspecified order.
143    #[commonware_macros::stability(ALPHA)]
144    fn for_each_value(&self, mut f: impl FnMut(&V)) {
145        for (p, partition) in self.partitions.iter().enumerate() {
146            for v in partition.values_iter() {
147                f(v);
148            }
149            if let Some(inner) = self.spilled_partition(p) {
150                for vals in inner.values() {
151                    for v in vals {
152                        f(v);
153                    }
154                }
155            }
156        }
157    }
158
159    /// Spill partition `i` to the side-table if its sorted array has reached the threshold.
160    fn maybe_spill(&mut self, i: usize) {
161        if self.partitions[i].len() < self.threshold {
162            return;
163        }
164        let inner: BTreeMap<T::Key, Vec<V>> = self.partitions[i].drain_runs().into_iter().collect();
165        self.spilled.insert(i, inner);
166        self.spills.inc();
167    }
168
169    /// The `BTreeMap` of spilled partition `i`, or `None` if `i` has not spilled. The empty-map
170    /// check skips hashing `i` in the common case where no partition has ever spilled.
171    fn spilled_partition(&self, i: usize) -> Option<&BTreeMap<T::Key, Vec<V>>> {
172        if self.spilled.is_empty() {
173            return None;
174        }
175        self.spilled.get(&i)
176    }
177
178    /// The values for translated key `k` in partition `i` (empty if absent), from whichever
179    /// representation the partition currently uses.
180    fn partition_values(&self, i: usize, k: &T::Key) -> &[V] {
181        if !self.partitions[i].is_empty() {
182            return self.partitions[i].values(k);
183        }
184        self.spilled_partition(i)
185            .and_then(|inner| inner.get(k))
186            .map_or(&[], Vec::as_slice)
187    }
188
189    /// Values of the smallest key in partition `i` (None if the partition is empty).
190    fn partition_first(&self, i: usize) -> Option<&[V]> {
191        self.partitions[i].first_values().or_else(|| {
192            self.spilled_partition(i)?
193                .first_key_value()
194                .map(|(_, v)| v.as_slice())
195        })
196    }
197
198    /// Values of the largest key in partition `i` (None if the partition is empty).
199    fn partition_last(&self, i: usize) -> Option<&[V]> {
200        self.partitions[i].last_values().or_else(|| {
201            self.spilled_partition(i)?
202                .last_key_value()
203                .map(|(_, v)| v.as_slice())
204        })
205    }
206
207    /// Whether the index currently holds no keys.
208    fn is_empty(&self) -> bool {
209        self.keys.get() == 0
210    }
211
212    /// Values of the smallest key strictly greater than `k` in partition `i`.
213    fn partition_next_after(&self, i: usize, k: &T::Key) -> Option<&[V]> {
214        self.partitions[i].next_values_after(k).or_else(|| {
215            self.spilled_partition(i)?
216                .range((Bound::Excluded(*k), Bound::Unbounded))
217                .next()
218                .map(|(_, v)| v.as_slice())
219        })
220    }
221
222    /// Values of the largest key strictly less than `k` in partition `i`.
223    fn partition_prev_before(&self, i: usize, k: &T::Key) -> Option<&[V]> {
224        self.partitions[i].prev_values_before(k).or_else(|| {
225            self.spilled_partition(i)?
226                .range((Bound::Unbounded, Bound::Excluded(*k)))
227                .next_back()
228                .map(|(_, v)| v.as_slice())
229        })
230    }
231
232    /// Number of partitions currently spilled to the side-table.
233    #[cfg(test)]
234    pub(crate) fn spilled_count(&self) -> usize {
235        self.spilled.len()
236    }
237
238    /// Cumulative value of the `spills` metric.
239    #[cfg(test)]
240    fn spills(&self) -> usize {
241        self.spills.get() as usize
242    }
243
244    /// Mutable cursor over the values of sub-key `sub` in the partition at local slot `i`, if the
245    /// key exists (see [Unordered::get_mut]).
246    fn get_mut_slot(&mut self, i: usize, sub: &[u8]) -> Option<Cursor<'_, T::Key, V>> {
247        let k = self.translator.transform(sub);
248        self.maybe_spill(i);
249        if !self.partitions[i].is_empty() {
250            let run = self.partitions[i].run_range(&k);
251            if run.is_empty() {
252                return None;
253            }
254            return Some(Cursor::soa(
255                &mut self.partitions[i],
256                k,
257                run,
258                &self.keys,
259                &self.items,
260                &self.pruned,
261            ));
262        }
263
264        // Hand out a spilled cursor if the partition has spilled and holds `k`.
265        if self
266            .spilled_partition(i)
267            .is_some_and(|inner| inner.contains_key(&k))
268        {
269            return Some(Cursor::spilled(
270                &mut self.spilled,
271                i,
272                k,
273                &self.keys,
274                &self.items,
275                &self.pruned,
276            ));
277        }
278
279        // Partition is genuinely empty.
280        None
281    }
282
283    /// Mutable cursor over the values of sub-key `sub` in the partition at local slot `i` if the
284    /// key exists, otherwise inserts `value` for it and returns `None` (see
285    /// [Unordered::get_mut_or_insert]).
286    fn get_mut_or_insert_slot(
287        &mut self,
288        i: usize,
289        sub: &[u8],
290        value: V,
291    ) -> Option<Cursor<'_, T::Key, V>> {
292        let k = self.translator.transform(sub);
293        self.maybe_spill(i);
294        if !self.partitions[i].is_empty() {
295            let run = self.partitions[i].run_range(&k);
296            if !run.is_empty() {
297                return Some(Cursor::soa(
298                    &mut self.partitions[i],
299                    k,
300                    run,
301                    &self.keys,
302                    &self.items,
303                    &self.pruned,
304                ));
305            }
306            self.partitions[i].insert_at(run.end, k, value);
307            self.keys.inc();
308            self.items.inc();
309            self.maybe_spill(i);
310            return None;
311        }
312
313        // Partition i is empty. If it's because it has spilled, serve or create the key in its
314        // `BTreeMap`.
315        if let Some(inner) = self.spilled_partition(i) {
316            if inner.contains_key(&k) {
317                return Some(Cursor::spilled(
318                    &mut self.spilled,
319                    i,
320                    k,
321                    &self.keys,
322                    &self.items,
323                    &self.pruned,
324                ));
325            }
326            self.spilled.get_mut(&i).unwrap().insert(k, vec![value]);
327            self.keys.inc();
328            self.items.inc();
329            return None;
330        }
331
332        // Partition i is genuinely empty: start a fresh sorted array.
333        self.partitions[i].insert_at(0, k, value);
334        self.keys.inc();
335        self.items.inc();
336        self.maybe_spill(i);
337
338        None
339    }
340}
341
342#[commonware_macros::stability(ALPHA)]
343impl<T: Translator, V: Send + Sync + 'static, const P: usize> Partitioned for Index<T, V, P> {
344    type Range = RangeIndex<T, V, P>;
345
346    fn partition_count(&self) -> usize {
347        self.partitions.len()
348    }
349
350    fn partition_of(key: &[u8]) -> usize {
351        partition_index_and_sub_key::<P>(key).0
352    }
353
354    /// The range matches this index's translator and spill threshold. It allocates only `count`
355    /// partition slots, so per-worker memory is the range rather than the full `2^(8*P)`, which
356    /// is what makes a large `P` affordable.
357    fn new_range(&self, offset: usize, count: usize) -> RangeIndex<T, V, P> {
358        let partitions = (0..count)
359            .map(|_| Partition::default())
360            .collect::<Vec<_>>()
361            .into_boxed_slice();
362        RangeIndex {
363            index: Self {
364                translator: self.translator.clone(),
365                partitions,
366                spilled: HashMap::new(),
367                threshold: self.threshold,
368                keys: self.keys.clone(),
369                items: self.items.clone(),
370                pruned: self.pruned.clone(),
371                spills: self.spills.clone(),
372            },
373            offset,
374        }
375    }
376
377    /// Moves the worker's partitions and spilled entries wholesale. Metrics need no adjustment,
378    /// since the worker updated this index's handles directly.
379    fn install_range(&mut self, mut worker: RangeIndex<T, V, P>) {
380        let lo = worker.offset;
381        let len = worker.index.partitions.len();
382
383        // Probe the spilled side-table by its (usually empty) key set rather than once per slot.
384        assert!(
385            self.spilled.keys().all(|&p| p < lo || p >= lo + len),
386            "install target range must be empty"
387        );
388        for (local, partition) in worker.index.partitions.iter_mut().enumerate() {
389            let global = lo + local;
390            assert!(
391                self.partitions[global].is_empty(),
392                "install target range must be empty"
393            );
394            self.partitions[global] = std::mem::take(partition);
395        }
396
397        // Drain only the partitions that actually spilled (remapping local -> global), rather than
398        // probing every slot in the range.
399        for (local, inner) in worker.index.spilled.drain() {
400            self.spilled.insert(lo + local, inner);
401        }
402    }
403}
404
405/// A restricted view of an [Index] covering only a contiguous range of partitions, held by one
406/// parallel snapshot-build worker (created by [Index::new_range], folded back into a full index by
407/// [Index::install_range]). It exposes only the cursor operations, which map a key's global
408/// partition index to the worker's local slot. The other [Unordered] operations index partitions
409/// globally and are deliberately unavailable, so they cannot be miscalled on a worker.
410#[commonware_macros::stability(ALPHA)]
411pub(crate) struct RangeIndex<T: Translator, V: Send + Sync, const P: usize> {
412    /// The worker's partitions, addressed by local slot (`global partition - offset`).
413    index: Index<T, V, P>,
414
415    /// The first global partition index the worker covers.
416    offset: usize,
417}
418
419#[commonware_macros::stability(ALPHA)]
420impl<T: Translator, V: Send + Sync, const P: usize> PartitionRange for RangeIndex<T, V, P> {
421    type Value = V;
422    type Cursor<'a>
423        = Cursor<'a, T::Key, V>
424    where
425        Self: 'a;
426
427    fn get_mut(&mut self, key: &[u8]) -> Option<Cursor<'_, T::Key, V>> {
428        let (i, sub) = partition_index_and_sub_key::<P>(key);
429        self.index.get_mut_slot(i - self.offset, sub)
430    }
431
432    fn get_mut_or_insert(&mut self, key: &[u8], value: V) -> Option<Cursor<'_, T::Key, V>> {
433        let (i, sub) = partition_index_and_sub_key::<P>(key);
434        self.index
435            .get_mut_or_insert_slot(i - self.offset, sub, value)
436    }
437
438    fn for_each_value(&self, f: impl FnMut(&V)) {
439        self.index.for_each_value(f);
440    }
441}
442
443impl<T: Translator, V: Send + Sync, const P: usize> Factory for Index<T, V, P> {
444    type Translator = T;
445
446    fn new(ctx: impl Metrics, translator: T) -> Self {
447        Self::new(ctx, translator)
448    }
449}
450
451impl<T: Translator, V: Send + Sync, const P: usize> Unordered for Index<T, V, P> {
452    type Value = V;
453    type Cursor<'a>
454        = Cursor<'a, T::Key, V>
455    where
456        Self: 'a;
457
458    fn get<'a>(&'a self, key: &[u8]) -> impl Iterator<Item = &'a V> + Send + 'a
459    where
460        V: 'a,
461    {
462        let (i, sub) = partition_index_and_sub_key::<P>(key);
463        let k = self.translator.transform(sub);
464        self.partition_values(i, &k).iter()
465    }
466
467    fn get_many<'a, K: AsRef<[u8]>>(&'a self, keys: &[K], mut visit: impl FnMut(usize, &'a V))
468    where
469        V: 'a,
470    {
471        // Probe in (partition, translated-key) order so consecutive probes hit the same partition
472        // (one region of the 2^(8*P)-entry partition array) and the same value run within it,
473        // instead of scattering across partitions in input order.
474        let mut order: Vec<(usize, T::Key, usize)> = keys
475            .iter()
476            .enumerate()
477            .map(|(key_idx, key)| {
478                let (partition, sub) = partition_index_and_sub_key::<P>(key.as_ref());
479                (partition, self.translator.transform(sub), key_idx)
480            })
481            .collect();
482        order.sort_unstable();
483        for (partition, translated, key_idx) in order {
484            for value in self.partition_values(partition, &translated) {
485                visit(key_idx, value);
486            }
487        }
488    }
489
490    fn get_mut<'a>(&'a mut self, key: &[u8]) -> Option<Self::Cursor<'a>> {
491        let (i, sub) = partition_index_and_sub_key::<P>(key);
492        self.get_mut_slot(i, sub)
493    }
494
495    fn get_mut_or_insert<'a>(
496        &'a mut self,
497        key: &[u8],
498        value: Self::Value,
499    ) -> Option<Self::Cursor<'a>> {
500        let (i, sub) = partition_index_and_sub_key::<P>(key);
501        self.get_mut_or_insert_slot(i, sub, value)
502    }
503
504    fn insert(&mut self, key: &[u8], value: Self::Value) {
505        let (i, sub) = partition_index_and_sub_key::<P>(key);
506        let k = self.translator.transform(sub);
507        self.maybe_spill(i);
508        if !self.partitions[i].is_empty() {
509            let run = self.partitions[i].run_range(&k);
510            let new_key = run.is_empty();
511            self.partitions[i].insert_at(run.end, k, value);
512            self.items.inc();
513            if new_key {
514                self.keys.inc();
515            }
516            self.maybe_spill(i);
517            return;
518        }
519
520        // Route into the spilled partition's `BTreeMap`.
521        if !self.spilled.is_empty()
522            && let hash_map::Entry::Occupied(mut partition) = self.spilled.entry(i)
523        {
524            match partition.get_mut().entry(k) {
525                btree_map::Entry::Occupied(mut run) => run.get_mut().push(value),
526                btree_map::Entry::Vacant(run) => {
527                    run.insert(vec![value]);
528                    self.keys.inc();
529                }
530            }
531            self.items.inc();
532            return;
533        }
534
535        // Genuinely empty partition: start a fresh sorted array.
536        self.partitions[i].insert_at(0, k, value);
537        self.items.inc();
538        self.keys.inc();
539        self.maybe_spill(i);
540    }
541
542    fn insert_and_retain(
543        &mut self,
544        key: &[u8],
545        value: Self::Value,
546        should_retain: impl Fn(&Self::Value) -> bool,
547    ) {
548        let (i, _) = partition_index_and_sub_key::<P>(key);
549        if let Some(mut cursor) = self.get_mut(key) {
550            cursor.retain(&should_retain);
551            if should_retain(&value) {
552                cursor.insert(value);
553            }
554        } else if should_retain(&value) {
555            self.insert(key, value);
556        }
557        self.maybe_spill(i);
558    }
559
560    fn remove(&mut self, key: &[u8]) {
561        let (i, sub) = partition_index_and_sub_key::<P>(key);
562        let k = self.translator.transform(sub);
563        self.maybe_spill(i);
564        if !self.partitions[i].is_empty() {
565            let run = self.partitions[i].run_range(&k);
566            if run.is_empty() {
567                return;
568            }
569            let n = run.len();
570            self.partitions[i].remove_run(run);
571            self.keys.dec();
572            self.items.dec_by(n as i64);
573            self.pruned.inc_by(n as u64);
574            return;
575        }
576        // Partition i is empty here; if spilled, remove from its `BTreeMap` (and drop the
577        // partition entry, reverting to an empty sorted array, once its last key is gone).
578        if !self.spilled.is_empty()
579            && let hash_map::Entry::Occupied(mut partition) = self.spilled.entry(i)
580            && let Some(vals) = partition.get_mut().remove(&k)
581        {
582            let n = vals.len();
583            self.keys.dec();
584            self.items.dec_by(n as i64);
585            self.pruned.inc_by(n as u64);
586            if partition.get().is_empty() {
587                partition.remove();
588            }
589        }
590    }
591
592    #[cfg(test)]
593    fn keys(&self) -> usize {
594        self.keys.get() as usize
595    }
596
597    #[cfg(test)]
598    fn items(&self) -> usize {
599        self.items.get() as usize
600    }
601
602    #[cfg(test)]
603    fn pruned(&self) -> usize {
604        self.pruned.get() as usize
605    }
606}
607
608impl<T: Translator, V: Send + Sync, const P: usize> Ordered for Index<T, V, P> {
609    fn prev_translated_key<'a>(
610        &'a self,
611        key: &[u8],
612    ) -> Option<(impl Iterator<Item = &'a V> + Send + 'a, bool)>
613    where
614        V: 'a,
615    {
616        // Skip the all-partitions scan when there is nothing to find.
617        if self.is_empty() {
618            return None;
619        }
620
621        // The largest translated key strictly less than `k`: within the partition first, then the
622        // last key of the nearest lower partition, else cycle to the global last key.
623        let (i, sub) = partition_index_and_sub_key::<P>(key);
624        let k = self.translator.transform(sub);
625        if let Some(vals) = self.partition_prev_before(i, &k) {
626            return Some((vals.iter(), false));
627        }
628        for p in (0..i).rev() {
629            if let Some(vals) = self.partition_last(p) {
630                return Some((vals.iter(), false));
631            }
632        }
633        for p in (0..self.partitions.len()).rev() {
634            if let Some(vals) = self.partition_last(p) {
635                return Some((vals.iter(), true));
636            }
637        }
638        None
639    }
640
641    fn next_translated_key<'a>(
642        &'a self,
643        key: &[u8],
644    ) -> Option<(impl Iterator<Item = &'a V> + Send + 'a, bool)>
645    where
646        V: 'a,
647    {
648        // Skip the all-partitions scan when there is nothing to find.
649        if self.is_empty() {
650            return None;
651        }
652
653        // The smallest translated key strictly greater than `k`: within the partition first, then
654        // the first key of the nearest higher partition, else cycle to the global first key.
655        let (i, sub) = partition_index_and_sub_key::<P>(key);
656        let k = self.translator.transform(sub);
657        if let Some(vals) = self.partition_next_after(i, &k) {
658            return Some((vals.iter(), false));
659        }
660        for p in i + 1..self.partitions.len() {
661            if let Some(vals) = self.partition_first(p) {
662                return Some((vals.iter(), false));
663            }
664        }
665        for p in 0..self.partitions.len() {
666            if let Some(vals) = self.partition_first(p) {
667                return Some((vals.iter(), true));
668            }
669        }
670        None
671    }
672
673    fn first_translated_key<'a>(&'a self) -> Option<impl Iterator<Item = &'a V> + Send + 'a>
674    where
675        V: 'a,
676    {
677        // Skip the all-partitions scan when there is nothing to find.
678        if self.is_empty() {
679            return None;
680        }
681
682        // Scan partitions in ascending order for the global first key.
683        for p in 0..self.partitions.len() {
684            if let Some(vals) = self.partition_first(p) {
685                return Some(vals.iter());
686            }
687        }
688        None
689    }
690
691    fn last_translated_key<'a>(&'a self) -> Option<impl Iterator<Item = &'a V> + Send + 'a>
692    where
693        V: 'a,
694    {
695        // Skip the all-partitions scan when there is nothing to find.
696        if self.is_empty() {
697            return None;
698        }
699
700        // Scan partitions in descending order for the global last key.
701        for p in (0..self.partitions.len()).rev() {
702            if let Some(vals) = self.partition_last(p) {
703                return Some(vals.iter());
704            }
705        }
706        None
707    }
708}
709
710#[cfg(test)]
711mod tests {
712    use super::*;
713    use crate::translator::OneCap;
714    use commonware_formatting::hex;
715    use commonware_macros::test_traced;
716    use commonware_runtime::{Runner, Supervisor as _, deterministic};
717
718    fn new_index(context: deterministic::Context) -> Index<OneCap, u64, 1> {
719        Index::new(context, OneCap)
720    }
721
722    /// Index with a tiny spill threshold: a partition spills once it holds two entries. With
723    /// `OneCap` + P=1 the key byte selects the partition and the next byte is the translated key,
724    /// so keys sharing a first byte land in one partition.
725    fn new_index_spilling(context: deterministic::Context) -> Index<OneCap, u64, 1> {
726        Index::with_threshold(context, OneCap, 2)
727    }
728
729    #[test_traced]
730    fn test_empty_and_sparse_nav() {
731        deterministic::Runner::default().start(|context| async move {
732            let mut index = new_index(context);
733
734            // Empty index: every ordered navigation returns None (via the empty fast path, without
735            // scanning all partitions).
736            assert!(index.first_translated_key().is_none());
737            assert!(index.last_translated_key().is_none());
738            assert!(index.prev_translated_key(&[0x80, 0x00]).is_none());
739            assert!(index.next_translated_key(&[0x80, 0x00]).is_none());
740
741            // Two keys in widely separated partitions (0x05 and 0xF0): neighbor scans must still
742            // cross the large empty gap between them.
743            index.insert(&[0x05, 0x01], 1);
744            index.insert(&[0xF0, 0x02], 2);
745            assert_eq!(index.keys(), 2);
746            assert_eq!(
747                index
748                    .first_translated_key()
749                    .unwrap()
750                    .copied()
751                    .collect::<Vec<_>>(),
752                vec![1]
753            );
754            assert_eq!(
755                index
756                    .last_translated_key()
757                    .unwrap()
758                    .copied()
759                    .collect::<Vec<_>>(),
760                vec![2]
761            );
762
763            // Forward across the gap, then wrap from the global last key.
764            let (it, wrapped) = index.next_translated_key(&[0x05, 0x01]).unwrap();
765            assert_eq!((it.copied().collect::<Vec<_>>(), wrapped), (vec![2], false));
766            let (it, wrapped) = index.next_translated_key(&[0xF0, 0x02]).unwrap();
767            assert_eq!((it.copied().collect::<Vec<_>>(), wrapped), (vec![1], true));
768
769            // Backward across the gap, then wrap from the global first key.
770            let (it, wrapped) = index.prev_translated_key(&[0xF0, 0x02]).unwrap();
771            assert_eq!((it.copied().collect::<Vec<_>>(), wrapped), (vec![1], false));
772            let (it, wrapped) = index.prev_translated_key(&[0x05, 0x01]).unwrap();
773            assert_eq!((it.copied().collect::<Vec<_>>(), wrapped), (vec![2], true));
774
775            // A query landing in an empty partition between the two finds both neighbors.
776            let (it, wrapped) = index.prev_translated_key(&[0x80, 0x00]).unwrap();
777            assert_eq!((it.copied().collect::<Vec<_>>(), wrapped), (vec![1], false));
778            let (it, wrapped) = index.next_translated_key(&[0x80, 0x00]).unwrap();
779            assert_eq!((it.copied().collect::<Vec<_>>(), wrapped), (vec![2], false));
780
781            // Removing all keys returns to the empty fast path.
782            index.remove(&[0x05, 0x01]);
783            index.remove(&[0xF0, 0x02]);
784            assert_eq!(index.keys(), 0);
785            assert!(index.prev_translated_key(&[0x80, 0x00]).is_none());
786            assert!(index.next_translated_key(&[0x80, 0x00]).is_none());
787        });
788    }
789
790    #[test_traced]
791    fn test_spill_transition() {
792        deterministic::Runner::default().start(|context| async move {
793            let mut index = new_index_spilling(context);
794            // Distinct translated keys in one partition (prefix 0x10).
795            index.insert(&[0x10, 0x01], 1);
796            assert_eq!(index.spilled_count(), 0);
797            index.insert(&[0x10, 0x02], 2); // second entry crosses the threshold -> spills
798            assert_eq!(index.spilled_count(), 1);
799            index.insert(&[0x10, 0x03], 3); // routed straight into the spilled map
800            assert_eq!(index.spilled_count(), 1);
801            assert_eq!(index.keys(), 3);
802            assert_eq!(index.items(), 3);
803
804            // Values are served correctly from the spilled representation in append order.
805            assert_eq!(
806                index.get(&[0x10, 0x01]).copied().collect::<Vec<_>>(),
807                vec![1]
808            );
809            index.insert(&[0x10, 0x02], 22);
810            assert_eq!(
811                index.get(&[0x10, 0x02]).copied().collect::<Vec<_>>(),
812                vec![2, 22]
813            );
814            assert_eq!(index.items(), 4);
815
816            // A different prefix lands in its own (still inline) partition.
817            index.insert(&[0x20, 0x05], 5);
818            assert_eq!(index.spilled_count(), 1);
819            assert_eq!(
820                index.get(&[0x20, 0x05]).copied().collect::<Vec<_>>(),
821                vec![5]
822            );
823        });
824    }
825
826    #[test_traced]
827    fn test_spill_after_cursor_growth() {
828        deterministic::Runner::default().start(|context| async move {
829            let mut index = new_index_spilling(context);
830            let key = [0x10, 0x01];
831
832            index.insert(&key, 1);
833            {
834                let mut cursor = index.get_mut(&key).unwrap();
835                assert_eq!(cursor.next().copied(), Some(1));
836                assert_eq!(cursor.next(), None);
837                cursor.insert(2);
838            }
839            assert_eq!(index.spilled_count(), 0);
840
841            // The next index mutation spills the over-full inline partition.
842            index.insert(&key, 3);
843            assert_eq!(index.spilled_count(), 1);
844            assert_eq!(index.get(&key).copied().collect::<Vec<_>>(), vec![1, 2, 3]);
845
846            // Mutations that own their cursor can spill as soon as they release it.
847            let other = [0x20, 0x01];
848            index.insert(&other, 4);
849            index.insert_and_retain(&other, 5, |_| true);
850            assert_eq!(index.spilled_count(), 2);
851            assert_eq!(index.get(&other).copied().collect::<Vec<_>>(), vec![4, 5]);
852
853            // `get_mut` spills an over-full partition before handing out a cursor, so the cursor
854            // serves the spilled representation.
855            let third = [0x30, 0x01];
856            index.insert(&third, 6);
857            {
858                let mut cursor = index.get_mut(&third).unwrap();
859                assert_eq!(cursor.next().copied(), Some(6));
860                assert_eq!(cursor.next(), None);
861                cursor.insert(7);
862            }
863            assert_eq!(index.spilled_count(), 2);
864            {
865                let mut cursor = index.get_mut(&third).unwrap();
866                assert_eq!(cursor.next().copied(), Some(6));
867                assert_eq!(cursor.next().copied(), Some(7));
868                assert_eq!(cursor.next(), None);
869            }
870            assert_eq!(index.spilled_count(), 3);
871            assert_eq!(index.get(&third).copied().collect::<Vec<_>>(), vec![6, 7]);
872
873            // `remove` spills an over-full partition before access, even for an absent key.
874            let fourth = [0x40, 0x01];
875            index.insert(&fourth, 8);
876            {
877                let mut cursor = index.get_mut(&fourth).unwrap();
878                assert_eq!(cursor.next().copied(), Some(8));
879                assert_eq!(cursor.next(), None);
880                cursor.insert(9);
881            }
882            assert_eq!(index.spilled_count(), 3);
883            index.remove(&[0x40, 0x02]);
884            assert_eq!(index.spilled_count(), 4);
885            assert_eq!(index.get(&fourth).copied().collect::<Vec<_>>(), vec![8, 9]);
886        });
887    }
888
889    #[test_traced]
890    fn test_spill_after_get_mut_or_insert_cursor_growth() {
891        deterministic::Runner::default().start(|context| async move {
892            let mut index = new_index_spilling(context);
893            let key = [0x10, 0x01];
894
895            index.insert(&key, 1);
896            {
897                let mut cursor = index.get_mut_or_insert(&key, 2).unwrap();
898                assert_eq!(cursor.next().copied(), Some(1));
899                assert_eq!(cursor.next(), None);
900                cursor.insert(2);
901            }
902            assert_eq!(index.spilled_count(), 0);
903
904            // The next replay-style update spills before returning another collision cursor.
905            assert!(index.get_mut_or_insert(&key, 3).is_some());
906            assert_eq!(index.spilled_count(), 1);
907            assert_eq!(index.get(&key).copied().collect::<Vec<_>>(), vec![1, 2]);
908        });
909    }
910
911    #[test_traced]
912    fn test_spill_nav() {
913        deterministic::Runner::default().start(|context| async move {
914            let mut index = new_index_spilling(context);
915            // Partition 0x10: keys 0x01, 0x02 -> spills. Partition 0x20: key 0x05 -> inline.
916            // Partition 0x30: keys 0x07, 0x08 -> spills. Nav must cross spilled<->inline boundaries.
917            index.insert(&[0x10, 0x01], 1);
918            index.insert(&[0x10, 0x02], 2);
919            index.insert(&[0x20, 0x05], 5);
920            index.insert(&[0x30, 0x07], 7);
921            index.insert(&[0x30, 0x08], 8);
922            assert_eq!(index.spilled_count(), 2); // 0x10 and 0x30; 0x20 stays inline
923
924            assert_eq!(index.first_translated_key().unwrap().next(), Some(&1));
925            assert_eq!(index.last_translated_key().unwrap().next(), Some(&8));
926
927            // Within a spilled partition.
928            let (mut it, wrapped) = index.next_translated_key(&[0x10, 0x01]).unwrap();
929            assert!(!wrapped);
930            assert_eq!(it.next(), Some(&2));
931            // Spilled -> inline boundary.
932            let (mut it, wrapped) = index.next_translated_key(&[0x10, 0x02]).unwrap();
933            assert!(!wrapped);
934            assert_eq!(it.next(), Some(&5));
935            // Inline -> spilled boundary.
936            let (mut it, wrapped) = index.next_translated_key(&[0x20, 0x05]).unwrap();
937            assert!(!wrapped);
938            assert_eq!(it.next(), Some(&7));
939            // Spilled -> inline boundary, backwards.
940            let (mut it, wrapped) = index.prev_translated_key(&[0x30, 0x07]).unwrap();
941            assert!(!wrapped);
942            assert_eq!(it.next(), Some(&5));
943            // Inline -> spilled boundary, backwards.
944            let (mut it, wrapped) = index.prev_translated_key(&[0x20, 0x05]).unwrap();
945            assert!(!wrapped);
946            assert_eq!(it.next(), Some(&2));
947            // Wrap-around from the global last key.
948            let (mut it, wrapped) = index.next_translated_key(&[0x30, 0x08]).unwrap();
949            assert!(wrapped);
950            assert_eq!(it.next(), Some(&1));
951        });
952    }
953
954    #[test_traced]
955    fn test_spill_despill_on_full_drain() {
956        deterministic::Runner::default().start(|context| async move {
957            let mut index = new_index_spilling(context);
958            index.insert(&[0x10, 0x01], 1);
959            index.insert(&[0x10, 0x02], 2); // spills
960            assert_eq!(index.spilled_count(), 1);
961
962            index.remove(&[0x10, 0x01]);
963            assert_eq!(index.spilled_count(), 1); // 0x02 still present
964            index.remove(&[0x10, 0x02]);
965            assert_eq!(index.spilled_count(), 0); // last key removed -> de-spilled
966            assert_eq!(index.keys(), 0);
967
968            // The partition reverts to an inline sorted array.
969            index.insert(&[0x10, 0x09], 9);
970            assert_eq!(index.spilled_count(), 0);
971            assert_eq!(
972                index.get(&[0x10, 0x09]).copied().collect::<Vec<_>>(),
973                vec![9]
974            );
975        });
976    }
977
978    #[test_traced]
979    fn test_spill_full_lifecycle() {
980        deterministic::Runner::default().start(|context| async move {
981            let mut index = new_index_spilling(context);
982
983            // Empty.
984            assert_eq!(index.spilled_count(), 0);
985            assert_eq!(index.keys(), 0);
986            assert_eq!(index.items(), 0);
987
988            // Empty -> inline: one entry stays below the threshold (2).
989            index.insert(&[0x10, 0x01], 1);
990            assert_eq!(index.spilled_count(), 0);
991
992            // Inline -> spilled: a second distinct key crosses the threshold.
993            index.insert(&[0x10, 0x02], 2);
994            assert_eq!(index.spilled_count(), 1);
995            assert_eq!(index.spills(), 1);
996            assert_eq!(index.keys(), 2);
997            assert_eq!(index.items(), 2);
998
999            // Spilled -> empty, draining both keys through a cursor over the spilled representation
1000            // (the cursor de-spill path); the partition reverts only once its last key is gone.
1001            {
1002                let mut cursor = index.get_mut(&[0x10, 0x01]).unwrap();
1003                assert_eq!(cursor.next().copied(), Some(1));
1004                cursor.delete();
1005            }
1006            assert_eq!(index.spilled_count(), 1); // 0x02 still present
1007            {
1008                let mut cursor = index.get_mut(&[0x10, 0x02]).unwrap();
1009                assert_eq!(cursor.next().copied(), Some(2));
1010                cursor.delete();
1011            }
1012            assert_eq!(index.spilled_count(), 0); // de-spilled back to empty
1013            assert_eq!(index.spills(), 1); // cumulative: a de-spill does not decrement it
1014            assert_eq!(index.keys(), 0);
1015            assert_eq!(index.items(), 0);
1016
1017            // Empty -> inline -> spilled a second time: a de-spilled partition is fully reusable.
1018            index.insert(&[0x10, 0x03], 3);
1019            assert_eq!(index.spilled_count(), 0);
1020            index.insert(&[0x10, 0x04], 4);
1021            assert_eq!(index.spilled_count(), 1);
1022            assert_eq!(index.spills(), 2); // a fresh spill of the same partition counts again
1023            assert_eq!(
1024                index.get(&[0x10, 0x03]).copied().collect::<Vec<_>>(),
1025                vec![3]
1026            );
1027            assert_eq!(
1028                index.get(&[0x10, 0x04]).copied().collect::<Vec<_>>(),
1029                vec![4]
1030            );
1031
1032            // Spilled -> empty again, this time via `remove` (the other de-spill path).
1033            index.remove(&[0x10, 0x03]);
1034            assert_eq!(index.spilled_count(), 1); // 0x04 still present
1035            index.remove(&[0x10, 0x04]);
1036            assert_eq!(index.spilled_count(), 0);
1037            assert_eq!(index.spills(), 2); // still cumulative after a second spill + de-spill
1038            assert_eq!(index.keys(), 0);
1039            assert_eq!(index.items(), 0);
1040
1041            // Every removed value was counted once: 2 via cursor delete + 2 via remove.
1042            assert_eq!(index.pruned(), 4);
1043        });
1044    }
1045
1046    #[test_traced]
1047    fn test_spill_counts_live() {
1048        deterministic::Runner::default().start(|context| async move {
1049            // The full index that build workers install into. It spills once a partition holds
1050            // two entries.
1051            let mut full = new_index_spilling(context.child("full"));
1052            assert_eq!(full.spills(), 0);
1053
1054            // A build worker covering the whole partition range (offset 0, so the inner index's
1055            // globally-addressed methods are usable directly). It holds clones of the full
1056            // index's metric handles, so every spill event counts there as it happens, including
1057            // one whose partition later de-spills (fully drained via remove).
1058            let mut worker = full.new_range(0, full.partition_count());
1059            worker.get_mut_or_insert(&[0x10, 0x01], 1);
1060            worker.get_mut_or_insert(&[0x10, 0x02], 2); // second key in partition 0x10 -> spills
1061            worker.index.insert(&[0x20, 0x01], 3);
1062            worker.index.insert(&[0x20, 0x02], 4); // partition 0x20 spills too...
1063            worker.index.remove(&[0x20, 0x01]);
1064            worker.index.remove(&[0x20, 0x02]); // ...then fully drains, de-spilling
1065            assert_eq!(worker.index.spilled_count(), 1);
1066            assert_eq!(full.spills(), 2); // cumulative: the de-spilled partition still counts
1067
1068            // Installing moves the structures without touching the already-live counts.
1069            full.install_range(worker);
1070            assert_eq!(full.spilled_count(), 1);
1071            assert_eq!(full.spills(), 2);
1072        });
1073    }
1074
1075    #[test_traced]
1076    fn test_worker_prunes_count_live() {
1077        deterministic::Runner::default().start(|context| async move {
1078            let mut full = new_index(context.child("full"));
1079            assert_eq!(full.pruned(), 0);
1080
1081            // A worker covering the whole partition range (offset 0, so the inner index's
1082            // globally-addressed methods are usable directly). Give a key two values, then delete
1083            // both through a cursor (the same path the parallel build's deletes take). The worker
1084            // holds clones of the full index's metric handles, so the prunes count there
1085            // immediately, matching what the serial build records.
1086            let mut worker = full.new_range(0, full.partition_count());
1087            worker.index.insert(&[0x10, 0x01], 1);
1088            worker.index.insert(&[0x10, 0x01], 2);
1089            {
1090                let mut cursor = worker.get_mut(&[0x10, 0x01]).unwrap();
1091                while cursor.next().is_some() {
1092                    cursor.delete();
1093                }
1094            }
1095            assert_eq!(full.pruned(), 2);
1096
1097            // Installing moves the structures without touching the already-live counts.
1098            full.install_range(worker);
1099            assert_eq!(full.pruned(), 2);
1100        });
1101    }
1102
1103    /// A worker's value walk must visit every value it holds exactly once, whichever
1104    /// representation each partition uses (inline sorted arrays or the spilled side-table),
1105    /// since the snapshot build derives the activity bitmap and active-key counts from it.
1106    #[test_traced]
1107    fn test_range_for_each_value_visits_all_values_once() {
1108        deterministic::Runner::default().start(|context| async move {
1109            let full = new_index_spilling(context.child("full"));
1110
1111            // Two distinct keys spill partition 0x80 (threshold 2), then a translated-key
1112            // collision appends a second value to a spilled run. Partition 0x81 holds one key
1113            // and stays inline.
1114            let mut worker = full.new_range(0x80, 2);
1115            assert!(worker.get_mut_or_insert(&[0x80, 0x01], 1).is_none());
1116            assert!(worker.get_mut_or_insert(&[0x80, 0x02, 0xAA], 2).is_none());
1117            assert!(worker.get_mut_or_insert(&[0x80, 0x02, 0xBB], 3).is_some());
1118            {
1119                let mut cursor = worker.get_mut(&[0x80, 0x02, 0xBB]).unwrap();
1120                cursor.next();
1121                cursor.insert(3);
1122            }
1123            assert!(worker.get_mut_or_insert(&[0x81, 0x07], 4).is_none());
1124            assert_eq!(worker.index.spilled_count(), 1);
1125
1126            let mut seen = Vec::new();
1127            worker.for_each_value(|v| seen.push(*v));
1128            seen.sort_unstable();
1129            assert_eq!(seen, vec![1, 2, 3, 4]);
1130        });
1131    }
1132
1133    /// A worker with a nonzero offset must land its partitions AND its spilled entries at the
1134    /// global slots `offset + local` when installed. The multi-worker equivalence tests never
1135    /// spill inside a worker range (their per-partition load stays below the spill threshold), so
1136    /// this is the only coverage of the spilled remap off offset zero.
1137    #[test_traced]
1138    fn test_install_range_nonzero_offset() {
1139        deterministic::Runner::default().start(|context| async move {
1140            // The full index spills once a partition holds two entries, as does the worker
1141            // (`new_range` copies the threshold).
1142            let mut full = new_index_spilling(context.child("full"));
1143
1144            // A worker covering global partitions [0x80, 0x82): two distinct keys spill partition
1145            // 0x80, and a third key stays inline in partition 0x81.
1146            let mut worker = full.new_range(0x80, 2);
1147            worker.get_mut_or_insert(&[0x80, 0x01], 1);
1148            worker.get_mut_or_insert(&[0x80, 0x02], 2);
1149            worker.get_mut_or_insert(&[0x81, 0x07], 3);
1150            assert_eq!(worker.index.spilled_count(), 1);
1151            assert_eq!(worker.index.keys(), 3);
1152
1153            // Installing must remap both the inline partition and the spilled entry to the
1154            // worker's global range.
1155            full.install_range(worker);
1156            assert_eq!(full.spilled_count(), 1);
1157            assert_eq!(full.keys(), 3);
1158            assert_eq!(full.items(), 3);
1159            assert_eq!(
1160                full.get(&[0x80, 0x01]).copied().collect::<Vec<_>>(),
1161                vec![1]
1162            );
1163            assert_eq!(
1164                full.get(&[0x80, 0x02]).copied().collect::<Vec<_>>(),
1165                vec![2]
1166            );
1167            assert_eq!(
1168                full.get(&[0x81, 0x07]).copied().collect::<Vec<_>>(),
1169                vec![3]
1170            );
1171
1172            // Nothing may land at the worker's local slots (global partitions 0x00 and 0x01),
1173            // which is exactly where a remap that dropped the offset would file these entries.
1174            assert!(full.get(&[0x00, 0x01]).next().is_none());
1175            assert!(full.get(&[0x01, 0x07]).next().is_none());
1176        });
1177    }
1178
1179    #[test_traced]
1180    fn test_spill_get_mut_or_insert() {
1181        deterministic::Runner::default().start(|context| async move {
1182            let mut index = new_index_spilling(context);
1183            index.insert(&[0x10, 0x01], 1);
1184            index.insert(&[0x10, 0x02], 2); // second distinct key crosses the threshold -> spills
1185            assert_eq!(index.spilled_count(), 1);
1186            assert_eq!(index.keys(), 2);
1187            assert_eq!(index.items(), 2);
1188
1189            // Existing key in a spilled partition: returns a cursor over its values; the passed
1190            // value is not inserted.
1191            {
1192                let mut cursor = index.get_mut_or_insert(&[0x10, 0x01], 99).unwrap();
1193                assert_eq!(cursor.next().copied(), Some(1));
1194                assert!(cursor.next().is_none());
1195            }
1196            assert_eq!(index.keys(), 2);
1197            assert_eq!(index.items(), 2);
1198            assert_eq!(
1199                index.get(&[0x10, 0x01]).copied().collect::<Vec<_>>(),
1200                vec![1]
1201            );
1202
1203            // Absent key in a spilled partition: inserts it as a new key and returns None (the
1204            // partition stays spilled).
1205            assert!(index.get_mut_or_insert(&[0x10, 0x03], 3).is_none());
1206            assert_eq!(index.spilled_count(), 1);
1207            assert_eq!(index.keys(), 3);
1208            assert_eq!(index.items(), 3);
1209            assert_eq!(
1210                index.get(&[0x10, 0x03]).copied().collect::<Vec<_>>(),
1211                vec![3]
1212            );
1213        });
1214    }
1215
1216    #[test_traced]
1217    #[should_panic(expected = "must call Cursor::next()")]
1218    fn test_spill_cursor_delete_before_next_panics() {
1219        deterministic::Runner::default().start(|context| async move {
1220            let mut index = new_index_spilling(context);
1221            index.insert(&[0x10, 0x01], 1);
1222            index.insert(&[0x10, 0x02], 2); // spills
1223            let mut cursor = index.get_mut(&[0x10, 0x01]).unwrap(); // cursor over the spilled partition
1224            cursor.delete();
1225        });
1226    }
1227
1228    #[test_traced]
1229    fn test_soa_basic() {
1230        deterministic::Runner::default().start(|context| async move {
1231            let mut index = new_index(context);
1232            assert_eq!(index.keys(), 0);
1233
1234            let key = b"duplicate".as_slice();
1235            index.insert(key, 1);
1236            index.insert(key, 2);
1237            index.insert(key, 3);
1238            assert_eq!(index.keys(), 1);
1239            assert_eq!(index.items(), 3);
1240            assert_eq!(index.get(key).copied().collect::<Vec<_>>(), vec![1, 2, 3]);
1241
1242            {
1243                let mut cursor = index.get_mut(key).unwrap();
1244                assert_eq!(*cursor.next().unwrap(), 1);
1245                assert_eq!(*cursor.next().unwrap(), 2);
1246                assert_eq!(*cursor.next().unwrap(), 3);
1247                assert!(cursor.next().is_none());
1248            }
1249
1250            index.insert(key, 3);
1251            index.insert(key, 4);
1252            index.retain(key, |i| *i != 3);
1253            assert_eq!(index.get(key).copied().collect::<Vec<_>>(), vec![1, 2, 4]);
1254
1255            index.retain(key, |_| false);
1256            assert_eq!(
1257                index.get(key).copied().collect::<Vec<_>>(),
1258                Vec::<u64>::new()
1259            );
1260            assert_eq!(index.keys(), 0);
1261            assert!(index.get_mut(key).is_none());
1262
1263            // No-op on a missing key.
1264            index.retain(key, |_| false);
1265        });
1266    }
1267
1268    #[test_traced]
1269    fn test_soa_cursor_find() {
1270        deterministic::Runner::default().start(|context| async move {
1271            let mut index = new_index(context);
1272            let key = b"test_key";
1273            for v in [10u64, 20, 30, 40] {
1274                index.insert(key, v);
1275            }
1276
1277            {
1278                let mut cursor = index.get_mut(key).unwrap();
1279                assert!(cursor.find(|&v| v == 30));
1280                cursor.update(35);
1281            }
1282            let values: Vec<u64> = index.get(key).copied().collect();
1283            assert!(values.contains(&35) && !values.contains(&30));
1284
1285            {
1286                let mut cursor = index.get_mut(key).unwrap();
1287                assert!(!cursor.find(|&v| v == 100));
1288                assert!(cursor.next().is_none());
1289            }
1290
1291            {
1292                let mut cursor = index.get_mut(key).unwrap();
1293                assert!(cursor.find(|&v| v == 20));
1294                cursor.delete();
1295            }
1296            let values: Vec<u64> = index.get(key).copied().collect();
1297            assert!(!values.contains(&20));
1298            assert_eq!(values.len(), 3);
1299        });
1300    }
1301
1302    #[test_traced]
1303    fn test_soa_get_many_and_partitions() {
1304        deterministic::Runner::default().start(|context| async move {
1305            let mut index = new_index(context);
1306            // "ab"/"abX" share a partition+translated key; "zz" is a different partition.
1307            index.insert(b"ab", 1);
1308            index.insert(b"ab", 2);
1309            index.insert(b"abX", 3);
1310            index.insert(b"zz", 4);
1311
1312            let keys: Vec<&[u8]> = vec![b"zz", b"missing", b"ab", b"zz"];
1313            let mut visits: Vec<Vec<u64>> = vec![Vec::new(); keys.len()];
1314            index.get_many(&keys, |key_idx, value| visits[key_idx].push(*value));
1315            assert_eq!(visits[0], vec![4]);
1316            assert!(visits[1].is_empty());
1317            assert_eq!(visits[2], vec![1, 2, 3]);
1318            assert_eq!(visits[3], vec![4]);
1319        });
1320    }
1321
1322    #[test_traced]
1323    fn test_soa_insert_and_retain() {
1324        deterministic::Runner::default().start(|context| async move {
1325            let mut index = new_index(context);
1326            // Keep both: new value appends to the run.
1327            index.insert(b"k", 1u64);
1328            index.insert_and_retain(b"k", 2, |_| true);
1329            assert_eq!(index.get(b"k").copied().collect::<Vec<_>>(), vec![1, 2]);
1330
1331            // Drop the new value: no-op.
1332            index.insert_and_retain(b"k", 9, |v| *v != 9);
1333            assert_eq!(index.get(b"k").copied().collect::<Vec<_>>(), vec![1, 2]);
1334
1335            // Drop everything.
1336            index.insert_and_retain(b"k", 9, |_| false);
1337            assert!(index.get_mut(b"k").is_none());
1338            assert_eq!(index.keys(), 0);
1339
1340            // Vacant key: insert only if retained.
1341            index.insert_and_retain(b"new", 7, |_| true);
1342            assert_eq!(index.get(b"new").copied().collect::<Vec<_>>(), vec![7]);
1343            assert_eq!(index.keys(), 1);
1344        });
1345    }
1346
1347    #[test_traced]
1348    fn test_soa_remove() {
1349        deterministic::Runner::default().start(|context| async move {
1350            let mut index = new_index(context);
1351            index.insert(b"k", 1u64);
1352            index.insert(b"k", 2);
1353            index.insert(b"other", 3);
1354            assert_eq!(index.items(), 3);
1355            assert_eq!(index.keys(), 2);
1356
1357            index.remove(b"k");
1358            assert!(index.get_mut(b"k").is_none());
1359            assert_eq!(index.keys(), 1);
1360            assert_eq!(index.items(), 1);
1361            assert_eq!(index.pruned(), 2);
1362            assert_eq!(index.get(b"other").copied().collect::<Vec<_>>(), vec![3]);
1363
1364            index.remove(b"missing"); // no-op
1365            assert_eq!(index.keys(), 1);
1366        });
1367    }
1368
1369    #[test_traced]
1370    fn test_soa_ordered() {
1371        deterministic::Runner::default().start(|context| async move {
1372            let mut index = new_index(context);
1373            assert!(index.first_translated_key().is_none());
1374            assert!(index.last_translated_key().is_none());
1375            assert!(index.next_translated_key(b"key").is_none());
1376            assert!(index.prev_translated_key(b"key").is_none());
1377
1378            // With OneCap + P=1, the full key orders as (prefix byte, first sub-key byte).
1379            let k1 = &hex!("0x0b02AA"); // -> partition 0b, sub-key 02
1380            let k2 = &hex!("0x1c04CC"); // -> partition 1c, sub-key 04
1381            let k2_collides = &hex!("0x1c0411"); // same (1c, 04) as k2
1382            let k3 = &hex!("0x2d06EE"); // -> partition 2d, sub-key 06
1383            index.insert(k1, 1);
1384            index.insert(k2, 21);
1385            index.insert(k2_collides, 22);
1386            index.insert(k3, 3);
1387            assert_eq!(index.keys(), 3);
1388
1389            assert_eq!(index.first_translated_key().unwrap().next(), Some(&1));
1390            assert_eq!(index.last_translated_key().unwrap().next(), Some(&3));
1391
1392            // From before the first key: the first key, not wrapped.
1393            let (mut it, wrapped) = index.next_translated_key(&[0x00]).unwrap();
1394            assert!(!wrapped);
1395            assert_eq!(it.next(), Some(&1));
1396            assert_eq!(it.next(), None);
1397
1398            // From k1's bucket: jumps partitions to k2's collision run.
1399            let (mut it, wrapped) = index.next_translated_key(&hex!("0x0b02F2")).unwrap();
1400            assert!(!wrapped);
1401            assert_eq!(it.next(), Some(&21));
1402            assert_eq!(it.next(), Some(&22));
1403            assert_eq!(it.next(), None);
1404
1405            // From the last key: cycles to the first.
1406            let (mut it, wrapped) = index.next_translated_key(k3).unwrap();
1407            assert!(wrapped);
1408            assert_eq!(it.next(), Some(&1));
1409
1410            // From the first key going backwards: cycles to the last.
1411            let (mut it, wrapped) = index.prev_translated_key(k1).unwrap();
1412            assert!(wrapped);
1413            assert_eq!(it.next(), Some(&3));
1414
1415            // Previous bucket below 1d is 1c's collision run.
1416            let (mut it, wrapped) = index.prev_translated_key(&hex!("0x1d0102")).unwrap();
1417            assert!(!wrapped);
1418            assert_eq!(it.next(), Some(&21));
1419            assert_eq!(it.next(), Some(&22));
1420            assert_eq!(it.next(), None);
1421        });
1422    }
1423
1424    #[test_traced]
1425    fn test_soa_ordered_exhaustive_traversal() {
1426        deterministic::Runner::default().start(|context| async move {
1427            let mut index = new_index(context);
1428
1429            // A grid of (prefix, sub-key) keys spanning several partitions, including the edge
1430            // bytes 0x00/0xFF, each a distinct translated key (OneCap + P=1 orders by
1431            // (prefix, first sub-key byte)). `keys` is built in ascending order.
1432            let prefixes = [0x00u8, 0x05, 0xAA, 0xFF];
1433            let subkeys = [0x00u8, 0x80, 0xFF];
1434            let mut keys: Vec<[u8; 2]> = Vec::new();
1435            for &p in &prefixes {
1436                for &s in &subkeys {
1437                    keys.push([p, s]);
1438                }
1439            }
1440            let value_of = |k: &[u8; 2]| ((k[0] as u64) << 8) | k[1] as u64;
1441            let n = keys.len();
1442
1443            // Insert scrambled to exercise sorted-array maintenance regardless of insertion order.
1444            let mut scrambled = keys.clone();
1445            scrambled.reverse();
1446            scrambled.rotate_left(5);
1447            for k in &scrambled {
1448                index.insert(k, value_of(k));
1449            }
1450            assert_eq!(index.keys(), n);
1451
1452            assert_eq!(
1453                index.first_translated_key().unwrap().next(),
1454                Some(&value_of(&keys[0]))
1455            );
1456            assert_eq!(
1457                index.last_translated_key().unwrap().next(),
1458                Some(&value_of(&keys[n - 1]))
1459            );
1460
1461            // For every key, `next` is its successor and `prev` its predecessor, wrapping at the
1462            // ends. This walks run_starting_at / run_ending_at across every partition boundary.
1463            for i in 0..n {
1464                let next = value_of(&keys[(i + 1) % n]);
1465                let (mut it, wrapped) = index.next_translated_key(&keys[i]).unwrap();
1466                assert_eq!(wrapped, i + 1 == n, "next wrap at index {i}");
1467                assert_eq!(it.next(), Some(&next), "next at {i}");
1468                assert_eq!(it.next(), None);
1469
1470                let prev = value_of(&keys[(i + n - 1) % n]);
1471                let (mut it, wrapped) = index.prev_translated_key(&keys[i]).unwrap();
1472                assert_eq!(wrapped, i == 0, "prev wrap at index {i}");
1473                assert_eq!(it.next(), Some(&prev), "prev at {i}");
1474                assert_eq!(it.next(), None);
1475            }
1476        });
1477    }
1478}