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