Fix defining multiple tables on the same struct in Rust modules (#2103

)

crates/bindings-macro/src/lib.rs

@@ -7,13 +7,13 @@ mod sats;
 mod table;
 mod util;
 
-use crate::util::cvt_attr;
 use proc_macro::TokenStream as StdTokenStream;
 use proc_macro2::TokenStream;
 use quote::quote;
 use std::time::Duration;
-use syn::parse::ParseStream;
 use syn::ItemFn;
+use syn::{parse::ParseStream, Attribute};
+use util::{cvt_attr, ok_or_compile_error};
 
 mod sym {
     /// A symbol known at compile-time against
@@ -159,6 +159,22 @@ pub fn reducer(args: StdTokenStream, item: StdTokenStream) -> StdTokenStream {
     })
 }
 
+/// It turns out to be shockingly difficult to construct an [`Attribute`].
+/// That type is not [`Parse`], instead having two distinct methods
+/// for parsing "inner" vs "outer" attributes.
+///
+/// We need this [`Attribute`] in [`table`] so that we can "pushnew" it
+/// onto the end of a list of attributes. See comments within [`table`].
+fn derive_table_helper_attr() -> Attribute {
+    let source = quote!(#[derive(spacetimedb::__TableHelper)]);
+
+    syn::parse::Parser::parse2(Attribute::parse_outer, source)
+        .unwrap()
+        .into_iter()
+        .next()
+        .unwrap()
+}
+
 /// Generates code for treating this struct type as a table.
 ///
 /// Among other things, this derives `Serialize`, `Deserialize`,
@@ -233,18 +249,46 @@ pub fn reducer(args: StdTokenStream, item: StdTokenStream) -> StdTokenStream {
 #[proc_macro_attribute]
 pub fn table(args: StdTokenStream, item: StdTokenStream) -> StdTokenStream {
     // put this on the struct so we don't get unknown attribute errors
-    let extra_attr = quote!(#[derive(spacetimedb::__TableHelper)]);
-    cvt_attr::<syn::DeriveInput>(args, item, extra_attr, |args, item| {
-        let args = table::TableArgs::parse(args, &item.ident)?;
-        table::table_impl(args, item)
+    let derive_table_helper: syn::Attribute = derive_table_helper_attr();
+
+    ok_or_compile_error(|| {
+        let item = TokenStream::from(item);
+        let mut derive_input: syn::DeriveInput = syn::parse2(item.clone())?;
+
+        // Add `derive(__TableHelper)` only if it's not already in the attributes of the `derive_input.`
+        // If multiple `#[table]` attributes are applied to the same `struct` item,
+        // this will ensure that we don't emit multiple conflicting implementations
+        // for traits like `SpacetimeType`, `Serialize` and `Deserialize`.
+        //
+        // We need to push at the end, rather than the beginning,
+        // because rustc expands attribute macros (including derives) top-to-bottom,
+        // and we need *all* `#[table]` attributes *before* the `derive(__TableHelper)`.
+        // This way, the first `table` will insert a `derive(__TableHelper)`,
+        // and all subsequent `#[table]`s on the same `struct` will see it,
+        // and not add another.
+        //
+        // Note, thank goodness, that `syn`'s `PartialEq` impls (provided with the `extra-traits` feature)
+        // skip any [`Span`]s contained in the items,
+        // thereby comparing for syntactic rather than structural equality. This shouldn't matter,
+        // since we expect that the `derive_table_helper` will always have the same [`Span`]s,
+        // but it's nice to know.
+        if !derive_input.attrs.contains(&derive_table_helper) {
+            derive_input.attrs.push(derive_table_helper);
+        }
+
+        let args = table::TableArgs::parse(args.into(), &derive_input.ident)?;
+        let generated = table::table_impl(args, &derive_input)?;
+        Ok(TokenStream::from_iter([quote!(#derive_input), generated]))
     })
 }
 
+/// Special alias for `derive(SpacetimeType)`, aka [`schema_type`], for use by [`table`].
+///
 /// Provides helper attributes for `#[spacetimedb::table]`, so that we don't get unknown attribute errors.
 #[doc(hidden)]
 #[proc_macro_derive(__TableHelper, attributes(sats, unique, auto_inc, primary_key, index))]
-pub fn table_helper(_input: StdTokenStream) -> StdTokenStream {
-    Default::default()
+pub fn table_helper(input: StdTokenStream) -> StdTokenStream {
+    schema_type(input)
 }
 
 #[proc_macro]

crates/bindings-macro/src/table.rs

@@ -1,4 +1,4 @@
-use crate::sats::{self, derive_deserialize, derive_satstype, derive_serialize};
+use crate::sats;
 use crate::sym;
 use crate::util::{check_duplicate, check_duplicate_msg, ident_to_litstr, match_meta};
 use heck::ToSnakeCase;
@@ -543,10 +543,6 @@ pub(crate) fn table_impl(mut args: TableArgs, item: &syn::DeriveInput) -> syn::R
 
     let tablehandle_ident = format_ident!("{}__TableHandle", table_ident);
 
-    let deserialize_impl = derive_deserialize(&sats_ty);
-    let serialize_impl = derive_serialize(&sats_ty);
-    let schema_impl = derive_satstype(&sats_ty);
-
     // Generate `integrate_generated_columns`
     // which will integrate all generated auto-inc col values into `_row`.
     let integrate_gen_col = sequenced_columns.iter().map(|col| {
@@ -666,18 +662,14 @@ pub(crate) fn table_impl(mut args: TableArgs, item: &syn::DeriveInput) -> syn::R
         }
     };
 
-    let row_type_to_table = quote!(<#row_type as spacetimedb::table::__MapRowTypeToTable>::Table);
-
     // Output all macro data
     let trait_def = quote_spanned! {table_ident.span()=>
         #[allow(non_camel_case_types, dead_code)]
         #vis trait #table_ident {
-            fn #table_ident(&self) -> &#row_type_to_table;
+            fn #table_ident(&self) -> &#tablehandle_ident;
         }
         impl #table_ident for spacetimedb::Local {
-            fn #table_ident(&self) -> &#row_type_to_table {
-                #[allow(non_camel_case_types)]
-                type #tablehandle_ident = #row_type_to_table;
+            fn #table_ident(&self) -> &#tablehandle_ident {
                 &#tablehandle_ident {}
             }
         }
@@ -697,17 +689,9 @@ pub(crate) fn table_impl(mut args: TableArgs, item: &syn::DeriveInput) -> syn::R
 
         #trait_def
 
-        #[cfg(doc)]
         #tablehandle_def
 
         const _: () = {
-            #[cfg(not(doc))]
-            #tablehandle_def
-
-            impl spacetimedb::table::__MapRowTypeToTable for #row_type {
-                type Table = #tablehandle_ident;
-            }
-
             impl #tablehandle_ident {
                 #(#index_accessors)*
             }
@@ -723,10 +707,6 @@ pub(crate) fn table_impl(mut args: TableArgs, item: &syn::DeriveInput) -> syn::R
 
             #describe_table_func
         };
-
-        #schema_impl
-        #deserialize_impl
-        #serialize_impl
     };
 
     if std::env::var("PROC_MACRO_DEBUG").is_ok() {

crates/bindings-macro/src/util.rs

@@ -21,6 +21,16 @@ pub(crate) fn cvt_attr<Item: Parse + quote::ToTokens>(
     TokenStream::from_iter([extra_attr, item, generated]).into()
 }
 
+/// Run `f`, converting `Err` returns into a compile error.
+///
+/// This helper allows code within the closure `f` to use `?` for early return.
+pub(crate) fn ok_or_compile_error<Res: Into<StdTokenStream>>(f: impl FnOnce() -> syn::Result<Res>) -> StdTokenStream {
+    match f() {
+        Ok(ok) => ok.into(),
+        Err(e) => e.into_compile_error().into(),
+    }
+}
+
 pub(crate) fn ident_to_litstr(ident: &Ident) -> syn::LitStr {
     syn::LitStr::new(&ident.to_string(), ident.span())
 }

crates/bindings/src/table.rs

@@ -149,11 +149,6 @@ pub struct ScheduleDesc<'a> {
     pub scheduled_at_column: u16,
 }
 
-#[doc(hidden)]
-pub trait __MapRowTypeToTable {
-    type Table: Table;
-}
-
 /// A UNIQUE constraint violation on a table was attempted.
 // TODO: add column name for better error message
 #[derive(Debug)]