spacetimedb_table/table_index/
mod.rs

1//! Table indexes with specialized key types.
2//!
3//! Indexes could be implemented as `MultiMap<AlgebraicValue, RowPointer>` (and once were),
4//! but that results in wasted memory and spurious comparisons and branches
5//! because the keys must always be homogeneous at a more specific type than `AlgebraicValue`.
6//!
7//! As an optimization, we hoist the enum out of the keys to the index itself.
8//! This is a sizeable improvement for integer keys,
9//! as e.g. `u64::cmp` is much faster than `AlgebraicValue::cmp`.
10//!
11//! This results in some pretty ugly code, where types that would be structs
12//! are instead enums with similar-looking variants for each specialized key type,
13//! and methods that interact with those enums have matches with similar-looking arms.
14//! Some day we may devise a better solution, but this is good enough for now.
15//
16// I (pgoldman 2024-02-05) suspect, but have not measured, that there's no real reason
17// to have a `ProductType` variant, which would apply to multi-column indexes.
18// I believe `ProductValue::cmp` to not be meaningfully faster than `AlgebraicValue::cmp`.
19// Eventually, we will likely want to compile comparison functions and representations
20// for `ProductValue`-keyed indexes which take advantage of type information,
21// since we know when creating the index the number and type of all the indexed columns.
22// This may involve a bytecode compiler, a tree of closures, or a native JIT.
23///
24/// We also represent unique indices more compactly than non-unique ones, avoiding the multi-map.
25/// Additionally, beyond our btree indices,
26/// we support direct unique indices, where key are indices into `Vec`s.
27use super::indexes::RowPointer;
28use super::table::RowRef;
29use crate::{read_column::ReadColumn, static_assert_size, MemoryUsage};
30use core::ops::RangeBounds;
31use spacetimedb_primitives::ColList;
32use spacetimedb_sats::{
33    algebraic_value::Packed, i256, product_value::InvalidFieldError, sum_value::SumTag, u256, AlgebraicType,
34    AlgebraicValue, ProductType,
35};
36
37mod key_size;
38mod multimap;
39pub mod unique_direct_fixed_cap_index;
40pub mod unique_direct_index;
41pub mod uniquemap;
42
43pub use key_size::KeySize;
44use spacetimedb_schema::def::IndexAlgorithm;
45use unique_direct_fixed_cap_index::{UniqueDirectFixedCapIndex, UniqueDirectFixedCapIndexRangeIter};
46use unique_direct_index::{UniqueDirectIndex, UniqueDirectIndexPointIter, UniqueDirectIndexRangeIter};
47
48type BtreeIndex<K> = multimap::MultiMap<K, RowPointer>;
49type BtreeIndexPointIter<'a> = multimap::MultiMapPointIter<'a, RowPointer>;
50type BtreeIndexRangeIter<'a, K> = multimap::MultiMapRangeIter<'a, K, RowPointer>;
51type BtreeUniqueIndex<K> = uniquemap::UniqueMap<K, RowPointer>;
52type BtreeUniqueIndexPointIter<'a> = uniquemap::UniqueMapPointIter<'a, RowPointer>;
53type BtreeUniqueIndexRangeIter<'a, K> = uniquemap::UniqueMapRangeIter<'a, K, RowPointer>;
54
55/// A point iterator over a [`TypedIndex`], with a specialized key type.
56///
57/// See module docs for info about specialization.
58enum TypedIndexPointIter<'a> {
59    BTree(BtreeIndexPointIter<'a>),
60    UniqueBTree(BtreeUniqueIndexPointIter<'a>),
61    UniqueDirect(UniqueDirectIndexPointIter),
62}
63
64impl Iterator for TypedIndexPointIter<'_> {
65    type Item = RowPointer;
66    fn next(&mut self) -> Option<Self::Item> {
67        match self {
68            Self::BTree(this) => this.next().copied(),
69            Self::UniqueBTree(this) => this.next().copied(),
70            Self::UniqueDirect(this) => this.next(),
71        }
72    }
73}
74
75/// An iterator over rows matching a certain [`AlgebraicValue`] on the [`TableIndex`].
76pub struct TableIndexPointIter<'a> {
77    /// The iterator seeking for matching values.
78    iter: TypedIndexPointIter<'a>,
79}
80
81impl Iterator for TableIndexPointIter<'_> {
82    type Item = RowPointer;
83
84    fn next(&mut self) -> Option<Self::Item> {
85        self.iter.next()
86    }
87}
88
89/// A ranged iterator over a [`TypedIndex`], with a specialized key type.
90///
91/// See module docs for info about specialization.
92enum TypedIndexRangeIter<'a> {
93    // All the non-unique btree index iterators.
94    BtreeBool(BtreeIndexRangeIter<'a, bool>),
95    BtreeU8(BtreeIndexRangeIter<'a, u8>),
96    BtreeI8(BtreeIndexRangeIter<'a, i8>),
97    BtreeU16(BtreeIndexRangeIter<'a, u16>),
98    BtreeI16(BtreeIndexRangeIter<'a, i16>),
99    BtreeU32(BtreeIndexRangeIter<'a, u32>),
100    BtreeI32(BtreeIndexRangeIter<'a, i32>),
101    BtreeU64(BtreeIndexRangeIter<'a, u64>),
102    BtreeI64(BtreeIndexRangeIter<'a, i64>),
103    BtreeU128(BtreeIndexRangeIter<'a, Packed<u128>>),
104    BtreeI128(BtreeIndexRangeIter<'a, Packed<i128>>),
105    BtreeU256(BtreeIndexRangeIter<'a, u256>),
106    BtreeI256(BtreeIndexRangeIter<'a, i256>),
107    BtreeString(BtreeIndexRangeIter<'a, Box<str>>),
108    BtreeAV(BtreeIndexRangeIter<'a, AlgebraicValue>),
109
110    // All the unique btree index iterators.
111    UniqueBtreeBool(BtreeUniqueIndexRangeIter<'a, bool>),
112    UniqueBtreeU8(BtreeUniqueIndexRangeIter<'a, u8>),
113    UniqueBtreeI8(BtreeUniqueIndexRangeIter<'a, i8>),
114    UniqueBtreeU16(BtreeUniqueIndexRangeIter<'a, u16>),
115    UniqueBtreeI16(BtreeUniqueIndexRangeIter<'a, i16>),
116    UniqueBtreeU32(BtreeUniqueIndexRangeIter<'a, u32>),
117    UniqueBtreeI32(BtreeUniqueIndexRangeIter<'a, i32>),
118    UniqueBtreeU64(BtreeUniqueIndexRangeIter<'a, u64>),
119    UniqueBtreeI64(BtreeUniqueIndexRangeIter<'a, i64>),
120    UniqueBtreeU128(BtreeUniqueIndexRangeIter<'a, Packed<u128>>),
121    UniqueBtreeI128(BtreeUniqueIndexRangeIter<'a, Packed<i128>>),
122    UniqueBtreeU256(BtreeUniqueIndexRangeIter<'a, u256>),
123    UniqueBtreeI256(BtreeUniqueIndexRangeIter<'a, i256>),
124    UniqueBtreeString(BtreeUniqueIndexRangeIter<'a, Box<str>>),
125    UniqueBtreeAV(BtreeUniqueIndexRangeIter<'a, AlgebraicValue>),
126
127    UniqueDirect(UniqueDirectIndexRangeIter<'a>),
128    UniqueDirectU8(UniqueDirectFixedCapIndexRangeIter<'a>),
129}
130
131impl Iterator for TypedIndexRangeIter<'_> {
132    type Item = RowPointer;
133    fn next(&mut self) -> Option<Self::Item> {
134        match self {
135            Self::BtreeBool(this) => this.next().copied(),
136            Self::BtreeU8(this) => this.next().copied(),
137            Self::BtreeI8(this) => this.next().copied(),
138            Self::BtreeU16(this) => this.next().copied(),
139            Self::BtreeI16(this) => this.next().copied(),
140            Self::BtreeU32(this) => this.next().copied(),
141            Self::BtreeI32(this) => this.next().copied(),
142            Self::BtreeU64(this) => this.next().copied(),
143            Self::BtreeI64(this) => this.next().copied(),
144            Self::BtreeU128(this) => this.next().copied(),
145            Self::BtreeI128(this) => this.next().copied(),
146            Self::BtreeU256(this) => this.next().copied(),
147            Self::BtreeI256(this) => this.next().copied(),
148            Self::BtreeString(this) => this.next().copied(),
149            Self::BtreeAV(this) => this.next().copied(),
150
151            Self::UniqueBtreeBool(this) => this.next().copied(),
152            Self::UniqueBtreeU8(this) => this.next().copied(),
153            Self::UniqueBtreeI8(this) => this.next().copied(),
154            Self::UniqueBtreeU16(this) => this.next().copied(),
155            Self::UniqueBtreeI16(this) => this.next().copied(),
156            Self::UniqueBtreeU32(this) => this.next().copied(),
157            Self::UniqueBtreeI32(this) => this.next().copied(),
158            Self::UniqueBtreeU64(this) => this.next().copied(),
159            Self::UniqueBtreeI64(this) => this.next().copied(),
160            Self::UniqueBtreeU128(this) => this.next().copied(),
161            Self::UniqueBtreeI128(this) => this.next().copied(),
162            Self::UniqueBtreeU256(this) => this.next().copied(),
163            Self::UniqueBtreeI256(this) => this.next().copied(),
164            Self::UniqueBtreeString(this) => this.next().copied(),
165            Self::UniqueBtreeAV(this) => this.next().copied(),
166
167            Self::UniqueDirect(this) => this.next(),
168            Self::UniqueDirectU8(this) => this.next(),
169        }
170    }
171}
172
173/// An iterator over rows matching a range of [`AlgebraicValue`]s on the [`TableIndex`].
174pub struct TableIndexRangeIter<'a> {
175    /// The iterator seeking for matching values.
176    iter: TypedIndexRangeIter<'a>,
177}
178
179impl Iterator for TableIndexRangeIter<'_> {
180    type Item = RowPointer;
181
182    fn next(&mut self) -> Option<Self::Item> {
183        self.iter.next()
184    }
185}
186
187/// An index from a key type determined at runtime to `RowPointer`(s).
188///
189/// See module docs for info about specialization.
190#[derive(Debug, PartialEq, Eq)]
191enum TypedIndex {
192    // All the non-unique btree index types.
193    BtreeBool(BtreeIndex<bool>),
194    BtreeU8(BtreeIndex<u8>),
195    BtreeSumTag(BtreeIndex<u8>),
196    BtreeI8(BtreeIndex<i8>),
197    BtreeU16(BtreeIndex<u16>),
198    BtreeI16(BtreeIndex<i16>),
199    BtreeU32(BtreeIndex<u32>),
200    BtreeI32(BtreeIndex<i32>),
201    BtreeU64(BtreeIndex<u64>),
202    BtreeI64(BtreeIndex<i64>),
203    BtreeU128(BtreeIndex<Packed<u128>>),
204    BtreeI128(BtreeIndex<Packed<i128>>),
205    BtreeU256(BtreeIndex<u256>),
206    BtreeI256(BtreeIndex<i256>),
207    BtreeString(BtreeIndex<Box<str>>),
208    BtreeAV(BtreeIndex<AlgebraicValue>),
209
210    // All the unique btree index types.
211    UniqueBtreeBool(BtreeUniqueIndex<bool>),
212    UniqueBtreeU8(BtreeUniqueIndex<u8>),
213    UniqueBtreeSumTag(BtreeUniqueIndex<u8>),
214    UniqueBtreeI8(BtreeUniqueIndex<i8>),
215    UniqueBtreeU16(BtreeUniqueIndex<u16>),
216    UniqueBtreeI16(BtreeUniqueIndex<i16>),
217    UniqueBtreeU32(BtreeUniqueIndex<u32>),
218    UniqueBtreeI32(BtreeUniqueIndex<i32>),
219    UniqueBtreeU64(BtreeUniqueIndex<u64>),
220    UniqueBtreeI64(BtreeUniqueIndex<i64>),
221    UniqueBtreeU128(BtreeUniqueIndex<Packed<u128>>),
222    UniqueBtreeI128(BtreeUniqueIndex<Packed<i128>>),
223    UniqueBtreeU256(BtreeUniqueIndex<u256>),
224    UniqueBtreeI256(BtreeUniqueIndex<i256>),
225    UniqueBtreeString(BtreeUniqueIndex<Box<str>>),
226    UniqueBtreeAV(BtreeUniqueIndex<AlgebraicValue>),
227
228    // All the unique direct index types.
229    UniqueDirectU8(UniqueDirectIndex),
230    UniqueDirectSumTag(UniqueDirectFixedCapIndex),
231    UniqueDirectU16(UniqueDirectIndex),
232    UniqueDirectU32(UniqueDirectIndex),
233    UniqueDirectU64(UniqueDirectIndex),
234}
235
236impl MemoryUsage for TypedIndex {
237    fn heap_usage(&self) -> usize {
238        match self {
239            TypedIndex::BtreeBool(this) => this.heap_usage(),
240            TypedIndex::BtreeU8(this) | TypedIndex::BtreeSumTag(this) => this.heap_usage(),
241            TypedIndex::BtreeI8(this) => this.heap_usage(),
242            TypedIndex::BtreeU16(this) => this.heap_usage(),
243            TypedIndex::BtreeI16(this) => this.heap_usage(),
244            TypedIndex::BtreeU32(this) => this.heap_usage(),
245            TypedIndex::BtreeI32(this) => this.heap_usage(),
246            TypedIndex::BtreeU64(this) => this.heap_usage(),
247            TypedIndex::BtreeI64(this) => this.heap_usage(),
248            TypedIndex::BtreeU128(this) => this.heap_usage(),
249            TypedIndex::BtreeI128(this) => this.heap_usage(),
250            TypedIndex::BtreeU256(this) => this.heap_usage(),
251            TypedIndex::BtreeI256(this) => this.heap_usage(),
252            TypedIndex::BtreeString(this) => this.heap_usage(),
253            TypedIndex::BtreeAV(this) => this.heap_usage(),
254
255            TypedIndex::UniqueBtreeBool(this) => this.heap_usage(),
256            TypedIndex::UniqueBtreeU8(this) | TypedIndex::UniqueBtreeSumTag(this) => this.heap_usage(),
257            TypedIndex::UniqueBtreeI8(this) => this.heap_usage(),
258            TypedIndex::UniqueBtreeU16(this) => this.heap_usage(),
259            TypedIndex::UniqueBtreeI16(this) => this.heap_usage(),
260            TypedIndex::UniqueBtreeU32(this) => this.heap_usage(),
261            TypedIndex::UniqueBtreeI32(this) => this.heap_usage(),
262            TypedIndex::UniqueBtreeU64(this) => this.heap_usage(),
263            TypedIndex::UniqueBtreeI64(this) => this.heap_usage(),
264            TypedIndex::UniqueBtreeU128(this) => this.heap_usage(),
265            TypedIndex::UniqueBtreeI128(this) => this.heap_usage(),
266            TypedIndex::UniqueBtreeU256(this) => this.heap_usage(),
267            TypedIndex::UniqueBtreeI256(this) => this.heap_usage(),
268            TypedIndex::UniqueBtreeString(this) => this.heap_usage(),
269            TypedIndex::UniqueBtreeAV(this) => this.heap_usage(),
270
271            TypedIndex::UniqueDirectSumTag(this) => this.heap_usage(),
272            TypedIndex::UniqueDirectU8(this)
273            | TypedIndex::UniqueDirectU16(this)
274            | TypedIndex::UniqueDirectU32(this)
275            | TypedIndex::UniqueDirectU64(this) => this.heap_usage(),
276        }
277    }
278}
279
280fn as_tag(av: &AlgebraicValue) -> Option<&u8> {
281    av.as_sum().map(|s| &s.tag)
282}
283
284impl TypedIndex {
285    /// Returns a new index with keys being of `key_type` and the index possibly `is_unique`.
286    fn new(key_type: &AlgebraicType, index_algo: &IndexAlgorithm, is_unique: bool) -> Self {
287        use TypedIndex::*;
288
289        if let IndexAlgorithm::Direct(_) = index_algo {
290            assert!(is_unique);
291            return match key_type {
292                AlgebraicType::U8 => Self::UniqueDirectU8(<_>::default()),
293                AlgebraicType::U16 => Self::UniqueDirectU16(<_>::default()),
294                AlgebraicType::U32 => Self::UniqueDirectU32(<_>::default()),
295                AlgebraicType::U64 => Self::UniqueDirectU64(<_>::default()),
296                // For a plain enum, use `u8` as the native type.
297                AlgebraicType::Sum(sum) if sum.is_simple_enum() => {
298                    UniqueDirectSumTag(UniqueDirectFixedCapIndex::new(sum.variants.len()))
299                }
300                _ => unreachable!("unexpected key type {key_type:?} for direct index"),
301            };
302        }
303
304        // If the index is on a single column of a primitive type, string, or plain enum,
305        // use a homogeneous map with a native key type.
306        if is_unique {
307            match key_type {
308                AlgebraicType::Bool => UniqueBtreeBool(<_>::default()),
309                AlgebraicType::I8 => UniqueBtreeI8(<_>::default()),
310                AlgebraicType::U8 => UniqueBtreeU8(<_>::default()),
311                AlgebraicType::I16 => UniqueBtreeI16(<_>::default()),
312                AlgebraicType::U16 => UniqueBtreeU16(<_>::default()),
313                AlgebraicType::I32 => UniqueBtreeI32(<_>::default()),
314                AlgebraicType::U32 => UniqueBtreeU32(<_>::default()),
315                AlgebraicType::I64 => UniqueBtreeI64(<_>::default()),
316                AlgebraicType::U64 => UniqueBtreeU64(<_>::default()),
317                AlgebraicType::I128 => UniqueBtreeI128(<_>::default()),
318                AlgebraicType::U128 => UniqueBtreeU128(<_>::default()),
319                AlgebraicType::I256 => UniqueBtreeI256(<_>::default()),
320                AlgebraicType::U256 => UniqueBtreeU256(<_>::default()),
321                AlgebraicType::String => UniqueBtreeString(<_>::default()),
322                // For a plain enum, use `u8` as the native type.
323                // We use a direct index here
324                AlgebraicType::Sum(sum) if sum.is_simple_enum() => UniqueBtreeSumTag(<_>::default()),
325
326                // The index is either multi-column,
327                // or we don't care to specialize on the key type,
328                // so use a map keyed on `AlgebraicValue`.
329                _ => UniqueBtreeAV(<_>::default()),
330            }
331        } else {
332            match key_type {
333                AlgebraicType::Bool => BtreeBool(<_>::default()),
334                AlgebraicType::I8 => BtreeI8(<_>::default()),
335                AlgebraicType::U8 => BtreeU8(<_>::default()),
336                AlgebraicType::I16 => BtreeI16(<_>::default()),
337                AlgebraicType::U16 => BtreeU16(<_>::default()),
338                AlgebraicType::I32 => BtreeI32(<_>::default()),
339                AlgebraicType::U32 => BtreeU32(<_>::default()),
340                AlgebraicType::I64 => BtreeI64(<_>::default()),
341                AlgebraicType::U64 => BtreeU64(<_>::default()),
342                AlgebraicType::I128 => BtreeI128(<_>::default()),
343                AlgebraicType::U128 => BtreeU128(<_>::default()),
344                AlgebraicType::I256 => BtreeI256(<_>::default()),
345                AlgebraicType::U256 => BtreeU256(<_>::default()),
346                AlgebraicType::String => BtreeString(<_>::default()),
347
348                // For a plain enum, use `u8` as the native type.
349                AlgebraicType::Sum(sum) if sum.is_simple_enum() => BtreeSumTag(<_>::default()),
350
351                // The index is either multi-column,
352                // or we don't care to specialize on the key type,
353                // so use a map keyed on `AlgebraicValue`.
354                _ => BtreeAV(<_>::default()),
355            }
356        }
357    }
358
359    /// Clones the structure of this index but not the indexed elements,
360    /// so the returned index is empty.
361    fn clone_structure(&self) -> Self {
362        use TypedIndex::*;
363        match self {
364            BtreeBool(_) => BtreeBool(<_>::default()),
365            BtreeU8(_) => BtreeU8(<_>::default()),
366            BtreeSumTag(_) => BtreeSumTag(<_>::default()),
367            BtreeI8(_) => BtreeI8(<_>::default()),
368            BtreeU16(_) => BtreeU16(<_>::default()),
369            BtreeI16(_) => BtreeI16(<_>::default()),
370            BtreeU32(_) => BtreeU32(<_>::default()),
371            BtreeI32(_) => BtreeI32(<_>::default()),
372            BtreeU64(_) => BtreeU64(<_>::default()),
373            BtreeI64(_) => BtreeI64(<_>::default()),
374            BtreeU128(_) => BtreeU128(<_>::default()),
375            BtreeI128(_) => BtreeI128(<_>::default()),
376            BtreeU256(_) => BtreeU256(<_>::default()),
377            BtreeI256(_) => BtreeI256(<_>::default()),
378            BtreeString(_) => BtreeString(<_>::default()),
379            BtreeAV(_) => BtreeAV(<_>::default()),
380            UniqueBtreeBool(_) => UniqueBtreeBool(<_>::default()),
381            UniqueBtreeU8(_) => UniqueBtreeU8(<_>::default()),
382            UniqueBtreeSumTag(_) => UniqueBtreeSumTag(<_>::default()),
383            UniqueBtreeI8(_) => UniqueBtreeI8(<_>::default()),
384            UniqueBtreeU16(_) => UniqueBtreeU16(<_>::default()),
385            UniqueBtreeI16(_) => UniqueBtreeI16(<_>::default()),
386            UniqueBtreeU32(_) => UniqueBtreeU32(<_>::default()),
387            UniqueBtreeI32(_) => UniqueBtreeI32(<_>::default()),
388            UniqueBtreeU64(_) => UniqueBtreeU64(<_>::default()),
389            UniqueBtreeI64(_) => UniqueBtreeI64(<_>::default()),
390            UniqueBtreeU128(_) => UniqueBtreeU128(<_>::default()),
391            UniqueBtreeI128(_) => UniqueBtreeI128(<_>::default()),
392            UniqueBtreeU256(_) => UniqueBtreeU256(<_>::default()),
393            UniqueBtreeI256(_) => UniqueBtreeI256(<_>::default()),
394            UniqueBtreeString(_) => UniqueBtreeString(<_>::default()),
395            UniqueBtreeAV(_) => UniqueBtreeAV(<_>::default()),
396            UniqueDirectU8(_) => UniqueDirectU8(<_>::default()),
397            UniqueDirectSumTag(idx) => UniqueDirectSumTag(idx.clone_structure()),
398            UniqueDirectU16(_) => UniqueDirectU16(<_>::default()),
399            UniqueDirectU32(_) => UniqueDirectU32(<_>::default()),
400            UniqueDirectU64(_) => UniqueDirectU64(<_>::default()),
401        }
402    }
403
404    /// Returns whether this is a unique index or not.
405    fn is_unique(&self) -> bool {
406        use TypedIndex::*;
407        match self {
408            BtreeBool(_) | BtreeU8(_) | BtreeSumTag(_) | BtreeI8(_) | BtreeU16(_) | BtreeI16(_) | BtreeU32(_)
409            | BtreeI32(_) | BtreeU64(_) | BtreeI64(_) | BtreeU128(_) | BtreeI128(_) | BtreeU256(_) | BtreeI256(_)
410            | BtreeString(_) | BtreeAV(_) => false,
411            UniqueBtreeBool(_)
412            | UniqueBtreeU8(_)
413            | UniqueBtreeSumTag(_)
414            | UniqueBtreeI8(_)
415            | UniqueBtreeU16(_)
416            | UniqueBtreeI16(_)
417            | UniqueBtreeU32(_)
418            | UniqueBtreeI32(_)
419            | UniqueBtreeU64(_)
420            | UniqueBtreeI64(_)
421            | UniqueBtreeU128(_)
422            | UniqueBtreeI128(_)
423            | UniqueBtreeU256(_)
424            | UniqueBtreeI256(_)
425            | UniqueBtreeString(_)
426            | UniqueBtreeAV(_)
427            | UniqueDirectU8(_)
428            | UniqueDirectSumTag(_)
429            | UniqueDirectU16(_)
430            | UniqueDirectU32(_)
431            | UniqueDirectU64(_) => true,
432        }
433    }
434
435    /// Add the row referred to by `row_ref` to the index `self`,
436    /// which must be keyed at `cols`.
437    ///
438    /// The returned `usize` is the number of bytes used by the key.
439    /// [`TableIndex::check_and_insert`] will use this
440    /// to update the counter for [`TableIndex::num_key_bytes`].
441    /// We want to store said counter outside of the [`TypedIndex`] enum,
442    /// but we can only compute the size using type info within the [`TypedIndex`],
443    /// so we have to return the size across this boundary.
444    ///
445    /// Returns `Errs(existing_row)` if this index was a unique index that was violated.
446    /// The index is not inserted to in that case.
447    ///
448    /// # Safety
449    ///
450    /// 1. Caller promises that `cols` matches what was given at construction (`Self::new`).
451    /// 2. Caller promises that the projection of `row_ref`'s type's equals the index's key type.
452    unsafe fn insert(&mut self, cols: &ColList, row_ref: RowRef<'_>) -> Result<usize, RowPointer> {
453        fn project_to_singleton_key<T: ReadColumn>(cols: &ColList, row_ref: RowRef<'_>) -> T {
454            // Extract the column.
455            let col_pos = cols.as_singleton();
456            // SAFETY: Caller promised that `cols` matches what was given at construction (`Self::new`).
457            // In the case of `.clone_structure()`, the structure is preserved,
458            // so the promise is also preserved.
459            // This entails, that because we reached here, that `cols` is singleton.
460            let col_pos = unsafe { col_pos.unwrap_unchecked() }.idx();
461
462            // Extract the layout of the column.
463            let col_layouts = &row_ref.row_layout().product().elements;
464            // SAFETY:
465            // - Caller promised that projecting the `row_ref`'s type/layout to `self.indexed_columns`
466            //   gives us the index's key type.
467            //   This entails that each `ColId` in `self.indexed_columns`
468            //   must be in-bounds of `row_ref`'s layout.
469            let col_layout = unsafe { col_layouts.get_unchecked(col_pos) };
470
471            // Read the column in `row_ref`.
472            // SAFETY:
473            // - `col_layout` was just derived from the row layout.
474            // - Caller promised that the type-projection of the row type/layout
475            //   equals the index's key type.
476            //   We've reached here, so the index's key type is compatible with `T`.
477            // - `self` is a valid row so offsetting to `col_layout` is valid.
478            unsafe { T::unchecked_read_column(row_ref, col_layout) }
479        }
480
481        fn mm_insert_at_type<T: Ord + ReadColumn + KeySize>(
482            this: &mut BtreeIndex<T>,
483            cols: &ColList,
484            row_ref: RowRef<'_>,
485        ) -> Result<usize, RowPointer> {
486            let key: T = project_to_singleton_key(cols, row_ref);
487            let key_size = key.key_size_in_bytes();
488            this.insert(key, row_ref.pointer());
489            Ok(key_size)
490        }
491        fn um_insert_at_type<T: Ord + ReadColumn + KeySize>(
492            this: &mut BtreeUniqueIndex<T>,
493            cols: &ColList,
494            row_ref: RowRef<'_>,
495        ) -> Result<usize, RowPointer> {
496            let key: T = project_to_singleton_key(cols, row_ref);
497            let key_size = key.key_size_in_bytes();
498            this.insert(key, row_ref.pointer())
499                .map_err(|ptr| *ptr)
500                .map(|_| key_size)
501        }
502        fn direct_insert_at_type<T: ReadColumn>(
503            this: &mut UniqueDirectIndex,
504            cols: &ColList,
505            row_ref: RowRef<'_>,
506            to_usize: impl FnOnce(T) -> usize,
507        ) -> Result<usize, RowPointer> {
508            let key: T = project_to_singleton_key(cols, row_ref);
509            let key = to_usize(key);
510            let key_size = key.key_size_in_bytes();
511            this.insert(key, row_ref.pointer()).map(|_| key_size)
512        }
513        fn direct_u8_insert_at_type<T: ReadColumn>(
514            this: &mut UniqueDirectFixedCapIndex,
515            cols: &ColList,
516            row_ref: RowRef<'_>,
517            to_u8: impl FnOnce(T) -> usize,
518        ) -> Result<usize, RowPointer> {
519            let key: T = project_to_singleton_key(cols, row_ref);
520            let key = to_u8(key);
521            let key_size = key.key_size_in_bytes();
522            this.insert(key, row_ref.pointer()).map(|_| key_size)
523        }
524        match self {
525            Self::BtreeBool(idx) => mm_insert_at_type(idx, cols, row_ref),
526            Self::BtreeU8(idx) => mm_insert_at_type(idx, cols, row_ref),
527            Self::BtreeSumTag(idx) => {
528                let SumTag(key) = project_to_singleton_key(cols, row_ref);
529                let key_size = key.key_size_in_bytes();
530                idx.insert(key, row_ref.pointer());
531                Ok(key_size)
532            }
533            Self::BtreeI8(idx) => mm_insert_at_type(idx, cols, row_ref),
534            Self::BtreeU16(idx) => mm_insert_at_type(idx, cols, row_ref),
535            Self::BtreeI16(idx) => mm_insert_at_type(idx, cols, row_ref),
536            Self::BtreeU32(idx) => mm_insert_at_type(idx, cols, row_ref),
537            Self::BtreeI32(idx) => mm_insert_at_type(idx, cols, row_ref),
538            Self::BtreeU64(idx) => mm_insert_at_type(idx, cols, row_ref),
539            Self::BtreeI64(idx) => mm_insert_at_type(idx, cols, row_ref),
540            Self::BtreeU128(idx) => mm_insert_at_type(idx, cols, row_ref),
541            Self::BtreeI128(idx) => mm_insert_at_type(idx, cols, row_ref),
542            Self::BtreeU256(idx) => mm_insert_at_type(idx, cols, row_ref),
543            Self::BtreeI256(idx) => mm_insert_at_type(idx, cols, row_ref),
544            Self::BtreeString(idx) => mm_insert_at_type(idx, cols, row_ref),
545            Self::BtreeAV(this) => {
546                // SAFETY: Caller promised that any `col` in `cols` is in-bounds of `row_ref`'s layout.
547                let key = unsafe { row_ref.project_unchecked(cols) };
548                let key_size = key.key_size_in_bytes();
549                this.insert(key, row_ref.pointer());
550                Ok(key_size)
551            }
552            Self::UniqueBtreeBool(idx) => um_insert_at_type(idx, cols, row_ref),
553            Self::UniqueBtreeU8(idx) => um_insert_at_type(idx, cols, row_ref),
554            Self::UniqueBtreeSumTag(idx) => {
555                let SumTag(key) = project_to_singleton_key(cols, row_ref);
556                let key_size = key.key_size_in_bytes();
557                idx.insert(key, row_ref.pointer()).map_err(|ptr| *ptr).map(|_| key_size)
558            }
559            Self::UniqueBtreeI8(idx) => um_insert_at_type(idx, cols, row_ref),
560            Self::UniqueBtreeU16(idx) => um_insert_at_type(idx, cols, row_ref),
561            Self::UniqueBtreeI16(idx) => um_insert_at_type(idx, cols, row_ref),
562            Self::UniqueBtreeU32(idx) => um_insert_at_type(idx, cols, row_ref),
563            Self::UniqueBtreeI32(idx) => um_insert_at_type(idx, cols, row_ref),
564            Self::UniqueBtreeU64(idx) => um_insert_at_type(idx, cols, row_ref),
565            Self::UniqueBtreeI64(idx) => um_insert_at_type(idx, cols, row_ref),
566            Self::UniqueBtreeU128(idx) => um_insert_at_type(idx, cols, row_ref),
567            Self::UniqueBtreeI128(idx) => um_insert_at_type(idx, cols, row_ref),
568            Self::UniqueBtreeU256(idx) => um_insert_at_type(idx, cols, row_ref),
569            Self::UniqueBtreeI256(idx) => um_insert_at_type(idx, cols, row_ref),
570            Self::UniqueBtreeString(idx) => um_insert_at_type(idx, cols, row_ref),
571            Self::UniqueBtreeAV(this) => {
572                // SAFETY: Caller promised that any `col` in `cols` is in-bounds of `row_ref`'s layout.
573                let key = unsafe { row_ref.project_unchecked(cols) };
574                let key_size = key.key_size_in_bytes();
575                this.insert(key, row_ref.pointer())
576                    .map_err(|ptr| *ptr)
577                    .map(|_| key_size)
578            }
579            Self::UniqueDirectSumTag(idx) => direct_u8_insert_at_type(idx, cols, row_ref, |SumTag(tag)| tag as usize),
580            Self::UniqueDirectU8(idx) => direct_insert_at_type(idx, cols, row_ref, |k: u8| k as usize),
581            Self::UniqueDirectU16(idx) => direct_insert_at_type(idx, cols, row_ref, |k: u16| k as usize),
582            Self::UniqueDirectU32(idx) => direct_insert_at_type(idx, cols, row_ref, |k: u32| k as usize),
583            Self::UniqueDirectU64(idx) => direct_insert_at_type(idx, cols, row_ref, |k: u64| k as usize),
584        }
585    }
586
587    /// Remove the row referred to by `row_ref` from the index `self`,
588    /// which must be keyed at `cols`.
589    ///
590    /// If `cols` is inconsistent with `self`,
591    /// or the `row_ref` has a row type other than that used for `self`,
592    /// this will behave oddly; it may return an error, do nothing,
593    /// or remove the wrong value from the index.
594    /// Note, however, that it will not invoke undefined behavior.
595    ///
596    /// If the row was present and has been deleted, returns `Ok(Some(key_size_in_bytes))`,
597    /// where `key_size_in_bytes` is the size of the key.
598    /// [`TableIndex::delete`] will use this
599    /// to update the counter for [`TableIndex::num_key_bytes`].
600    /// We want to store said counter outside of the [`TypedIndex`] enum,
601    /// but we can only compute the size using type info within the [`TypedIndex`],
602    /// so we have to return the size across this boundary.
603    // TODO(centril): make this unsafe and use unchecked conversions.
604    fn delete(&mut self, cols: &ColList, row_ref: RowRef<'_>) -> Result<Option<usize>, InvalidFieldError> {
605        fn mm_delete_at_type<T: Ord + ReadColumn + KeySize>(
606            this: &mut BtreeIndex<T>,
607            cols: &ColList,
608            row_ref: RowRef<'_>,
609        ) -> Result<Option<usize>, InvalidFieldError> {
610            let col_pos = cols.as_singleton().unwrap();
611            let key: T = row_ref.read_col(col_pos).map_err(|_| col_pos)?;
612            let key_size = key.key_size_in_bytes();
613            Ok(this.delete(&key, &row_ref.pointer()).then_some(key_size))
614        }
615        fn um_delete_at_type<T: Ord + ReadColumn + KeySize>(
616            this: &mut BtreeUniqueIndex<T>,
617            cols: &ColList,
618            row_ref: RowRef<'_>,
619        ) -> Result<Option<usize>, InvalidFieldError> {
620            let col_pos = cols.as_singleton().unwrap();
621            let key: T = row_ref.read_col(col_pos).map_err(|_| col_pos)?;
622            let key_size = key.key_size_in_bytes();
623            Ok(this.delete(&key).then_some(key_size))
624        }
625        fn direct_delete_at_type<T: ReadColumn>(
626            this: &mut UniqueDirectIndex,
627            cols: &ColList,
628            row_ref: RowRef<'_>,
629            to_usize: impl FnOnce(T) -> usize,
630        ) -> Result<Option<usize>, InvalidFieldError> {
631            let col_pos = cols.as_singleton().unwrap();
632            let key: T = row_ref.read_col(col_pos).map_err(|_| col_pos)?;
633            let key = to_usize(key);
634            let key_size = key.key_size_in_bytes();
635            Ok(this.delete(key).then_some(key_size))
636        }
637        fn direct_u8_delete_at_type<T: ReadColumn>(
638            this: &mut UniqueDirectFixedCapIndex,
639            cols: &ColList,
640            row_ref: RowRef<'_>,
641            to_u8: impl FnOnce(T) -> usize,
642        ) -> Result<Option<usize>, InvalidFieldError> {
643            let col_pos = cols.as_singleton().unwrap();
644            let key: T = row_ref.read_col(col_pos).map_err(|_| col_pos)?;
645            let key = to_u8(key);
646            let key_size = key.key_size_in_bytes();
647            Ok(this.delete(key).then_some(key_size))
648        }
649
650        match self {
651            Self::BtreeBool(this) => mm_delete_at_type(this, cols, row_ref),
652            Self::BtreeU8(this) => mm_delete_at_type(this, cols, row_ref),
653            Self::BtreeSumTag(this) => {
654                let col_pos = cols.as_singleton().unwrap();
655                let SumTag(key) = row_ref.read_col(col_pos).map_err(|_| col_pos)?;
656                let key_size = key.key_size_in_bytes();
657                Ok(this.delete(&key, &row_ref.pointer()).then_some(key_size))
658            }
659            Self::BtreeI8(this) => mm_delete_at_type(this, cols, row_ref),
660            Self::BtreeU16(this) => mm_delete_at_type(this, cols, row_ref),
661            Self::BtreeI16(this) => mm_delete_at_type(this, cols, row_ref),
662            Self::BtreeU32(this) => mm_delete_at_type(this, cols, row_ref),
663            Self::BtreeI32(this) => mm_delete_at_type(this, cols, row_ref),
664            Self::BtreeU64(this) => mm_delete_at_type(this, cols, row_ref),
665            Self::BtreeI64(this) => mm_delete_at_type(this, cols, row_ref),
666            Self::BtreeU128(this) => mm_delete_at_type(this, cols, row_ref),
667            Self::BtreeI128(this) => mm_delete_at_type(this, cols, row_ref),
668            Self::BtreeU256(this) => mm_delete_at_type(this, cols, row_ref),
669            Self::BtreeI256(this) => mm_delete_at_type(this, cols, row_ref),
670            Self::BtreeString(this) => mm_delete_at_type(this, cols, row_ref),
671            Self::BtreeAV(this) => {
672                let key = row_ref.project(cols)?;
673                let key_size = key.key_size_in_bytes();
674                Ok(this.delete(&key, &row_ref.pointer()).then_some(key_size))
675            }
676            Self::UniqueBtreeBool(this) => um_delete_at_type(this, cols, row_ref),
677            Self::UniqueBtreeU8(this) => um_delete_at_type(this, cols, row_ref),
678            Self::UniqueBtreeSumTag(this) => {
679                let col_pos = cols.as_singleton().unwrap();
680                let SumTag(key) = row_ref.read_col(col_pos).map_err(|_| col_pos)?;
681                let key_size = key.key_size_in_bytes();
682                Ok(this.delete(&key).then_some(key_size))
683            }
684            Self::UniqueBtreeI8(this) => um_delete_at_type(this, cols, row_ref),
685            Self::UniqueBtreeU16(this) => um_delete_at_type(this, cols, row_ref),
686            Self::UniqueBtreeI16(this) => um_delete_at_type(this, cols, row_ref),
687            Self::UniqueBtreeU32(this) => um_delete_at_type(this, cols, row_ref),
688            Self::UniqueBtreeI32(this) => um_delete_at_type(this, cols, row_ref),
689            Self::UniqueBtreeU64(this) => um_delete_at_type(this, cols, row_ref),
690            Self::UniqueBtreeI64(this) => um_delete_at_type(this, cols, row_ref),
691            Self::UniqueBtreeU128(this) => um_delete_at_type(this, cols, row_ref),
692            Self::UniqueBtreeI128(this) => um_delete_at_type(this, cols, row_ref),
693            Self::UniqueBtreeU256(this) => um_delete_at_type(this, cols, row_ref),
694            Self::UniqueBtreeI256(this) => um_delete_at_type(this, cols, row_ref),
695            Self::UniqueBtreeString(this) => um_delete_at_type(this, cols, row_ref),
696            Self::UniqueBtreeAV(this) => {
697                let key = row_ref.project(cols)?;
698                let key_size = key.key_size_in_bytes();
699                Ok(this.delete(&key).then_some(key_size))
700            }
701            Self::UniqueDirectSumTag(this) => direct_u8_delete_at_type(this, cols, row_ref, |SumTag(k)| k as usize),
702            Self::UniqueDirectU8(this) => direct_delete_at_type(this, cols, row_ref, |k: u8| k as usize),
703            Self::UniqueDirectU16(this) => direct_delete_at_type(this, cols, row_ref, |k: u16| k as usize),
704            Self::UniqueDirectU32(this) => direct_delete_at_type(this, cols, row_ref, |k: u32| k as usize),
705            Self::UniqueDirectU64(this) => direct_delete_at_type(this, cols, row_ref, |k: u64| k as usize),
706        }
707    }
708
709    fn seek_point(&self, key: &AlgebraicValue) -> TypedIndexPointIter<'_> {
710        fn mm_iter_at_type<'a, T: Ord>(
711            this: &'a BtreeIndex<T>,
712            key: &AlgebraicValue,
713            av_as_t: impl Fn(&AlgebraicValue) -> Option<&T>,
714        ) -> BtreeIndexPointIter<'a> {
715            this.values_in_point(av_as_t(key).expect("key does not conform to key type of index"))
716        }
717        fn um_iter_at_type<'a, T: Ord>(
718            this: &'a BtreeUniqueIndex<T>,
719            key: &AlgebraicValue,
720            av_as_t: impl Fn(&AlgebraicValue) -> Option<&T>,
721        ) -> BtreeUniqueIndexPointIter<'a> {
722            this.values_in_point(av_as_t(key).expect("key does not conform to key type of index"))
723        }
724        fn direct_iter_at_type<T>(
725            this: &UniqueDirectIndex,
726            key: &AlgebraicValue,
727            av_as_t: impl Fn(&AlgebraicValue) -> Option<&T>,
728            to_usize: impl Copy + FnOnce(&T) -> usize,
729        ) -> UniqueDirectIndexPointIter {
730            let av_as_t = |v| av_as_t(v).expect("key does not conform to key type of index");
731            this.seek_point(to_usize(av_as_t(key)))
732        }
733
734        use TypedIndex::*;
735        use TypedIndexPointIter::*;
736        match self {
737            BtreeBool(this) => BTree(mm_iter_at_type(this, key, AlgebraicValue::as_bool)),
738            BtreeU8(this) => BTree(mm_iter_at_type(this, key, AlgebraicValue::as_u8)),
739            BtreeSumTag(this) => BTree(mm_iter_at_type(this, key, as_tag)),
740            BtreeI8(this) => BTree(mm_iter_at_type(this, key, AlgebraicValue::as_i8)),
741            BtreeU16(this) => BTree(mm_iter_at_type(this, key, AlgebraicValue::as_u16)),
742            BtreeI16(this) => BTree(mm_iter_at_type(this, key, AlgebraicValue::as_i16)),
743            BtreeU32(this) => BTree(mm_iter_at_type(this, key, AlgebraicValue::as_u32)),
744            BtreeI32(this) => BTree(mm_iter_at_type(this, key, AlgebraicValue::as_i32)),
745            BtreeU64(this) => BTree(mm_iter_at_type(this, key, AlgebraicValue::as_u64)),
746            BtreeI64(this) => BTree(mm_iter_at_type(this, key, AlgebraicValue::as_i64)),
747            BtreeU128(this) => BTree(mm_iter_at_type(this, key, AlgebraicValue::as_u128)),
748            BtreeI128(this) => BTree(mm_iter_at_type(this, key, AlgebraicValue::as_i128)),
749            BtreeU256(this) => BTree(mm_iter_at_type(this, key, |av| av.as_u256().map(|x| &**x))),
750            BtreeI256(this) => BTree(mm_iter_at_type(this, key, |av| av.as_i256().map(|x| &**x))),
751            BtreeString(this) => BTree(mm_iter_at_type(this, key, AlgebraicValue::as_string)),
752            BtreeAV(this) => BTree(this.values_in_point(key)),
753
754            UniqueBtreeBool(this) => UniqueBTree(um_iter_at_type(this, key, AlgebraicValue::as_bool)),
755            UniqueBtreeU8(this) => UniqueBTree(um_iter_at_type(this, key, AlgebraicValue::as_u8)),
756            UniqueBtreeSumTag(this) => UniqueBTree(um_iter_at_type(this, key, as_tag)),
757            UniqueBtreeI8(this) => UniqueBTree(um_iter_at_type(this, key, AlgebraicValue::as_i8)),
758            UniqueBtreeU16(this) => UniqueBTree(um_iter_at_type(this, key, AlgebraicValue::as_u16)),
759            UniqueBtreeI16(this) => UniqueBTree(um_iter_at_type(this, key, AlgebraicValue::as_i16)),
760            UniqueBtreeU32(this) => UniqueBTree(um_iter_at_type(this, key, AlgebraicValue::as_u32)),
761            UniqueBtreeI32(this) => UniqueBTree(um_iter_at_type(this, key, AlgebraicValue::as_i32)),
762            UniqueBtreeU64(this) => UniqueBTree(um_iter_at_type(this, key, AlgebraicValue::as_u64)),
763            UniqueBtreeI64(this) => UniqueBTree(um_iter_at_type(this, key, AlgebraicValue::as_i64)),
764            UniqueBtreeU128(this) => UniqueBTree(um_iter_at_type(this, key, AlgebraicValue::as_u128)),
765            UniqueBtreeI128(this) => UniqueBTree(um_iter_at_type(this, key, AlgebraicValue::as_i128)),
766            UniqueBtreeU256(this) => UniqueBTree(um_iter_at_type(this, key, |av| av.as_u256().map(|x| &**x))),
767            UniqueBtreeI256(this) => UniqueBTree(um_iter_at_type(this, key, |av| av.as_i256().map(|x| &**x))),
768            UniqueBtreeString(this) => UniqueBTree(um_iter_at_type(this, key, AlgebraicValue::as_string)),
769            UniqueBtreeAV(this) => UniqueBTree(this.values_in_point(key)),
770
771            UniqueDirectSumTag(this) => {
772                let key = as_tag(key).expect("key does not conform to key type of index");
773                UniqueDirect(this.seek_point(*key as usize))
774            }
775            UniqueDirectU8(this) => {
776                UniqueDirect(direct_iter_at_type(this, key, AlgebraicValue::as_u8, |k| *k as usize))
777            }
778            UniqueDirectU16(this) => {
779                UniqueDirect(direct_iter_at_type(this, key, AlgebraicValue::as_u16, |k| *k as usize))
780            }
781            UniqueDirectU32(this) => {
782                UniqueDirect(direct_iter_at_type(this, key, AlgebraicValue::as_u32, |k| *k as usize))
783            }
784            UniqueDirectU64(this) => {
785                UniqueDirect(direct_iter_at_type(this, key, AlgebraicValue::as_u64, |k| *k as usize))
786            }
787        }
788    }
789
790    fn seek_range(&self, range: &impl RangeBounds<AlgebraicValue>) -> TypedIndexRangeIter<'_> {
791        fn mm_iter_at_type<'a, T: Ord>(
792            this: &'a BtreeIndex<T>,
793            range: &impl RangeBounds<AlgebraicValue>,
794            av_as_t: impl Fn(&AlgebraicValue) -> Option<&T>,
795        ) -> BtreeIndexRangeIter<'a, T> {
796            let av_as_t = |v| av_as_t(v).expect("bound does not conform to key type of index");
797            let start = range.start_bound().map(av_as_t);
798            let end = range.end_bound().map(av_as_t);
799            this.values_in_range(&(start, end))
800        }
801        fn um_iter_at_type<'a, T: Ord>(
802            this: &'a BtreeUniqueIndex<T>,
803            range: &impl RangeBounds<AlgebraicValue>,
804            av_as_t: impl Fn(&AlgebraicValue) -> Option<&T>,
805        ) -> BtreeUniqueIndexRangeIter<'a, T> {
806            let av_as_t = |v| av_as_t(v).expect("bound does not conform to key type of index");
807            let start = range.start_bound().map(av_as_t);
808            let end = range.end_bound().map(av_as_t);
809            this.values_in_range(&(start, end))
810        }
811        fn direct_iter_at_type<'a, T>(
812            this: &'a UniqueDirectIndex,
813            range: &impl RangeBounds<AlgebraicValue>,
814            av_as_t: impl Fn(&AlgebraicValue) -> Option<&T>,
815            to_usize: impl Copy + FnOnce(&T) -> usize,
816        ) -> UniqueDirectIndexRangeIter<'a> {
817            let av_as_t = |v| av_as_t(v).expect("bound does not conform to key type of index");
818            let start = range.start_bound().map(av_as_t).map(to_usize);
819            let end = range.end_bound().map(av_as_t).map(to_usize);
820            this.seek_range(&(start, end))
821        }
822
823        use TypedIndexRangeIter::*;
824        match self {
825            Self::BtreeBool(this) => BtreeBool(mm_iter_at_type(this, range, AlgebraicValue::as_bool)),
826            Self::BtreeU8(this) => BtreeU8(mm_iter_at_type(this, range, AlgebraicValue::as_u8)),
827            Self::BtreeSumTag(this) => BtreeU8(mm_iter_at_type(this, range, as_tag)),
828            Self::BtreeI8(this) => BtreeI8(mm_iter_at_type(this, range, AlgebraicValue::as_i8)),
829            Self::BtreeU16(this) => BtreeU16(mm_iter_at_type(this, range, AlgebraicValue::as_u16)),
830            Self::BtreeI16(this) => BtreeI16(mm_iter_at_type(this, range, AlgebraicValue::as_i16)),
831            Self::BtreeU32(this) => BtreeU32(mm_iter_at_type(this, range, AlgebraicValue::as_u32)),
832            Self::BtreeI32(this) => BtreeI32(mm_iter_at_type(this, range, AlgebraicValue::as_i32)),
833            Self::BtreeU64(this) => BtreeU64(mm_iter_at_type(this, range, AlgebraicValue::as_u64)),
834            Self::BtreeI64(this) => BtreeI64(mm_iter_at_type(this, range, AlgebraicValue::as_i64)),
835            Self::BtreeU128(this) => BtreeU128(mm_iter_at_type(this, range, AlgebraicValue::as_u128)),
836            Self::BtreeI128(this) => BtreeI128(mm_iter_at_type(this, range, AlgebraicValue::as_i128)),
837            Self::BtreeU256(this) => BtreeU256(mm_iter_at_type(this, range, |av| av.as_u256().map(|x| &**x))),
838            Self::BtreeI256(this) => BtreeI256(mm_iter_at_type(this, range, |av| av.as_i256().map(|x| &**x))),
839            Self::BtreeString(this) => BtreeString(mm_iter_at_type(this, range, AlgebraicValue::as_string)),
840            Self::BtreeAV(this) => BtreeAV(this.values_in_range(range)),
841
842            Self::UniqueBtreeBool(this) => UniqueBtreeBool(um_iter_at_type(this, range, AlgebraicValue::as_bool)),
843            Self::UniqueBtreeU8(this) => UniqueBtreeU8(um_iter_at_type(this, range, AlgebraicValue::as_u8)),
844            Self::UniqueBtreeSumTag(this) => UniqueBtreeU8(um_iter_at_type(this, range, as_tag)),
845            Self::UniqueBtreeI8(this) => UniqueBtreeI8(um_iter_at_type(this, range, AlgebraicValue::as_i8)),
846            Self::UniqueBtreeU16(this) => UniqueBtreeU16(um_iter_at_type(this, range, AlgebraicValue::as_u16)),
847            Self::UniqueBtreeI16(this) => UniqueBtreeI16(um_iter_at_type(this, range, AlgebraicValue::as_i16)),
848            Self::UniqueBtreeU32(this) => UniqueBtreeU32(um_iter_at_type(this, range, AlgebraicValue::as_u32)),
849            Self::UniqueBtreeI32(this) => UniqueBtreeI32(um_iter_at_type(this, range, AlgebraicValue::as_i32)),
850            Self::UniqueBtreeU64(this) => UniqueBtreeU64(um_iter_at_type(this, range, AlgebraicValue::as_u64)),
851            Self::UniqueBtreeI64(this) => UniqueBtreeI64(um_iter_at_type(this, range, AlgebraicValue::as_i64)),
852            Self::UniqueBtreeU128(this) => UniqueBtreeU128(um_iter_at_type(this, range, AlgebraicValue::as_u128)),
853            Self::UniqueBtreeI128(this) => UniqueBtreeI128(um_iter_at_type(this, range, AlgebraicValue::as_i128)),
854            Self::UniqueBtreeU256(this) => {
855                UniqueBtreeU256(um_iter_at_type(this, range, |av| av.as_u256().map(|x| &**x)))
856            }
857            Self::UniqueBtreeI256(this) => {
858                UniqueBtreeI256(um_iter_at_type(this, range, |av| av.as_i256().map(|x| &**x)))
859            }
860            Self::UniqueBtreeString(this) => UniqueBtreeString(um_iter_at_type(this, range, AlgebraicValue::as_string)),
861            Self::UniqueBtreeAV(this) => UniqueBtreeAV(this.values_in_range(range)),
862
863            Self::UniqueDirectSumTag(this) => {
864                let av_as_t = |v| as_tag(v).copied().expect("bound does not conform to key type of index") as usize;
865                let start = range.start_bound().map(av_as_t);
866                let end = range.end_bound().map(av_as_t);
867                let iter = this.seek_range(&(start, end));
868                UniqueDirectU8(iter)
869            }
870            Self::UniqueDirectU8(this) => {
871                UniqueDirect(direct_iter_at_type(this, range, AlgebraicValue::as_u8, |k| *k as usize))
872            }
873            Self::UniqueDirectU16(this) => {
874                UniqueDirect(direct_iter_at_type(this, range, AlgebraicValue::as_u16, |k| {
875                    *k as usize
876                }))
877            }
878            Self::UniqueDirectU32(this) => {
879                UniqueDirect(direct_iter_at_type(this, range, AlgebraicValue::as_u32, |k| {
880                    *k as usize
881                }))
882            }
883            Self::UniqueDirectU64(this) => {
884                UniqueDirect(direct_iter_at_type(this, range, AlgebraicValue::as_u64, |k| {
885                    *k as usize
886                }))
887            }
888        }
889    }
890
891    fn clear(&mut self) {
892        match self {
893            Self::BtreeBool(this) => this.clear(),
894            Self::BtreeU8(this) | Self::BtreeSumTag(this) => this.clear(),
895            Self::BtreeI8(this) => this.clear(),
896            Self::BtreeU16(this) => this.clear(),
897            Self::BtreeI16(this) => this.clear(),
898            Self::BtreeU32(this) => this.clear(),
899            Self::BtreeI32(this) => this.clear(),
900            Self::BtreeU64(this) => this.clear(),
901            Self::BtreeI64(this) => this.clear(),
902            Self::BtreeU128(this) => this.clear(),
903            Self::BtreeI128(this) => this.clear(),
904            Self::BtreeU256(this) => this.clear(),
905            Self::BtreeI256(this) => this.clear(),
906            Self::BtreeString(this) => this.clear(),
907            Self::BtreeAV(this) => this.clear(),
908
909            Self::UniqueBtreeBool(this) => this.clear(),
910            Self::UniqueBtreeU8(this) | Self::UniqueBtreeSumTag(this) => this.clear(),
911            Self::UniqueBtreeI8(this) => this.clear(),
912            Self::UniqueBtreeU16(this) => this.clear(),
913            Self::UniqueBtreeI16(this) => this.clear(),
914            Self::UniqueBtreeU32(this) => this.clear(),
915            Self::UniqueBtreeI32(this) => this.clear(),
916            Self::UniqueBtreeU64(this) => this.clear(),
917            Self::UniqueBtreeI64(this) => this.clear(),
918            Self::UniqueBtreeU128(this) => this.clear(),
919            Self::UniqueBtreeI128(this) => this.clear(),
920            Self::UniqueBtreeU256(this) => this.clear(),
921            Self::UniqueBtreeI256(this) => this.clear(),
922            Self::UniqueBtreeString(this) => this.clear(),
923            Self::UniqueBtreeAV(this) => this.clear(),
924
925            Self::UniqueDirectSumTag(this) => this.clear(),
926            Self::UniqueDirectU8(this)
927            | Self::UniqueDirectU16(this)
928            | Self::UniqueDirectU32(this)
929            | Self::UniqueDirectU64(this) => this.clear(),
930        }
931    }
932
933    #[allow(unused)] // used only by tests
934    fn is_empty(&self) -> bool {
935        self.len() == 0
936    }
937
938    #[allow(unused)] // used only by tests
939    fn len(&self) -> usize {
940        match self {
941            Self::BtreeBool(this) => this.len(),
942            Self::BtreeU8(this) | Self::BtreeSumTag(this) => this.len(),
943            Self::BtreeI8(this) => this.len(),
944            Self::BtreeU16(this) => this.len(),
945            Self::BtreeI16(this) => this.len(),
946            Self::BtreeU32(this) => this.len(),
947            Self::BtreeI32(this) => this.len(),
948            Self::BtreeU64(this) => this.len(),
949            Self::BtreeI64(this) => this.len(),
950            Self::BtreeU128(this) => this.len(),
951            Self::BtreeI128(this) => this.len(),
952            Self::BtreeU256(this) => this.len(),
953            Self::BtreeI256(this) => this.len(),
954            Self::BtreeString(this) => this.len(),
955            Self::BtreeAV(this) => this.len(),
956
957            Self::UniqueBtreeBool(this) => this.len(),
958            Self::UniqueBtreeU8(this) | Self::UniqueBtreeSumTag(this) => this.len(),
959            Self::UniqueBtreeI8(this) => this.len(),
960            Self::UniqueBtreeU16(this) => this.len(),
961            Self::UniqueBtreeI16(this) => this.len(),
962            Self::UniqueBtreeU32(this) => this.len(),
963            Self::UniqueBtreeI32(this) => this.len(),
964            Self::UniqueBtreeU64(this) => this.len(),
965            Self::UniqueBtreeI64(this) => this.len(),
966            Self::UniqueBtreeU128(this) => this.len(),
967            Self::UniqueBtreeI128(this) => this.len(),
968            Self::UniqueBtreeU256(this) => this.len(),
969            Self::UniqueBtreeI256(this) => this.len(),
970            Self::UniqueBtreeString(this) => this.len(),
971            Self::UniqueBtreeAV(this) => this.len(),
972
973            Self::UniqueDirectSumTag(this) => this.len(),
974            Self::UniqueDirectU8(this)
975            | Self::UniqueDirectU16(this)
976            | Self::UniqueDirectU32(this)
977            | Self::UniqueDirectU64(this) => this.len(),
978        }
979    }
980
981    fn num_keys(&self) -> usize {
982        match self {
983            Self::BtreeBool(this) => this.num_keys(),
984            Self::BtreeU8(this) | Self::BtreeSumTag(this) => this.num_keys(),
985            Self::BtreeI8(this) => this.num_keys(),
986            Self::BtreeU16(this) => this.num_keys(),
987            Self::BtreeI16(this) => this.num_keys(),
988            Self::BtreeU32(this) => this.num_keys(),
989            Self::BtreeI32(this) => this.num_keys(),
990            Self::BtreeU64(this) => this.num_keys(),
991            Self::BtreeI64(this) => this.num_keys(),
992            Self::BtreeU128(this) => this.num_keys(),
993            Self::BtreeI128(this) => this.num_keys(),
994            Self::BtreeU256(this) => this.num_keys(),
995            Self::BtreeI256(this) => this.num_keys(),
996            Self::BtreeString(this) => this.num_keys(),
997            Self::BtreeAV(this) => this.num_keys(),
998
999            Self::UniqueBtreeBool(this) => this.num_keys(),
1000            Self::UniqueBtreeU8(this) | Self::UniqueBtreeSumTag(this) => this.num_keys(),
1001            Self::UniqueBtreeI8(this) => this.num_keys(),
1002            Self::UniqueBtreeU16(this) => this.num_keys(),
1003            Self::UniqueBtreeI16(this) => this.num_keys(),
1004            Self::UniqueBtreeU32(this) => this.num_keys(),
1005            Self::UniqueBtreeI32(this) => this.num_keys(),
1006            Self::UniqueBtreeU64(this) => this.num_keys(),
1007            Self::UniqueBtreeI64(this) => this.num_keys(),
1008            Self::UniqueBtreeU128(this) => this.num_keys(),
1009            Self::UniqueBtreeI128(this) => this.num_keys(),
1010            Self::UniqueBtreeU256(this) => this.num_keys(),
1011            Self::UniqueBtreeI256(this) => this.num_keys(),
1012            Self::UniqueBtreeString(this) => this.num_keys(),
1013            Self::UniqueBtreeAV(this) => this.num_keys(),
1014
1015            Self::UniqueDirectSumTag(this) => this.num_keys(),
1016            Self::UniqueDirectU8(this)
1017            | Self::UniqueDirectU16(this)
1018            | Self::UniqueDirectU32(this)
1019            | Self::UniqueDirectU64(this) => this.num_keys(),
1020        }
1021    }
1022}
1023
1024/// An index on a set of [`ColId`]s of a table.
1025#[derive(Debug, PartialEq, Eq)]
1026pub struct TableIndex {
1027    /// The actual index, specialized for the appropriate key type.
1028    idx: TypedIndex,
1029    /// The key type of this index.
1030    /// This is the projection of the row type to the types of the columns indexed.
1031    // NOTE(centril): This is accessed in index scan ABIs for decoding, so don't `Box<_>` it.
1032    pub key_type: AlgebraicType,
1033
1034    /// The number of rows in this index.
1035    ///
1036    /// Memoized counter for [`Self::num_rows`].
1037    num_rows: u64,
1038
1039    /// The number of key bytes in this index.
1040    ///
1041    /// Memoized counter for [`Self::num_key_bytes`].
1042    /// See that method for more detailed documentation.
1043    num_key_bytes: u64,
1044
1045    /// Given a full row, typed at some `ty: ProductType`,
1046    /// these columns are the ones that this index indexes.
1047    /// Projecting the `ty` to `self.indexed_columns` yields the index's type `self.key_type`.
1048    pub indexed_columns: ColList,
1049}
1050
1051impl MemoryUsage for TableIndex {
1052    fn heap_usage(&self) -> usize {
1053        let Self {
1054            idx,
1055            key_type,
1056            num_rows,
1057            num_key_bytes,
1058            indexed_columns,
1059        } = self;
1060        idx.heap_usage()
1061            + key_type.heap_usage()
1062            + num_rows.heap_usage()
1063            + num_key_bytes.heap_usage()
1064            + indexed_columns.heap_usage()
1065    }
1066}
1067
1068static_assert_size!(TableIndex, 88);
1069
1070impl TableIndex {
1071    /// Returns a new possibly unique index, with `index_id` for a choice of indexing algorithm.
1072    pub fn new(
1073        row_type: &ProductType,
1074        index_algo: &IndexAlgorithm,
1075        is_unique: bool,
1076    ) -> Result<Self, InvalidFieldError> {
1077        let indexed_columns = index_algo.columns().to_owned();
1078        let key_type = row_type.project(&indexed_columns)?;
1079        let typed_index = TypedIndex::new(&key_type, index_algo, is_unique);
1080        Ok(Self {
1081            idx: typed_index,
1082            key_type,
1083            num_rows: 0,
1084            num_key_bytes: 0,
1085            indexed_columns,
1086        })
1087    }
1088
1089    /// Clones the structure of this index but not the indexed elements,
1090    /// so the returned index is empty.
1091    pub fn clone_structure(&self) -> Self {
1092        let key_type = self.key_type.clone();
1093        let idx = self.idx.clone_structure();
1094        let indexed_columns = self.indexed_columns.clone();
1095        Self {
1096            idx,
1097            key_type,
1098            num_rows: 0,
1099            num_key_bytes: 0,
1100            indexed_columns,
1101        }
1102    }
1103
1104    /// Returns whether this is a unique index or not.
1105    pub fn is_unique(&self) -> bool {
1106        self.idx.is_unique()
1107    }
1108
1109    /// Inserts `ptr` with the value `row` to this index.
1110    /// This index will extract the necessary values from `row` based on `self.indexed_columns`.
1111    ///
1112    /// Returns `Err(existing_row)` if this insertion would violate a unique constraint.
1113    ///
1114    /// # Safety
1115    ///
1116    /// Caller promises that projecting the `row_ref`'s type
1117    /// to the index's columns equals the index's key type.
1118    /// This is entailed by an index belonging to the table's schema.
1119    /// It also follows from `row_ref`'s type/layout
1120    /// being the same as passed in on `self`'s construction.
1121    pub unsafe fn check_and_insert(&mut self, row_ref: RowRef<'_>) -> Result<(), RowPointer> {
1122        // SAFETY:
1123        // 1. We're passing the same `ColList` that was provided during construction.
1124        // 2. Forward the caller's proof obligation.
1125        let res = unsafe { self.idx.insert(&self.indexed_columns, row_ref) };
1126        match res {
1127            Ok(key_size) => {
1128                // No existing row; the new row was inserted.
1129                // Update the `num_rows` and `num_key_bytes` counters
1130                // to account for the new insertion.
1131                self.num_rows += 1;
1132                self.num_key_bytes += key_size as u64;
1133                Ok(())
1134            }
1135            Err(e) => Err(e),
1136        }
1137    }
1138
1139    /// Deletes `row_ref` with its indexed value `row_ref.project(&self.indexed_columns)` from this index.
1140    ///
1141    /// Returns whether `ptr` was present.
1142    pub fn delete(&mut self, row_ref: RowRef<'_>) -> Result<bool, InvalidFieldError> {
1143        if let Some(size_in_bytes) = self.idx.delete(&self.indexed_columns, row_ref)? {
1144            // Was present, and deleted: update the `num_rows` and `num_key_bytes` counters.
1145            self.num_rows -= 1;
1146            self.num_key_bytes -= size_in_bytes as u64;
1147            Ok(true)
1148        } else {
1149            // Was not present: don't update counters.
1150            Ok(false)
1151        }
1152    }
1153
1154    /// Returns whether `value` is in this index.
1155    pub fn contains_any(&self, value: &AlgebraicValue) -> bool {
1156        self.seek_point(value).next().is_some()
1157    }
1158
1159    /// Returns the number of rows associated with this `value`.
1160    /// Returns `None` if 0.
1161    /// Returns `Some(1)` if the index is unique.
1162    pub fn count(&self, value: &AlgebraicValue) -> Option<usize> {
1163        match self.seek_point(value).count() {
1164            0 => None,
1165            n => Some(n),
1166        }
1167    }
1168
1169    /// Returns an iterator that yields all the `RowPointer`s for the given `key`.
1170    pub fn seek_point(&self, key: &AlgebraicValue) -> TableIndexPointIter<'_> {
1171        TableIndexPointIter {
1172            iter: self.idx.seek_point(key),
1173        }
1174    }
1175
1176    /// Returns an iterator over the [TableIndex] that yields all the `RowPointer`s
1177    /// that fall within the specified `range`.
1178    pub fn seek_range(&self, range: &impl RangeBounds<AlgebraicValue>) -> TableIndexRangeIter<'_> {
1179        TableIndexRangeIter {
1180            iter: self.idx.seek_range(range),
1181        }
1182    }
1183
1184    /// Extends [`TableIndex`] with `rows`.
1185    ///
1186    /// Returns the first unique constraint violation caused when adding this index, if any.
1187    ///
1188    /// # Safety
1189    ///
1190    /// Caller promises that projecting any of the `row_ref`'s type
1191    /// to the index's columns equals the index's key type.
1192    /// This is entailed by an index belonging to the table's schema.
1193    /// It also follows from `row_ref`'s type/layout
1194    /// being the same as passed in on `self`'s construction.
1195    pub unsafe fn build_from_rows<'table>(
1196        &mut self,
1197        rows: impl IntoIterator<Item = RowRef<'table>>,
1198    ) -> Result<(), RowPointer> {
1199        rows.into_iter()
1200            // SAFETY: Forward caller proof obligation.
1201            .try_for_each(|row_ref| unsafe { self.check_and_insert(row_ref) })
1202    }
1203
1204    /// Returns an error with the first unique constraint violation that
1205    /// would occur if `self` and `other` were to be merged.
1206    ///
1207    /// The closure `ignore` indicates whether a row in `self` should be ignored.
1208    pub fn can_merge(&self, other: &Self, ignore: impl Fn(&RowPointer) -> bool) -> Result<(), RowPointer> {
1209        use TypedIndex::*;
1210        match (&self.idx, &other.idx) {
1211            // For non-unique indices, it's always possible to merge.
1212            (BtreeBool(_), BtreeBool(_))
1213            | (BtreeU8(_), BtreeU8(_))
1214            | (BtreeSumTag(_), BtreeSumTag(_))
1215            | (BtreeI8(_), BtreeI8(_))
1216            | (BtreeU16(_), BtreeU16(_))
1217            | (BtreeI16(_), BtreeI16(_))
1218            | (BtreeU32(_), BtreeU32(_))
1219            | (BtreeI32(_), BtreeI32(_))
1220            | (BtreeU64(_), BtreeU64(_))
1221            | (BtreeI64(_), BtreeI64(_))
1222            | (BtreeU128(_), BtreeU128(_))
1223            | (BtreeI128(_), BtreeI128(_))
1224            | (BtreeU256(_), BtreeU256(_))
1225            | (BtreeI256(_), BtreeI256(_))
1226            | (BtreeString(_), BtreeString(_))
1227            | (BtreeAV(_), BtreeAV(_)) => Ok(()),
1228            // For unique indices, we'll need to see if everything in `other` can be added to `idx`.
1229            (UniqueBtreeBool(idx), UniqueBtreeBool(other)) => idx.can_merge(other, ignore).map_err(|ptr| *ptr),
1230            (UniqueBtreeU8(idx), UniqueBtreeU8(other)) => idx.can_merge(other, ignore).map_err(|ptr| *ptr),
1231            (UniqueBtreeSumTag(idx), UniqueBtreeSumTag(other)) => idx.can_merge(other, ignore).map_err(|ptr| *ptr),
1232            (UniqueBtreeI8(idx), UniqueBtreeI8(other)) => idx.can_merge(other, ignore).map_err(|ptr| *ptr),
1233            (UniqueBtreeU16(idx), UniqueBtreeU16(other)) => idx.can_merge(other, ignore).map_err(|ptr| *ptr),
1234            (UniqueBtreeI16(idx), UniqueBtreeI16(other)) => idx.can_merge(other, ignore).map_err(|ptr| *ptr),
1235            (UniqueBtreeU32(idx), UniqueBtreeU32(other)) => idx.can_merge(other, ignore).map_err(|ptr| *ptr),
1236            (UniqueBtreeI32(idx), UniqueBtreeI32(other)) => idx.can_merge(other, ignore).map_err(|ptr| *ptr),
1237            (UniqueBtreeU64(idx), UniqueBtreeU64(other)) => idx.can_merge(other, ignore).map_err(|ptr| *ptr),
1238            (UniqueBtreeI64(idx), UniqueBtreeI64(other)) => idx.can_merge(other, ignore).map_err(|ptr| *ptr),
1239            (UniqueBtreeU128(idx), UniqueBtreeU128(other)) => idx.can_merge(other, ignore).map_err(|ptr| *ptr),
1240            (UniqueBtreeI128(idx), UniqueBtreeI128(other)) => idx.can_merge(other, ignore).map_err(|ptr| *ptr),
1241            (UniqueBtreeU256(idx), UniqueBtreeU256(other)) => idx.can_merge(other, ignore).map_err(|ptr| *ptr),
1242            (UniqueBtreeI256(idx), UniqueBtreeI256(other)) => idx.can_merge(other, ignore).map_err(|ptr| *ptr),
1243            (UniqueBtreeString(idx), UniqueBtreeString(other)) => idx.can_merge(other, ignore).map_err(|ptr| *ptr),
1244            (UniqueBtreeAV(idx), UniqueBtreeAV(other)) => idx.can_merge(other, ignore).map_err(|ptr| *ptr),
1245            (UniqueDirectU8(idx), UniqueDirectU8(other)) => idx.can_merge(other, ignore),
1246            (UniqueDirectSumTag(idx), UniqueDirectSumTag(other)) => idx.can_merge(other, ignore),
1247            (UniqueDirectU16(idx), UniqueDirectU16(other)) => idx.can_merge(other, ignore),
1248            (UniqueDirectU32(idx), UniqueDirectU32(other)) => idx.can_merge(other, ignore),
1249            (UniqueDirectU64(idx), UniqueDirectU64(other)) => idx.can_merge(other, ignore),
1250
1251            _ => unreachable!("non-matching index kinds"),
1252        }
1253    }
1254
1255    /// Deletes all entries from the index, leaving it empty.
1256    ///
1257    /// When inserting a newly-created index into the committed state,
1258    /// we clear the tx state's index and insert it,
1259    /// rather than constructing a new `TableIndex`.
1260    pub fn clear(&mut self) {
1261        self.idx.clear();
1262        self.num_key_bytes = 0;
1263        self.num_rows = 0;
1264    }
1265
1266    /// The number of unique keys in this index.
1267    pub fn num_keys(&self) -> usize {
1268        self.idx.num_keys()
1269    }
1270
1271    /// The number of rows stored in this index.
1272    ///
1273    /// Note that, for non-unique indexes, this may be larger than [`Self::num_keys`].
1274    ///
1275    /// This method runs in constant time.
1276    pub fn num_rows(&self) -> u64 {
1277        self.num_rows
1278    }
1279
1280    /// The number of bytes stored in keys in this index.
1281    ///
1282    /// For non-unique indexes, duplicate keys are counted once for each row that refers to them,
1283    /// even though the internal storage may deduplicate them as an optimization.
1284    ///
1285    /// This method runs in constant time.
1286    ///
1287    /// See the [`KeySize`] trait for more details on how this method computes its result.
1288    pub fn num_key_bytes(&self) -> u64 {
1289        self.num_key_bytes
1290    }
1291}
1292
1293#[cfg(test)]
1294mod test {
1295    use super::*;
1296    use crate::page_pool::PagePool;
1297    use crate::{blob_store::HashMapBlobStore, table::test::table};
1298    use core::ops::Bound::*;
1299    use proptest::prelude::*;
1300    use proptest::{collection::vec, test_runner::TestCaseResult};
1301    use spacetimedb_data_structures::map::HashMap;
1302    use spacetimedb_primitives::ColId;
1303    use spacetimedb_sats::{
1304        product,
1305        proptest::{generate_product_value, generate_row_type},
1306        AlgebraicType, ProductType, ProductValue,
1307    };
1308    use spacetimedb_schema::def::BTreeAlgorithm;
1309
1310    fn gen_cols(ty_len: usize) -> impl Strategy<Value = ColList> {
1311        vec((0..ty_len as u16).prop_map_into::<ColId>(), 1..=ty_len)
1312            .prop_map(|cols| cols.into_iter().collect::<ColList>())
1313    }
1314
1315    fn gen_row_and_cols() -> impl Strategy<Value = (ProductType, ColList, ProductValue)> {
1316        generate_row_type(1..16).prop_flat_map(|ty| {
1317            (
1318                Just(ty.clone()),
1319                gen_cols(ty.elements.len()),
1320                generate_product_value(ty),
1321            )
1322        })
1323    }
1324
1325    fn new_index(row_type: &ProductType, cols: &ColList, is_unique: bool) -> TableIndex {
1326        let algo = BTreeAlgorithm { columns: cols.clone() }.into();
1327        TableIndex::new(row_type, &algo, is_unique).unwrap()
1328    }
1329
1330    /// Extracts from `row` the relevant column values according to what columns are indexed.
1331    fn get_fields(cols: &ColList, row: &ProductValue) -> AlgebraicValue {
1332        row.project(cols).unwrap()
1333    }
1334
1335    /// Returns whether indexing `row` again would violate a unique constraint, if any.
1336    fn violates_unique_constraint(index: &TableIndex, cols: &ColList, row: &ProductValue) -> bool {
1337        !index.is_unique() || index.contains_any(&get_fields(cols, row))
1338    }
1339
1340    /// Returns an iterator over the rows that would violate the unique constraint of this index,
1341    /// if `row` were inserted,
1342    /// or `None`, if this index doesn't have a unique constraint.
1343    fn get_rows_that_violate_unique_constraint<'a>(
1344        index: &'a TableIndex,
1345        row: &'a AlgebraicValue,
1346    ) -> Option<TableIndexPointIter<'a>> {
1347        index.is_unique().then(|| index.seek_point(row))
1348    }
1349
1350    proptest! {
1351        #![proptest_config(ProptestConfig { max_shrink_iters: 0x10000000, ..Default::default() })]
1352        #[test]
1353        fn remove_nonexistent_noop(((ty, cols, pv), is_unique) in (gen_row_and_cols(), any::<bool>())) {
1354            let mut index = new_index(&ty, &cols, is_unique);
1355            let mut table = table(ty);
1356            let pool = PagePool::new_for_test();
1357            let mut blob_store = HashMapBlobStore::default();
1358            let row_ref = table.insert(&pool, &mut blob_store, &pv).unwrap().1;
1359            prop_assert_eq!(index.delete(row_ref).unwrap(), false);
1360            prop_assert!(index.idx.is_empty());
1361        }
1362
1363        #[test]
1364        fn insert_delete_noop(((ty, cols, pv), is_unique) in (gen_row_and_cols(), any::<bool>())) {
1365            let mut index = new_index(&ty, &cols, is_unique);
1366            let mut table = table(ty);
1367            let pool = PagePool::new_for_test();
1368            let mut blob_store = HashMapBlobStore::default();
1369            let row_ref = table.insert(&pool, &mut blob_store, &pv).unwrap().1;
1370            let value = get_fields(&cols, &pv);
1371
1372            prop_assert_eq!(index.idx.len(), 0);
1373            prop_assert_eq!(index.contains_any(&value), false);
1374
1375            prop_assert_eq!(unsafe { index.check_and_insert(row_ref) }, Ok(()));
1376            prop_assert_eq!(index.idx.len(), 1);
1377            prop_assert_eq!(index.contains_any(&value), true);
1378
1379            prop_assert_eq!(index.delete(row_ref).unwrap(), true);
1380            prop_assert_eq!(index.idx.len(), 0);
1381            prop_assert_eq!(index.contains_any(&value), false);
1382        }
1383
1384        #[test]
1385        fn insert_again_violates_unique_constraint((ty, cols, pv) in gen_row_and_cols()) {
1386            let mut index = new_index(&ty, &cols, true);
1387            let mut table = table(ty);
1388            let pool = PagePool::new_for_test();
1389            let mut blob_store = HashMapBlobStore::default();
1390            let row_ref = table.insert(&pool, &mut blob_store, &pv).unwrap().1;
1391            let value = get_fields(&cols, &pv);
1392
1393            // Nothing in the index yet.
1394            prop_assert_eq!(index.idx.len(), 0);
1395            prop_assert_eq!(violates_unique_constraint(&index, &cols, &pv), false);
1396            prop_assert_eq!(
1397                get_rows_that_violate_unique_constraint(&index, &value).unwrap().collect::<Vec<_>>(),
1398                []
1399            );
1400
1401            // Insert.
1402            // SAFETY: `row_ref` has the same type as was passed in when constructing `index`.
1403            prop_assert_eq!(unsafe { index.check_and_insert(row_ref) }, Ok(()));
1404
1405            // Inserting again would be a problem.
1406            prop_assert_eq!(index.idx.len(), 1);
1407            prop_assert_eq!(violates_unique_constraint(&index, &cols, &pv), true);
1408            prop_assert_eq!(
1409                get_rows_that_violate_unique_constraint(&index, &value).unwrap().collect::<Vec<_>>(),
1410                [row_ref.pointer()]
1411            );
1412            // SAFETY: `row_ref` has the same type as was passed in when constructing `index`.
1413            prop_assert_eq!(unsafe { index.check_and_insert(row_ref) }, Err(row_ref.pointer()));
1414        }
1415
1416        #[test]
1417        fn seek_various_ranges(needle in 1..u64::MAX) {
1418            use AlgebraicValue::U64 as V;
1419
1420            let cols = 0.into();
1421            let ty = ProductType::from_iter([AlgebraicType::U64]);
1422            let mut index = new_index(&ty, &cols, true);
1423            let mut table = table(ty);
1424            let pool = PagePool::new_for_test();
1425            let mut blob_store = HashMapBlobStore::default();
1426
1427            let prev = needle - 1;
1428            let next = needle + 1;
1429            let range = prev..=next;
1430
1431            let mut val_to_ptr = HashMap::default();
1432
1433            // Insert `prev`, `needle`, and `next`.
1434            for x in range.clone() {
1435                let row = product![x];
1436                let row_ref = table.insert(&pool, &mut blob_store, &row).unwrap().1;
1437                val_to_ptr.insert(x, row_ref.pointer());
1438                // SAFETY: `row_ref` has the same type as was passed in when constructing `index`.
1439                prop_assert_eq!(unsafe { index.check_and_insert(row_ref) }, Ok(()));
1440            }
1441
1442            fn test_seek(index: &TableIndex, val_to_ptr: &HashMap<u64, RowPointer>, range: impl RangeBounds<AlgebraicValue>, expect: impl IntoIterator<Item = u64>) -> TestCaseResult {
1443                let mut ptrs_in_index = index.seek_range(&range).collect::<Vec<_>>();
1444                ptrs_in_index.sort();
1445                let mut expected_ptrs = expect.into_iter().map(|expected| val_to_ptr.get(&expected).unwrap()).copied().collect::<Vec<_>>();
1446                expected_ptrs.sort();
1447                prop_assert_eq!(
1448                    ptrs_in_index,
1449                    expected_ptrs
1450                );
1451                Ok(())
1452            }
1453
1454            // Test point ranges.
1455            for x in range.clone() {
1456                test_seek(&index, &val_to_ptr, V(x), [x])?;
1457            }
1458
1459            // Test `..` (`RangeFull`).
1460            test_seek(&index, &val_to_ptr, .., [prev, needle, next])?;
1461
1462            // Test `x..` (`RangeFrom`).
1463            test_seek(&index, &val_to_ptr, V(prev).., [prev, needle, next])?;
1464            test_seek(&index, &val_to_ptr, V(needle).., [needle, next])?;
1465            test_seek(&index, &val_to_ptr, V(next).., [next])?;
1466
1467            // Test `..x` (`RangeTo`).
1468            test_seek(&index, &val_to_ptr, ..V(prev), [])?;
1469            test_seek(&index, &val_to_ptr, ..V(needle), [prev])?;
1470            test_seek(&index, &val_to_ptr, ..V(next), [prev, needle])?;
1471
1472            // Test `..=x` (`RangeToInclusive`).
1473            test_seek(&index, &val_to_ptr, ..=V(prev), [prev])?;
1474            test_seek(&index, &val_to_ptr, ..=V(needle), [prev, needle])?;
1475            test_seek(&index, &val_to_ptr, ..=V(next), [prev, needle, next])?;
1476
1477            // Test `x..y` (`Range`).
1478            test_seek(&index, &val_to_ptr, V(prev)..V(prev), [])?;
1479            test_seek(&index, &val_to_ptr, V(prev)..V(needle), [prev])?;
1480            test_seek(&index, &val_to_ptr, V(prev)..V(next), [prev, needle])?;
1481            test_seek(&index, &val_to_ptr, V(needle)..V(next), [needle])?;
1482
1483            // Test `x..=y` (`RangeInclusive`).
1484            test_seek(&index, &val_to_ptr, V(prev)..=V(prev), [prev])?;
1485            test_seek(&index, &val_to_ptr, V(prev)..=V(needle), [prev, needle])?;
1486            test_seek(&index, &val_to_ptr, V(prev)..=V(next), [prev, needle, next])?;
1487            test_seek(&index, &val_to_ptr, V(needle)..=V(next), [needle, next])?;
1488            test_seek(&index, &val_to_ptr, V(next)..=V(next), [next])?;
1489
1490            // Test `(x, y]` (Exclusive start, inclusive end).
1491            test_seek(&index, &val_to_ptr, (Excluded(V(prev)), Included(V(prev))), [])?;
1492            test_seek(&index, &val_to_ptr, (Excluded(V(prev)), Included(V(needle))), [needle])?;
1493            test_seek(&index, &val_to_ptr, (Excluded(V(prev)), Included(V(next))), [needle, next])?;
1494
1495            // Test `(x, inf]` (Exclusive start, unbounded end).
1496            test_seek(&index, &val_to_ptr, (Excluded(V(prev)), Unbounded), [needle, next])?;
1497            test_seek(&index, &val_to_ptr, (Excluded(V(needle)), Unbounded), [next])?;
1498            test_seek(&index, &val_to_ptr, (Excluded(V(next)), Unbounded), [])?;
1499
1500            // Test `(x, y)` (Exclusive start, exclusive end).
1501            test_seek(&index, &val_to_ptr, (Excluded(V(prev)), Excluded(V(needle))), [])?;
1502            test_seek(&index, &val_to_ptr, (Excluded(V(prev)), Excluded(V(next))), [needle])?;
1503        }
1504    }
1505}