use specialized indices for sum tags

Author: Centril

crates/bindings-macro/src/sats.rs

@@ -143,7 +143,7 @@ pub(crate) fn derive_satstype(ty: &SatsType<'_>) -> TokenStream {
     let name = &ty.ident;
     let krate = &ty.krate;
 
-    let mut add_filterable_value = false;
+    let mut add_impls_for_plain_enum = false;
     let typ = match &ty.data {
         SatsTypeData::Product(fields) => {
             let fields = fields.iter().map(|field| {
@@ -169,7 +169,7 @@ pub(crate) fn derive_satstype(ty: &SatsType<'_>) -> TokenStream {
         SatsTypeData::Sum(variants) => {
             // To allow an enum, with all-unit variants, as an index key type,
             // add derive `Filterable` for the enum.
-            add_filterable_value = variants.iter().all(|var| var.ty.is_none());
+            add_impls_for_plain_enum = variants.iter().all(|var| var.ty.is_none());
 
             let unit = syn::Type::Tuple(syn::TypeTuple {
                 paren_token: Default::default(),
@@ -209,28 +209,40 @@ pub(crate) fn derive_satstype(ty: &SatsType<'_>) -> TokenStream {
     }
     let (_, typeid_ty_generics, _) = typeid_generics.split_for_impl();
 
-    let impl_filterable_value = if add_filterable_value {
+    let impl_plain_enum_extras = if add_impls_for_plain_enum {
         // These will mostly be empty as lifetime and type parameters must be constrained
         // but const parameters don't require constraining.
         let mut generics = ty.generics.clone();
         add_type_bounds(&mut generics, &quote!(#krate::FilterableValue));
         let (impl_generics, ty_generics, where_clause) = generics.split_for_impl();
-
         // As we don't have access to other `derive` attributes,
         // we don't know if `Copy` was derived,
         // so we won't impl for the owned type.
-        Some(quote! {
+        let filterable_impl = quote! {
             #[automatically_derived]
             impl #impl_generics #krate::FilterableValue for &#name #ty_generics #where_clause {
                 type Column = #name #ty_generics;
             }
+        };
+
+        let mut generics = ty.generics.clone();
+        add_type_bounds(&mut generics, &quote!(#krate::DirectIndexKey));
+        let (impl_generics, ty_generics, where_clause) = generics.split_for_impl();
+        let dik_impl = quote! {
+            #[automatically_derived]
+            impl #impl_generics #krate::DirectIndexKey for #name #ty_generics #where_clause {}
+        };
+
+        Some(quote! {
+            #filterable_impl
+            #dik_impl
         })
     } else {
         None
     };
 
     quote! {
-        #impl_filterable_value
+        #impl_plain_enum_extras
 
         #[automatically_derived]
         impl #impl_generics #krate::SpacetimeType for #name #ty_generics #where_clause {

crates/bindings-macro/src/table.rs

@@ -381,7 +381,7 @@ impl ValidatedIndex<'_> {
                 let col_ty = col.ty;
                 let typeck = quote_spanned!(col_ty.span()=>
                     const _: () = {
-                        spacetimedb::rt::assert_column_type_valid_for_direct_index::<#col_ty>();
+                        spacetimedb::spacetimedb_lib::assert_column_type_valid_for_direct_index::<#col_ty>();
                     };
                 );
                 (slice::from_ref(col), Some(typeck))

crates/bindings/src/rt.rs

@@ -346,26 +346,6 @@ pub fn register_table<T: Table>() {
     })
 }
 
-mod sealed_direct_index {
-    pub trait Sealed {}
-}
-#[diagnostic::on_unimplemented(
-    message = "column type must be a one of: `u8`, `u16`, `u32`, or `u64`",
-    label = "should be `u8`, `u16`, `u32`, or `u64`, not `{Self}`"
-)]
-pub trait DirectIndexKey: sealed_direct_index::Sealed {}
-impl sealed_direct_index::Sealed for u8 {}
-impl DirectIndexKey for u8 {}
-impl sealed_direct_index::Sealed for u16 {}
-impl DirectIndexKey for u16 {}
-impl sealed_direct_index::Sealed for u32 {}
-impl DirectIndexKey for u32 {}
-impl sealed_direct_index::Sealed for u64 {}
-impl DirectIndexKey for u64 {}
-
-/// Assert that `T` is a valid column to use direct index on.
-pub const fn assert_column_type_valid_for_direct_index<T: DirectIndexKey>() {}
-
 impl From<IndexAlgo<'_>> for RawIndexAlgorithm {
     fn from(algo: IndexAlgo<'_>) -> RawIndexAlgorithm {
         match algo {

crates/lib/src/direct_index_key.rs

@@ -0,0 +1,12 @@
+#[diagnostic::on_unimplemented(
+    message = "column type must be a one of: `u8`, `u16`, `u32`, `u64`, or plain `enum`",
+    label = "should be `u8`, `u16`, `u32`, `u64`, or plain `enum`, not `{Self}`"
+)]
+pub trait DirectIndexKey {}
+impl DirectIndexKey for u8 {}
+impl DirectIndexKey for u16 {}
+impl DirectIndexKey for u32 {}
+impl DirectIndexKey for u64 {}
+
+/// Assert that `T` is a valid column to use direct index on.
+pub const fn assert_column_type_valid_for_direct_index<T: DirectIndexKey>() {}

crates/lib/src/lib.rs

@@ -8,6 +8,7 @@ use std::collections::{btree_map, BTreeMap};
 
 pub mod connection_id;
 pub mod db;
+mod direct_index_key;
 pub mod error;
 mod filterable_value;
 pub mod identity;
@@ -27,6 +28,7 @@ pub mod type_value {
 }
 
 pub use connection_id::ConnectionId;
+pub use direct_index_key::{assert_column_type_valid_for_direct_index, DirectIndexKey};
 pub use filterable_value::{FilterableValue, IndexScanRangeBoundsTerminator, TermBound};
 pub use identity::Identity;
 pub use scheduler::ScheduleAt;

crates/sats/src/algebraic_value.rs

@@ -178,15 +178,13 @@ impl AlgebraicValue {
 
     /// Returns an [`AlgebraicValue`] representing a sum value with `tag` and `value`.
     pub fn sum(tag: u8, value: Self) -> Self {
-        let value = Box::new(value);
-        Self::Sum(SumValue { tag, value })
+        Self::Sum(SumValue::new(tag, value))
     }
 
     /// Returns an [`AlgebraicValue`] representing a sum value with `tag` and empty [AlgebraicValue::product], that is
     /// valid for simple enums without payload.
     pub fn enum_simple(tag: u8) -> Self {
-        let value = Box::new(AlgebraicValue::product(vec![]));
-        Self::Sum(SumValue { tag, value })
+        Self::Sum(SumValue::new_simple(tag))
     }
 
     /// Returns an [`AlgebraicValue`] representing a product value with the given `elements`.

crates/sats/src/convert.rs

@@ -1,3 +1,4 @@
+use crate::sum_value::SumTag;
 use crate::{i256, u256};
 use crate::{AlgebraicType, AlgebraicValue, ProductType, ProductValue};
 use spacetimedb_primitives::{ColId, ConstraintId, IndexId, ScheduleId, SequenceId, TableId};
@@ -25,6 +26,12 @@ impl From<AlgebraicType> for ProductType {
     }
 }
 
+impl From<()> for AlgebraicValue {
+    fn from((): ()) -> Self {
+        AlgebraicValue::unit()
+    }
+}
+
 macro_rules! built_in_into {
     ($native:ty, $kind:ident) => {
         impl From<$native> for AlgebraicValue {
@@ -45,6 +52,7 @@ built_in_into!(&str, String);
 built_in_into!(String, String);
 built_in_into!(&[u8], Bytes);
 built_in_into!(Box<[u8]>, Bytes);
+built_in_into!(SumTag, Sum);
 
 macro_rules! system_id {
     ($name:ident) => {

crates/sats/src/sum_value.rs

@@ -14,3 +14,27 @@ pub struct SumValue {
 impl crate::Value for SumValue {
     type Type = SumType;
 }
+
+impl SumValue {
+    /// Returns a new `SumValue` with the given `tag` and `value`.
+    pub fn new(tag: u8, value: impl Into<AlgebraicValue>) -> Self {
+        let value = Box::from(value.into());
+        Self { tag, value }
+    }
+
+    /// Returns a new `SumValue` with the given `tag` and unit value.
+    pub fn new_simple(tag: u8) -> Self {
+        Self::new(tag, ())
+    }
+}
+
+/// The tag of a `SumValue`.
+/// Can be used to read out the tag of a sum value without reading the payload.
+#[derive(Debug, Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Hash)]
+pub struct SumTag(pub u8);
+
+impl From<SumTag> for SumValue {
+    fn from(SumTag(tag): SumTag) -> Self {
+        SumValue::new_simple(tag)
+    }
+}

crates/schema/src/def/validate/v9.rs

@@ -625,9 +625,16 @@ impl TableValidator<'_, '_> {
                 let field = &self.product_type.elements[column.idx()];
                 let ty = &field.algebraic_type;
                 use AlgebraicType::*;
-                if let U8 | U16 | U32 | U64 = ty {
-                } else {
-                    return Err(ValidationError::DirectIndexOnNonUnsignedInt {
+                let is_bad_type = match ty {
+                    U8 | U16 | U32 | U64 => false,
+                    Ref(r) => self.module_validator.typespace[*r]
+                        .as_sum()
+                        .is_none_or(|s| !s.is_simple_enum()),
+                    Sum(sum) if sum.is_simple_enum() => false,
+                    _ => true,
+                };
+                if is_bad_type {
+                    return Err(ValidationError::DirectIndexOnBadType {
                         index: name.clone(),
                         column: field.name.clone().unwrap_or_else(|| column.idx().to_string().into()),
                         ty: ty.clone().into(),
@@ -1440,7 +1447,7 @@ mod tests {
             .finish();
         let result: Result<ModuleDef> = builder.finish().try_into();
 
-        expect_error_matching!(result, ValidationError::DirectIndexOnNonUnsignedInt { index, .. } => {
+        expect_error_matching!(result, ValidationError::DirectIndexOnBadType { index, .. } => {
             &index[..] == "Bananas_b_idx_direct"
         });
     }

crates/schema/src/error.rs

@@ -62,7 +62,7 @@ pub enum ValidationError {
     #[error("No index found to support unique constraint `{constraint}` for columns `{columns:?}`")]
     UniqueConstraintWithoutIndex { constraint: Box<str>, columns: ColSet },
     #[error("Direct index does not support type `{ty}` in column `{column}` in index `{index}`")]
-    DirectIndexOnNonUnsignedInt {
+    DirectIndexOnBadType {
         index: RawIdentifier,
         column: RawIdentifier,
         ty: PrettyAlgebraicType,

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

@@ -0,0 +1,10 @@
+# Seeds for failure cases proptest has generated in the past. It is
+# automatically read and these particular cases re-run before any
+# novel cases are generated.
+#
+# It is recommended to check this file in to source control so that
+# everyone who runs the test benefits from these saved cases.
+cc 471733dcdf80f2858344c15b6297070a312d2051ecf0eb56eaf9bf38cd51aab4 # shrinks to start = 38, end = 0
+cc 98241ddfa4ab59c91a264cbbd364619c1d02a912406c4402cd01fb0c10a72e10 # shrinks to start = 75, end = 0
+cc b058bb0beeaf3bb4f97854dff2bf835e2ea4eaffb1eb3c13c78652018b02b957 # shrinks to start = 130, end = 0, key = 0
+cc dc8757b81c7b7cffe9e4a9d43b75cd5814d5fb0667a0c1f401fd5f402fde3d97 # shrinks to start = 37, end = 37, key = 0

crates/table/src/memory_usage.rs

@@ -39,6 +39,12 @@ impl MemoryUsage for f64 {}
 impl MemoryUsage for spacetimedb_sats::F32 {}
 impl MemoryUsage for spacetimedb_sats::F64 {}
 
+impl<T: MemoryUsage + ?Sized> MemoryUsage for &T {
+    fn heap_usage(&self) -> usize {
+        (*self).heap_usage()
+    }
+}
+
 impl<T: MemoryUsage + ?Sized> MemoryUsage for Box<T> {
     fn heap_usage(&self) -> usize {
         mem::size_of_val::<T>(self) + T::heap_usage(self)
@@ -65,6 +71,12 @@ impl<T: MemoryUsage> MemoryUsage for [T] {
     }
 }
 
+impl<T: MemoryUsage, const N: usize> MemoryUsage for [T; N] {
+    fn heap_usage(&self) -> usize {
+        self.iter().map(T::heap_usage).sum()
+    }
+}
+
 impl MemoryUsage for str {}
 
 impl<T: MemoryUsage> MemoryUsage for Option<T> {

crates/table/src/read_column.rs

@@ -11,7 +11,9 @@ use crate::{
 };
 use spacetimedb_sats::{
     algebraic_value::{ser::ValueSerializer, Packed},
-    i256, u256, AlgebraicType, AlgebraicValue, ArrayValue, ProductType, ProductValue, SumValue,
+    i256,
+    sum_value::SumTag,
+    u256, AlgebraicType, AlgebraicValue, ArrayValue, ProductType, ProductValue, SumValue,
 };
 use std::{cell::Cell, mem};
 use thiserror::Error;
@@ -339,6 +341,29 @@ impl_read_column_via_from! {
     i256 => Box<i256>;
 }
 
+/// SAFETY: `is_compatible_type` only returns true for sum types,
+/// and any sum value stores the tag first in BFLATN.
+unsafe impl ReadColumn for SumTag {
+    fn is_compatible_type(ty: &AlgebraicTypeLayout) -> bool {
+        matches!(ty, AlgebraicTypeLayout::Sum(_))
+    }
+
+    unsafe fn unchecked_read_column(row_ref: RowRef<'_>, layout: &ProductTypeElementLayout) -> Self {
+        debug_assert!(Self::is_compatible_type(&layout.ty));
+
+        let (page, offset) = row_ref.page_and_offset();
+        let col_offset = offset + PageOffset(layout.offset);
+
+        let data = page.get_row_data(col_offset, Size(1));
+        let data: Result<[u8; 1], _> = data.try_into();
+        // SAFETY: `<[u8; 1] as TryFrom<&[u8]>` succeeds if and only if the slice's length is `1`.
+        // We used `1` as both the length of the slice and the array, so we know them to be equal.
+        let [data] = unsafe { data.unwrap_unchecked() };
+
+        Self(data)
+    }
+}
+
 #[cfg(test)]
 mod test {
     use super::*;
@@ -512,5 +537,21 @@ mod test {
 
         // Use a long string which will hit the blob store.
         read_column_long_string { AlgebraicType::String => Box<str> = "long string. ".repeat(2048).into() };
+
+        read_sum_value_plain { AlgebraicType::simple_enum(["a", "b"].into_iter()) => SumValue = SumValue::new_simple(1) };
+        read_sum_tag_plain { AlgebraicType::simple_enum(["a", "b"].into_iter()) => SumTag = SumTag(1) };
+    }
+
+    #[test]
+    fn read_sum_tag_from_sum_with_payload() {
+        let algebraic_type = AlgebraicType::sum([("a", AlgebraicType::U8), ("b", AlgebraicType::U16)]);
+
+        let mut blob_store = HashMapBlobStore::default();
+        let mut table = table(ProductType::from([algebraic_type]));
+
+        let val = SumValue::new(1, 42u16);
+        let (_, row_ref) = table.insert(&mut blob_store, &product![val.clone()]).unwrap();
+
+        assert_eq!(val.tag, row_ref.read_col::<SumTag>(0).unwrap().0);
     }
 }