Skip to main content

spacetime_bindings_macro_input/
table.rs

1use super::sats;
2use super::sats::SatsField;
3use super::sym;
4use super::util::{check_duplicate, check_duplicate_msg, match_meta};
5use core::slice;
6use ident_case::RenameRule;
7use proc_macro2::{Span, TokenStream};
8use quote::quote;
9use syn::meta::ParseNestedMeta;
10use syn::parse::Parse;
11use syn::parse::Parser as _;
12use syn::punctuated::Punctuated;
13use syn::spanned::Spanned;
14use syn::{Ident, Path, Token};
15
16pub struct TableArgs {
17    pub access: Option<TableAccess>,
18    pub scheduled: Option<ScheduledArg>,
19    pub accessor: Ident,
20    pub indices: Vec<IndexArg>,
21    pub event: Option<()>,
22}
23
24pub enum TableAccess {
25    Public(Span),
26    Private(Span),
27}
28
29pub struct ScheduledArg {
30    pub span: Span,
31    pub reducer_or_procedure: Path,
32    pub at: Option<Ident>,
33}
34
35pub struct IndexArg {
36    pub accessor: Ident,
37    pub is_unique: bool,
38    pub kind: IndexType,
39}
40
41impl IndexArg {
42    fn new(accessor: Ident, kind: IndexType) -> Self {
43        // We don't know if its unique yet.
44        // We'll discover this once we have collected constraints.
45        let is_unique = false;
46        Self {
47            accessor,
48            is_unique,
49            kind,
50        }
51    }
52}
53
54pub enum IndexType {
55    BTree { columns: Vec<Ident> },
56    Hash { columns: Vec<Ident> },
57    Direct { column: Ident },
58}
59
60impl TableArgs {
61    pub fn parse(input: TokenStream, item: &syn::DeriveInput) -> syn::Result<Self> {
62        let mut access = None;
63        let mut scheduled = None;
64        let mut accessor = None;
65        let mut indices = Vec::new();
66        let mut event = None;
67
68        syn::meta::parser(|meta| {
69            match_meta!(match meta {
70                sym::public => {
71                    check_duplicate_msg(&access, &meta, "already specified access level")?;
72                    access = Some(TableAccess::Public(meta.path.span()));
73                }
74                sym::private => {
75                    check_duplicate_msg(&access, &meta, "already specified access level")?;
76                    access = Some(TableAccess::Private(meta.path.span()));
77                }
78                sym::name => {}
79                sym::accessor => {
80                    check_duplicate(&accessor, &meta)?;
81                    let value = meta.value()?;
82                    accessor = Some(value.parse()?);
83                }
84                sym::index => indices.push(IndexArg::parse_meta(meta)?),
85                sym::scheduled => {
86                    check_duplicate(&scheduled, &meta)?;
87                    scheduled = Some(ScheduledArg::parse_meta(meta)?);
88                }
89                sym::event => {
90                    check_duplicate(&event, &meta)?;
91                    event = Some(());
92                }
93            });
94            Ok(())
95        })
96        .parse2(input)?;
97        let accessor: Ident = accessor.ok_or_else(|| {
98            let table = RenameRule::SnakeCase.apply_to_field(item.ident.to_string());
99            syn::Error::new(
100                Span::call_site(),
101                format_args!(
102                    "must specify table accessor, e.g. `#[spacetimedb::table(accessor = {table})]"
103                ),
104            )
105        })?;
106
107        Ok(TableArgs {
108            access,
109            scheduled,
110            accessor,
111            indices,
112            event,
113        })
114    }
115}
116
117impl ScheduledArg {
118    fn parse_meta(meta: ParseNestedMeta) -> syn::Result<Self> {
119        let span = meta.path.span();
120        let mut reducer_or_procedure = None;
121        let mut at = None;
122
123        meta.parse_nested_meta(|meta| {
124            if meta.input.peek(syn::Token![=]) || meta.input.peek(syn::token::Paren) {
125                match_meta!(match meta {
126                    sym::at => {
127                        check_duplicate(&at, &meta)?;
128                        let ident = meta.value()?.parse()?;
129                        at = Some(ident);
130                    }
131                })
132            } else {
133                check_duplicate_msg(
134                    &reducer_or_procedure,
135                    &meta,
136                    "can only specify one scheduled reducer or procedure",
137                )?;
138                reducer_or_procedure = Some(meta.path);
139            }
140            Ok(())
141        })?;
142
143        let reducer_or_procedure = reducer_or_procedure.ok_or_else(|| {
144            meta.error(
145                "must specify scheduled reducer or procedure associated with the table: scheduled(function_name)",
146            )
147        })?;
148        Ok(Self {
149            span,
150            reducer_or_procedure,
151            at,
152        })
153    }
154}
155
156impl IndexArg {
157    fn parse_meta(meta: ParseNestedMeta) -> syn::Result<Self> {
158        let mut accessor = None;
159        let mut algo = None;
160
161        meta.parse_nested_meta(|meta| {
162            match_meta!(match meta {
163                sym::name => {}
164                sym::accessor => {
165                    check_duplicate(&accessor, &meta)?;
166                    accessor = Some(meta.value()?.parse()?);
167                }
168                sym::btree => {
169                    check_duplicate_msg(&algo, &meta, "index algorithm specified twice")?;
170                    algo = Some(Self::parse_btree(meta)?);
171                }
172                sym::hash => {
173                    check_duplicate_msg(&algo, &meta, "index algorithm specified twice")?;
174                    algo = Some(Self::parse_hash(meta)?);
175                }
176                sym::direct => {
177                    check_duplicate_msg(&algo, &meta, "index algorithm specified twice")?;
178                    algo = Some(Self::parse_direct(meta)?);
179                }
180            });
181            Ok(())
182        })?;
183        let accessor = accessor
184            .ok_or_else(|| meta.error("missing index accessor, e.g. accessor = my_index"))?;
185        let kind = algo.ok_or_else(|| {
186            meta.error("missing index algorithm, e.g., `btree(columns = [col1, col2])`, `hash(columns = [col1, col2])` or `direct(column = col1)`")
187        })?;
188
189        Ok(IndexArg::new(accessor, kind))
190    }
191
192    fn parse_columns(meta: &ParseNestedMeta) -> syn::Result<Option<Vec<Ident>>> {
193        let mut columns = None;
194        meta.parse_nested_meta(|meta| {
195            match_meta!(match meta {
196                sym::columns => {
197                    check_duplicate(&columns, &meta)?;
198                    let value = meta.value()?;
199                    let inner;
200                    syn::bracketed!(inner in value);
201                    columns = Some(
202                        Punctuated::<Ident, Token![,]>::parse_terminated(&inner)?
203                            .into_iter()
204                            .collect::<Vec<_>>(),
205                    );
206                }
207            });
208            Ok(())
209        })?;
210        Ok(columns)
211    }
212
213    fn parse_btree(meta: ParseNestedMeta) -> syn::Result<IndexType> {
214        let columns = Self::parse_columns(&meta)?;
215        let columns = columns.ok_or_else(|| {
216            meta.error("must specify columns for btree index, e.g. `btree(columns = [col1, col2])`")
217        })?;
218        Ok(IndexType::BTree { columns })
219    }
220
221    fn parse_hash(meta: ParseNestedMeta) -> syn::Result<IndexType> {
222        let columns = Self::parse_columns(&meta)?;
223        let columns = columns.ok_or_else(|| {
224            meta.error("must specify columns for hash index, e.g. `hash(columns = [col1, col2])`")
225        })?;
226        Ok(IndexType::Hash { columns })
227    }
228
229    fn parse_direct(meta: ParseNestedMeta) -> syn::Result<IndexType> {
230        let mut column = None;
231        meta.parse_nested_meta(|meta| {
232            match_meta!(match meta {
233                sym::column => {
234                    check_duplicate(&column, &meta)?;
235                    let value = meta.value()?;
236                    let inner;
237                    syn::bracketed!(inner in value);
238                    column = Some(Ident::parse(&inner)?);
239                }
240            });
241            Ok(())
242        })?;
243        let column = column.ok_or_else(|| {
244            meta.error("must specify the column for direct index, e.g. `direct(column = col1)`")
245        })?;
246        Ok(IndexType::Direct { column })
247    }
248
249    /// Parses an inline `#[index(btree)]`, `#[index(hash)]` or `#[index(direct)]` attribute on a field.
250    fn parse_index_attr(field: &Ident, attr: &syn::Attribute) -> syn::Result<Self> {
251        let mut kind = None;
252        attr.parse_nested_meta(|meta| {
253            match_meta!(match meta {
254                sym::btree => {
255                    check_duplicate_msg(&kind, &meta, "index type specified twice")?;
256                    kind = Some(IndexType::BTree {
257                        columns: vec![field.clone()],
258                    });
259                }
260                sym::hash => {
261                    check_duplicate_msg(&kind, &meta, "index type specified twice")?;
262                    kind = Some(IndexType::Hash {
263                        columns: vec![field.clone()],
264                    });
265                }
266                sym::direct => {
267                    check_duplicate_msg(&kind, &meta, "index type specified twice")?;
268                    kind = Some(IndexType::Direct {
269                        column: field.clone(),
270                    })
271                }
272            });
273            Ok(())
274        })?;
275        let kind = kind.ok_or_else(|| {
276            syn::Error::new_spanned(
277                &attr.meta,
278                "must specify kind of index (`btree`, `hash` or `direct`)",
279            )
280        })?;
281        let accessor = field.clone();
282        Ok(IndexArg::new(accessor, kind))
283    }
284}
285
286pub struct ColumnArgs<'a> {
287    pub original_struct_name: Ident,
288    pub fields: Vec<SatsField<'a>>,
289    pub columns: Vec<Column<'a>>,
290    pub unique_columns: Vec<Column<'a>>,
291    pub sequenced_columns: Vec<Column<'a>>,
292    pub primary_key_column: Option<Column<'a>>,
293}
294
295impl<'a> ColumnArgs<'a> {
296    pub fn parse(
297        mut table: TableArgs,
298        item: &'a syn::DeriveInput,
299    ) -> syn::Result<(TableArgs, Self)> {
300        let sats_ty = sats::sats_type_from_derive(item, quote!(spacetimedb::spacetimedb_lib))?;
301
302        let original_struct_name = sats_ty.ident.clone();
303
304        let sats::SatsTypeData::Product(fields) = &sats_ty.data else {
305            return Err(syn::Error::new(
306                Span::call_site(),
307                "spacetimedb table must be a struct",
308            ));
309        };
310
311        if fields.len() > u16::MAX.into() {
312            return Err(syn::Error::new_spanned(
313                item,
314                "too many columns; the most a table can have is 2^16",
315            ));
316        }
317
318        let mut columns: Vec<Column<'a>> = vec![];
319        let mut unique_columns: Vec<Column<'a>> = vec![];
320        let mut sequenced_columns: Vec<Column<'a>> = vec![];
321        let mut primary_key_column: Option<Column<'a>> = None;
322
323        for (i, field) in fields.iter().enumerate() {
324            let col_num = i as u16;
325            let field_ident = field.ident.unwrap();
326
327            let mut unique = None;
328            let mut auto_inc = None;
329            let mut primary_key = None;
330            let mut default_value = None;
331
332            for attr in field.original_attrs {
333                let Some(attr) = ColumnAttr::parse(attr, field_ident)? else {
334                    continue;
335                };
336                match attr {
337                    ColumnAttr::Unique(span) => {
338                        check_duplicate(&unique, span)?;
339                        unique = Some(span);
340                    }
341                    ColumnAttr::AutoInc(span) => {
342                        check_duplicate(&auto_inc, span)?;
343                        auto_inc = Some(span);
344                    }
345                    ColumnAttr::PrimaryKey(span) => {
346                        check_duplicate(&primary_key, span)?;
347                        primary_key = Some(span);
348                    }
349                    ColumnAttr::Index(index_arg) => table.indices.push(index_arg),
350                    ColumnAttr::Default(expr, span) => {
351                        check_duplicate(&default_value, span)?;
352                        default_value = Some(expr);
353                    }
354                }
355            }
356
357            let column: Column<'a> = Column {
358                index: col_num,
359                ident: field_ident,
360                vis: field.vis,
361                ty: field.ty,
362                default_value,
363            };
364
365            if unique.is_some() || primary_key.is_some() {
366                unique_columns.push(column.clone());
367            }
368            if auto_inc.is_some() {
369                sequenced_columns.push(column.clone());
370            }
371            if let Some(span) = primary_key {
372                check_duplicate_msg(
373                    &primary_key_column,
374                    span,
375                    "can only have one primary key per table",
376                )?;
377                primary_key_column = Some(column.clone());
378            }
379
380            columns.push(column.clone());
381        }
382
383        // Mark all indices with a single column matching a unique constraint as unique.
384        // For all the unpaired unique columns, create a unique index.
385        for unique_col in &unique_columns {
386            if table.indices.iter_mut().any(|index| {
387                let covered_by_index = match &index.kind {
388                    IndexType::BTree { columns } => columns == slice::from_ref(unique_col.ident),
389                    IndexType::Hash { columns } => columns == slice::from_ref(unique_col.ident),
390                    IndexType::Direct { column } => column == unique_col.ident,
391                };
392                index.is_unique |= covered_by_index;
393                covered_by_index
394            }) {
395                continue;
396            }
397            // NOTE(centril): We pick `btree` here if the user does not specify otherwise,
398            // as it's the safest choice of index for the general case,
399            // even if isn't optimal in specific cases.
400            let accessor = unique_col.ident.clone();
401            let columns = vec![accessor.clone()];
402            table.indices.push(IndexArg {
403                accessor,
404                is_unique: true,
405                kind: IndexType::BTree { columns },
406            })
407        }
408
409        Ok((
410            table,
411            ColumnArgs {
412                original_struct_name,
413                fields: fields.to_vec(),
414                columns,
415                unique_columns,
416                sequenced_columns,
417                primary_key_column,
418            },
419        ))
420    }
421}
422
423#[derive(Clone)]
424pub struct Column<'a> {
425    pub index: u16,
426    pub vis: &'a syn::Visibility,
427    pub ident: &'a syn::Ident,
428    pub ty: &'a syn::Type,
429    pub default_value: Option<syn::Expr>,
430}
431
432pub enum ColumnAttr {
433    Unique(Span),
434    AutoInc(Span),
435    PrimaryKey(Span),
436    Index(IndexArg),
437    Default(syn::Expr, Span),
438}
439
440impl ColumnAttr {
441    fn parse(attr: &syn::Attribute, field_ident: &Ident) -> syn::Result<Option<Self>> {
442        let Some(ident) = attr.path().get_ident() else {
443            return Ok(None);
444        };
445        Ok(if ident == sym::index {
446            let index = IndexArg::parse_index_attr(field_ident, attr)?;
447            Some(ColumnAttr::Index(index))
448        } else if ident == sym::unique {
449            attr.meta.require_path_only()?;
450            Some(ColumnAttr::Unique(ident.span()))
451        } else if ident == sym::auto_inc {
452            attr.meta.require_path_only()?;
453            Some(ColumnAttr::AutoInc(ident.span()))
454        } else if ident == sym::primary_key {
455            attr.meta.require_path_only()?;
456            Some(ColumnAttr::PrimaryKey(ident.span()))
457        } else if ident == sym::default {
458            Some(ColumnAttr::Default(
459                attr.parse_args::<syn::Expr>()?,
460                ident.span(),
461            ))
462        } else {
463            None
464        })
465    }
466}