improve multi column index test coverage and necessary changes to get there

Author: Centril

crates/datastore/src/locking_tx_datastore/mut_tx.rs

@@ -1482,6 +1482,7 @@ impl MutTxId {
                     .map_err(|IndexCannotSeekRange| IndexError::IndexCannotSeekRange(index_id))?;
                 IndexScanPointOrRange::Range(iter)
             }
+            PointOrRange::Unsupported => return Err(IndexError::IndexCannotSeekRange(index_id).into()),
         };
         Ok((table_id, iter))
     }

crates/sats/src/algebraic_value/de.rs

@@ -1,5 +1,5 @@
 use crate::array_value::{ArrayValueIntoIter, ArrayValueIterCloned};
-use crate::{de, AlgebraicValue, SumValue};
+use crate::{de, AlgebraicValue, ProductValue, SumValue};
 use crate::{i256, u256};
 use derive_more::From;
 
@@ -23,6 +23,11 @@ impl ValueDeserializer {
         // SAFETY: The conversion is OK due to `repr(transparent)`.
         unsafe { &*(val as *const AlgebraicValue as *const ValueDeserializer) }
     }
+
+    pub fn from_product_ref(prod: &ProductValue) -> RefProductAccess<'_> {
+        let vals = prod.elements.iter();
+        RefProductAccess { vals }
+    }
 }
 
 /// Errors that can occur when deserializing the `AlgebraicValue`.
@@ -348,7 +353,7 @@ impl<'de> de::Deserializer<'de> for &'de ValueDeserializer {
 }
 
 /// Defines deserialization for [`&'de ValueDeserializer`] where product elements are in the input.
-struct RefProductAccess<'a> {
+pub struct RefProductAccess<'a> {
     /// The element values of the product as an iterator of borrowed values.
     vals: std::slice::Iter<'a, AlgebraicValue>,
 }

crates/sats/src/de/impls.rs

@@ -259,6 +259,17 @@ impl_deserialize!(
     de => Vec::<T>::validate(de)
 );
 
+/// The visitor merely valiates the slice.
+struct ValidatingSliceVisitor;
+
+impl<T: ToOwned + ?Sized> SliceVisitor<'_, T> for ValidatingSliceVisitor {
+    type Output = ();
+
+    fn visit<E: Error>(self, _: &T) -> Result<Self::Output, E> {
+        Ok(())
+    }
+}
+
 /// The visitor converts the slice to its owned version.
 struct OwnedSliceVisitor;
 
@@ -293,8 +304,16 @@ impl<const N: usize> SliceVisitor<'_, [u8]> for ByteArrayVisitor<N> {
     }
 }
 
-impl_deserialize!([] &'de str, de => de.deserialize_str(BorrowedSliceVisitor));
-impl_deserialize!([] &'de [u8], de => de.deserialize_bytes(BorrowedSliceVisitor));
+impl_deserialize!(
+    [] &'de str,
+    de => de.deserialize_str_slice(),
+    de => de.deserialize_str(ValidatingSliceVisitor)
+);
+impl_deserialize!(
+    [] &'de [u8],
+    de => de.deserialize_bytes(BorrowedSliceVisitor),
+    de => de.deserialize_bytes(ValidatingSliceVisitor)
+);
 
 /// The visitor returns the slice as-is and borrowed.
 pub(crate) struct BorrowedSliceVisitor;
@@ -610,7 +629,7 @@ impl<'de> DeserializeSeed<'de> for WithTypespace<'_, AlgebraicType> {
             AlgebraicType::U256 => u256::validate(de),
             AlgebraicType::F32 => f32::validate(de),
             AlgebraicType::F64 => f64::validate(de),
-            AlgebraicType::String => <Box<str>>::validate(de),
+            AlgebraicType::String => <&str>::validate(de),
         }
     }
 }

crates/table/proptest-regressions/table_index/mod.txt

@@ -10,3 +10,5 @@ cc c1e4c959a32f6ab8ef9c4e29d39a24ec47cb03524584606a7f1fa4563f0f8cca # shrinks to
 cc 4cb325be8b24c9efa5b1f20b9504d044d9dd110eb9e99355de4ca42f9cfc20b4 # shrinks to (ty, cols, pv) = (ProductType {None: Sum(SumType {"variant_0": Product(ProductType {})})}, [ColId(0)], ProductValue { elements: [Sum(SumValue { tag: 0, value: Product(ProductValue { elements: [] }) })] }), is_unique = false
 cc a166a3c619c7cae3938f4e0cfb4e7a96cddfbb7943efd0b74e8cbb99d7a1e6a8 # shrinks to (ty, cols, pv) = (ProductType {None: U8}, [ColId(0)], ProductValue { elements: [U8(0)] }), kind = Direct
 cc 05390c104810e7086fa5d3f3cac7f491a377ae6ba64431661fd94662e28d1fca # shrinks to (ty, cols, pv) = (ProductType {None: Sum(SumType {"variant_0": Product(ProductType {})})}, [ColId(0)], ProductValue { elements: [Sum(SumValue { tag: 0, value: Product(ProductValue { elements: [] }) })] }), kind = Direct
+cc 3b8115315ff39f3268c02ccd659b82444e6d0b12aae7e6c0feba956875dab5ab # shrinks to is_unique = false, (prefix_ty, prefix_val) = (Array(ArrayType { elem_ty: Product(ProductType {"field_0": String}) }), Array([ProductValue { elements: [String("")] }])), (middle_ty, [start, middle, end]) = (Bool, [Bool(false), Bool(false), Bool(false)]), (suffix_ty, suffix_val) = (Bool, Bool(false))
+cc aa8f01aec687cbe6ad36acc77c8c484ccd323a53dca132c6eb490104563d17af # shrinks to is_unique = false, index_kind = BTree, (prefix_ty, prefix_val) = (Bool, Bool(false)), (middle_ty, [included, excluded]) = (Product(ProductType {}), [Product(ProductValue { elements: [] }), Product(ProductValue { elements: [] })]), (suffix_ty, suffix_val) = (Bool, Bool(false))

