Plain enums as index keys with specialized indices (#2506)

Author: Centril

crates/bindings-csharp/Codegen.Tests/fixtures/diag/snapshots/Module#FFI.verified.cs

@@ -52,23 +52,6 @@
       }
     },
     {/*
-    [PrimaryKey]
-    public TestEnumWithExplicitValues PrimaryKeyField;
-                                      ^^^^^^^^^^^^^^^
-}
-*/
-      Message: Field PrimaryKeyField is marked as Unique but it has a type TestEnumWithExplicitValues which is not an equatable primitive.,
-      Severity: Error,
-      Descriptor: {
-        Id: STDB0003,
-        Title: Unique fields must be equatable,
-        MessageFormat: Field {0} is marked as Unique but it has a type {1} which is not an equatable primitive.,
-        Category: SpacetimeDB,
-        DefaultSeverity: Error,
-        IsEnabledByDefault: true
-      }
-    },
-    {/*
 [SpacetimeDB.Table]
 public partial record TestTableTaggedEnum : SpacetimeDB.TaggedEnum<(int X, int Y)> { }
                                           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
@@ -429,4 +412,4 @@ public partial struct TestScheduleIssues
       }
     }
   ]
-}
+}

crates/bindings-csharp/Codegen.Tests/fixtures/diag/snapshots/Module.verified.txt

@@ -96,9 +96,32 @@ or SpecialType.System_Int64
             diag.Report(ErrorDescriptor.AutoIncNotInteger, field);
         }
 
+        // Check whether this is a sum type without a payload.
+        var isAllUnitEnum = false;
+        if (type.TypeKind == Microsoft.CodeAnalysis.TypeKind.Enum)
+        {
+            isAllUnitEnum = true;
+        }
+        else if (type.BaseType?.OriginalDefinition.ToString() == "SpacetimeDB.TaggedEnum<Variants>")
+        {
+            if (
+                type.BaseType.TypeArguments.FirstOrDefault() is INamedTypeSymbol
+                {
+                    IsTupleType: true,
+                    TupleElements: var taggedEnumVariants
+                }
+            )
+            {
+                isAllUnitEnum = taggedEnumVariants.All(
+                    (field) => field.Type.ToString() == "SpacetimeDB.Unit"
+                );
+            }
+        }
+
         IsEquatable =
             (
                 isInteger
+                || isAllUnitEnum
                 || type.SpecialType switch
                 {
                     SpecialType.System_String or SpecialType.System_Boolean => true,

crates/bindings-csharp/Codegen/Module.cs

@@ -18,7 +18,6 @@ public interface IBTreeIndexBounds
 }
 
 public readonly struct Bound<T>(T min, T max)
-    where T : IEquatable<T>
 {
     public T Min => min;
     public T Max => max;
@@ -29,7 +28,6 @@ public readonly struct Bound<T>(T min, T max)
 }
 
 public readonly struct BTreeIndexBounds<T, TRW>(Bound<T> t) : IBTreeIndexBounds
-    where T : IEquatable<T>
     where TRW : struct, IReadWrite<T>
 {
     public ushort PrefixElems => 0;
@@ -50,7 +48,6 @@ public void REnd(BinaryWriter w)
 }
 
 public readonly struct BTreeIndexBounds<T, TRW, U, URW>((T t, Bound<U> u) b) : IBTreeIndexBounds
-    where U : IEquatable<U>
     where TRW : struct, IReadWrite<T>
     where URW : struct, IReadWrite<U>
 {
@@ -76,7 +73,6 @@ public void REnd(BinaryWriter w)
 
 public readonly struct BTreeIndexBounds<T, TRW, U, URW, V, VRW>((T t, U u, Bound<V> v) b)
     : IBTreeIndexBounds
-    where V : IEquatable<V>
     where TRW : struct, IReadWrite<T>
     where URW : struct, IReadWrite<U>
     where VRW : struct, IReadWrite<V>
@@ -105,7 +101,6 @@ public void REnd(BinaryWriter w)
 public readonly struct BTreeIndexBounds<T, TRW, U, URW, V, VRW, W, WRW>(
     (T t, U u, V v, Bound<W> w) b
 ) : IBTreeIndexBounds
-    where W : IEquatable<W>
     where TRW : struct, IReadWrite<T>
     where URW : struct, IReadWrite<U>
     where VRW : struct, IReadWrite<V>
@@ -136,7 +131,6 @@ public void REnd(BinaryWriter w)
 public readonly struct BTreeIndexBounds<T, TRW, U, URW, V, VRW, W, WRW, X, XRW>(
     (T t, U u, V v, W w, Bound<X> x) b
 ) : IBTreeIndexBounds
-    where X : IEquatable<X>
     where TRW : struct, IReadWrite<T>
     where URW : struct, IReadWrite<U>
     where VRW : struct, IReadWrite<V>
@@ -169,7 +163,6 @@ public void REnd(BinaryWriter w)
 public readonly struct BTreeIndexBounds<T, TRW, U, URW, V, VRW, W, WRW, X, XRW, Y, YRW>(
     (T t, U u, V v, W w, X x, Bound<Y> y) b
 ) : IBTreeIndexBounds
-    where Y : IEquatable<Y>
     where TRW : struct, IReadWrite<T>
     where URW : struct, IReadWrite<U>
     where VRW : struct, IReadWrite<V>
@@ -204,7 +197,6 @@ public void REnd(BinaryWriter w)
 public readonly struct BTreeIndexBounds<T, TRW, U, URW, V, VRW, W, WRW, X, XRW, Y, YRW, Z, ZRW>(
     (T t, U u, V v, W w, X x, Y y, Bound<Z> z) b
 ) : IBTreeIndexBounds
-    where Z : IEquatable<Z>
     where TRW : struct, IReadWrite<T>
     where URW : struct, IReadWrite<U>
     where VRW : struct, IReadWrite<V>
@@ -256,7 +248,6 @@ public readonly struct BTreeIndexBounds<
     A,
     ARW
 >((T t, U u, V v, W w, X x, Y y, Z z, Bound<A> a) b) : IBTreeIndexBounds
-    where A : IEquatable<A>
     where TRW : struct, IReadWrite<T>
     where URW : struct, IReadWrite<U>
     where VRW : struct, IReadWrite<V>
@@ -312,7 +303,6 @@ public readonly struct BTreeIndexBounds<
     B,
     BRW
 >((T t, U u, V v, W w, X x, Y y, Z z, A a, Bound<B> b) b) : IBTreeIndexBounds
-    where B : IEquatable<B>
     where TRW : struct, IReadWrite<T>
     where URW : struct, IReadWrite<U>
     where VRW : struct, IReadWrite<V>
@@ -372,7 +362,6 @@ public readonly struct BTreeIndexBounds<
     C,
     CRW
 >((T t, U u, V v, W w, X x, Y y, Z z, A a, B b, Bound<C> c) b) : IBTreeIndexBounds
-    where C : IEquatable<C>
     where TRW : struct, IReadWrite<T>
     where URW : struct, IReadWrite<U>
     where VRW : struct, IReadWrite<V>

crates/bindings-csharp/Runtime/Internal/Bounds.cs

@@ -85,7 +85,6 @@ out handle
 public abstract class UniqueIndex<Handle, Row, T, RW>(string name) : IndexBase<Row>(name)
     where Handle : ITableView<Handle, Row>
     where Row : IStructuralReadWrite, new()
-    where T : IEquatable<T>
     where RW : struct, BSATN.IReadWrite<T>
 {
     private static BTreeIndexBounds<T, RW> ToBounds(T key) => new(key);

crates/bindings-csharp/Runtime/Internal/IIndex.cs

@@ -143,6 +143,7 @@ pub(crate) fn derive_satstype(ty: &SatsType<'_>) -> TokenStream {
     let name = &ty.ident;
     let krate = &ty.krate;
 
+    let mut add_impls_for_plain_enum = false;
     let typ = match &ty.data {
         SatsTypeData::Product(fields) => {
             let fields = fields.iter().map(|field| {
@@ -166,6 +167,10 @@ 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_impls_for_plain_enum = variants.iter().all(|var| var.ty.is_none());
+
             let unit = syn::Type::Tuple(syn::TypeTuple {
                 paren_token: Default::default(),
                 elems: Default::default(),
@@ -186,8 +191,7 @@ pub(crate) fn derive_satstype(ty: &SatsType<'_>) -> TokenStream {
                     [#(#variants),*]
                 )
             )
-            // todo!()
-        } // syn::Data::Union(u) => return Err(syn::Error::new(u.union_token.span, "unions not supported")),
+        }
     };
 
     let mut sats_generics = ty.generics.clone();
@@ -205,7 +209,43 @@ pub(crate) fn derive_satstype(ty: &SatsType<'_>) -> TokenStream {
     }
     let (_, typeid_ty_generics, _) = typeid_generics.split_for_impl();
 
+    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();
+        // Assume that the type is `Copy`, as most all-unit enums will be.
+        let filterable_impl = quote! {
+            #[automatically_derived]
+            impl #impl_generics #krate::FilterableValue for #name #ty_generics #where_clause {
+                type Column = #name #ty_generics;
+            }
+            #[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_plain_enum_extras
+
         #[automatically_derived]
         impl #impl_generics #krate::SpacetimeType for #name #ty_generics #where_clause {
             fn make_type<S: #krate::sats::typespace::TypespaceBuilder>(__typespace: &mut S) -> #krate::sats::AlgebraicType {

crates/bindings-macro/src/sats.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-macro/src/table.rs

@@ -139,7 +139,7 @@ pub use spacetimedb_bindings_macro::client_visibility_filter;
 ///     for popular_user in by_popularity.filter((100, "a"..)) {
 ///         log::debug!("Popular user whose name starts with 'a': {:?}", popular_user);
 ///     }
-///     
+///
 ///     // For every `#[unique]` or `#[primary_key]` field,
 ///     // the table has an extra method that allows getting a
 ///     // corresponding `spacetimedb::UniqueColumn`.
@@ -359,7 +359,7 @@ pub use spacetimedb_bindings_macro::client_visibility_filter;
 /// ```ignore
 /// ctx.db.cities().latitude()
 /// ```
-///    
+///
 /// # Generated code
 ///
 /// For each `[table(name = {name})]` annotation on a type `{T}`, generates a struct
@@ -402,7 +402,7 @@ pub use spacetimedb_bindings_macro::client_visibility_filter;
 /// impl {name}Handle {
 ///     // For each `#[unique]` or `#[primary_key]` field `{field}` of type `{F}`:
 ///     fn {field}(&self) -> UniqueColumn<_, {F}, _> { /* ... */ };
-///     
+///
 ///     // For each named index `{index}` on fields of type `{(F1, ..., FN)}`:
 ///     fn {index}(&self) -> RangedIndex<_, {(F1, ..., FN)}, _>;
 /// }

crates/bindings/src/lib.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 {