crates/table/src/table_index/btree_index.rs

@@ -112,6 +112,8 @@ impl<K: Ord + KeySize> Index for BTreeIndex<K> {
         // `self.insert` always returns `Ok(_)`.
         Ok(())
     }
+
+    const IS_RANGED: bool = true;
 }
 
 impl<K: KeySize + Ord> BTreeIndex<K> {

crates/table/src/table_index/bytes_key.rs

@@ -5,12 +5,17 @@ use core::cmp::Ordering;
 use core::hash::{Hash, Hasher};
 use core::mem;
 use core::ops::Deref;
-use spacetimedb_lib::buffer::{CountWriter, TeeWriter};
+use spacetimedb_lib::de::ProductVisitor;
+use spacetimedb_lib::ser::{SerializeSeqProduct, Serializer as _};
 use spacetimedb_memory_usage::MemoryUsage;
 use spacetimedb_primitives::ColList;
+use spacetimedb_sats::algebraic_value::de::{ValueDeserializeError, ValueDeserializer};
 use spacetimedb_sats::bsatn::{DecodeError, Deserializer, Serializer};
+use spacetimedb_sats::buffer::{CountWriter, TeeWriter};
 use spacetimedb_sats::de::{DeserializeSeed, Error as _};
-use spacetimedb_sats::{i256, u256, AlgebraicType, AlgebraicValue, ProductTypeElement, Serialize as _, WithTypespace};
+use spacetimedb_sats::{
+    i256, u256, AlgebraicType, AlgebraicValue, ProductTypeElement, ProductValue, Serialize as _, WithTypespace,
+};
 
 /// A key for an all-primitive multi-column index
 /// serialized to a byte array.
@@ -418,6 +423,66 @@ impl<const N: usize> RangeCompatBytesKey<N> {
         Self::from_bytes_key(key, ty)
     }
 
+    /// Decodes `prefix`  in `AlgebraicValue` form to a [`RangeCompatBytesKey<N>`]
+    /// by serializing the prefix and massaging.
+    pub(super) fn from_algebraic_value_prefix(
+        prefix: &ProductValue,
+        prefix_types: &[ProductTypeElement],
+    ) -> Result<Self, ValueDeserializeError> {
+        // Validate the prefix.
+        WithTypespace::empty(prefix_types).validate_seq_product(ValueDeserializer::from_product_ref(prefix))?;
+
+        // Serialize the `prefix` and the `endpoint` over.
+        let bytes = BytesKey::via_serializer(|ser| {
+            prefix
+                .serialize(ser)
+                .expect("should've serialized to BSATN successfully");
+        });
+        let BytesKey { mut bytes, length } = bytes;
+
+        // Massage the bytes.
+        let mut slice = bytes.as_mut_slice();
+        for ty in prefix_types {
+            slice = Self::process_from_bytes_key(slice, &ty.algebraic_type);
+        }
+
+        Ok(Self::new(length as usize, bytes))
+    }
+
+    /// Decodes `prefix` and `endpoint` in `AlgebraicValue` form to a [`RangeCompatBytesKey<N>`]
+    /// by serializing over both and massaging if they fit into the key.
+    pub(super) fn from_algebraic_value_prefix_and_endpoint(
+        prefix: &ProductValue,
+        prefix_types: &[ProductTypeElement],
+        endpoint: &AlgebraicValue,
+        range_type: &AlgebraicType,
+    ) -> Result<Self, ValueDeserializeError> {
+        // Validate the values.
+        WithTypespace::empty(prefix_types).validate_seq_product(ValueDeserializer::from_product_ref(prefix))?;
+        WithTypespace::empty(range_type).validate(ValueDeserializer::from_ref(endpoint))?;
+
+        // Serialize the `prefix` and the `endpoint` over.
+        let bytes = BytesKey::via_serializer(|ser| {
+            (|| {
+                let mut ser = ser.serialize_seq_product(2)?;
+                ser.serialize_element(&prefix)?;
+                ser.serialize_element(&endpoint)?;
+                ser.end()
+            })()
+            .expect("should've serialized to BSATN successfully");
+        });
+        let BytesKey { mut bytes, length } = bytes;
+
+        // Massage the bytes.
+        let mut slice = bytes.as_mut_slice();
+        for ty in prefix_types {
+            slice = Self::process_from_bytes_key(slice, &ty.algebraic_type);
+        }
+        Self::process_from_bytes_key(slice, range_type);
+
+        Ok(Self::new(length as usize, bytes))
+    }
+
     /// Serializes `av` to a [`BytesKey<N>`].
     ///
     /// It's assumed that `av`

crates/table/src/table_index/hash_index.rs

@@ -118,6 +118,8 @@ impl<K: KeySize + Eq + Hash> Index for HashIndex<K> {
         // `self.insert` always returns `Ok(_)`.
         Ok(())
     }
+
+    const IS_RANGED: bool = false;
 }
 
 impl<K: KeySize + Eq + Hash> HashIndex<K> {

crates/table/src/table_index/index.rs

@@ -133,6 +133,9 @@ pub trait Index {
     /// The trait imposes no particular order.
     /// Implementations may provide a non-deterministic order.
     fn iter(&self) -> Self::Iter<'_>;
+
+    /// Whether the index is ranged or not.
+    const IS_RANGED: bool;
 }
 
 pub trait RangedIndex: Index {