Implement btree api on new abi

coolreader18 · coolreader18 · commit 1ffff4652d07 · 2024-09-19T14:43:28.000-05:00
diff --git a/crates/bindings-macro/src/lib.rs b/crates/bindings-macro/src/lib.rs
@@ -15,7 +15,7 @@ use heck::ToSnakeCase;
 use module::{derive_deserialize, derive_satstype, derive_serialize};
 use proc_macro::TokenStream as StdTokenStream;
 use proc_macro2::{Span, TokenStream};
-use quote::{format_ident, quote, quote_spanned};
+use quote::{format_ident, quote, quote_spanned, ToTokens};
 use std::collections::HashMap;
 use std::time::Duration;
 use syn::ext::IdentExt;
@@ -567,12 +567,8 @@ impl IndexArg {
         Ok(IndexArg { kind, name })
     }
 
-    fn to_desc_and_accessor(
-        &self,
-        index_index: u32,
-        cols: &[Column],
-    ) -> Result<(TokenStream, TokenStream), syn::Error> {
-        let (algo, accessor) = match &self.kind {
+    fn validate<'a>(&'a self, table_name: &str, cols: &'a [Column<'a>]) -> syn::Result<ValidatedIndex<'_>> {
+        let kind = match &self.kind {
             IndexType::BTree { columns } => {
                 let cols = columns
                     .iter()
@@ -585,28 +581,98 @@ impl IndexArg {
                     })
                     .collect::<syn::Result<Vec<_>>>()?;
 
+                ValidatedIndexType::BTree { cols }
+            }
+        };
+        let index_name = match &kind {
+            ValidatedIndexType::BTree { cols } => ([table_name, "btree"].into_iter())
+                .chain(cols.iter().map(|col| col.field.name.as_deref().unwrap()))
+                .collect::<Vec<_>>()
+                .join("_"),
+        };
+        Ok(ValidatedIndex {
+            index_name,
+            accessor_name: &self.name,
+            kind,
+        })
+    }
+}
+
+struct ValidatedIndex<'a> {
+    index_name: String,
+    accessor_name: &'a Ident,
+    kind: ValidatedIndexType<'a>,
+}
+
+enum ValidatedIndexType<'a> {
+    BTree { cols: Vec<&'a Column<'a>> },
+}
+
+impl ValidatedIndex<'_> {
+    fn desc(&self) -> TokenStream {
+        let algo = match &self.kind {
+            ValidatedIndexType::BTree { cols } => {
                 let col_ids = cols.iter().map(|col| col.index);
-                let algo = quote!(spacetimedb::table::IndexAlgo::BTree {
+                quote!(spacetimedb::table::IndexAlgo::BTree {
                     columns: &[#(#col_ids),*]
-                });
+                })
+            }
+        };
+        let index_name = &self.index_name;
+        let accessor_name = ident_to_litstr(self.accessor_name);
+        quote!(spacetimedb::table::IndexDesc {
+            name: #index_name,
+            accessor_name: #accessor_name,
+            algo: #algo,
+        })
+    }
 
-                let index_ident = &self.name;
+    fn accessor(&self, vis: &syn::Visibility, row_type_ident: &Ident) -> TokenStream {
+        match &self.kind {
+            ValidatedIndexType::BTree { cols } => {
+                let index_ident = self.accessor_name;
                 let col_tys = cols.iter().map(|col| col.ty);
-                let accessor = quote! {
-                    fn #index_ident(&self) -> spacetimedb::BTreeIndex<Self, (#(#col_tys,)*), #index_index> {
+                let doc_columns = cols
+                    .iter()
+                    .map(|col| {
+                        format!(
+                            "- [`{ident}`][{row_type_ident}#structfield.{ident}]: [`{ty}`]\n",
+                            ident = col.field.ident.unwrap(),
+                            ty = col.ty.to_token_stream()
+                        )
+                    })
+                    .collect::<String>();
+                let doc = format!(
+                    "Gets the `{index_ident}` [`BTreeIndex`][spacetimedb::BTreeIndex] as defined \
+                     on this table. \n\
+                     \n\
+                     This B-tree index is defined on the following columns, in order:\n\
+                     {doc_columns}"
+                );
+                quote! {
+                    #[doc = #doc]
+                    #vis fn #index_ident(&self) -> spacetimedb::BTreeIndex<Self, (#(#col_tys,)*), __indices::#index_ident> {
                         spacetimedb::BTreeIndex::__new()
                     }
-                };
+                }
+            }
+        }
+    }
 
-                (algo, accessor)
+    fn marker_type(&self) -> TokenStream {
+        let index_ident = self.accessor_name;
+        let index_name = &self.index_name;
+        quote! {
+            pub struct #index_ident;
+            impl spacetimedb::table::Index for #index_ident {
+                fn index_id() -> spacetimedb::table::IndexId {
+                    static INDEX_ID: std::sync::OnceLock<spacetimedb::table::IndexId> = std::sync::OnceLock::new();
+                    *INDEX_ID.get_or_init(|| {
+                        spacetimedb::sys::index_id_from_name(#index_name).unwrap()
+                    })
+                }
             }
-        };
-        let accessor_name = ident_to_litstr(&self.name);
-        let desc = quote!(spacetimedb::table::IndexDesc {
-            accessor_name: #accessor_name,
-            algo: #algo,
-        });
-        Ok((desc, accessor))
+        }
     }
 }
 
@@ -821,16 +887,15 @@ fn table_impl(mut args: TableArgs, mut item: MutItem<syn::DeriveInput>) -> syn::
 
     let row_type = quote!(#original_struct_ident);
 
-    let (indexes, index_accessors) = args
+    let indices = args
         .indices
         .iter()
-        .enumerate()
-        .map(|(i, index)| index.to_desc_and_accessor(i as u32, &columns))
-        // TODO: stabilized in 1.79
-        // .collect::<syn::Result<(Vec<_>, Vec<_>)>>()?;
-        .collect::<syn::Result<Vec<_>>>()?
-        .into_iter()
-        .unzip::<_, _, Vec<_>, Vec<_>>();
+        .map(|index| index.validate(&table_name, &columns))
+        .collect::<syn::Result<Vec<_>>>()?;
+
+    let index_descs = indices.iter().map(|index| index.desc());
+    let index_accessors = indices.iter().map(|index| index.accessor(vis, original_struct_ident));
+    let index_marker_types = indices.iter().map(|index| index.marker_type());
 
     let unique_field_accessors = unique_columns.iter().map(|unique| {
         let column_index = unique.index;
@@ -907,7 +972,7 @@ fn table_impl(mut args: TableArgs, mut item: MutItem<syn::DeriveInput>) -> syn::
             // the default value if not specified is Private
             #(const TABLE_ACCESS: spacetimedb::table::TableAccess = #table_access;)*
             const UNIQUE_COLUMNS: &'static [u16] = &[#(#unique_col_ids),*];
-            const INDEXES: &'static [spacetimedb::table::IndexDesc<'static>] = &[#(#indexes),*];
+            const INDEXES: &'static [spacetimedb::table::IndexDesc<'static>] = &[#(#index_descs),*];
             #(const PRIMARY_KEY: Option<u16> = Some(#primary_col_id);)*
             const SEQUENCES: &'static [u16] = &[#(#sequence_col_ids),*];
             #(const SCHEDULED_REDUCER_NAME: Option<&'static str> = Some(<#scheduled_reducer_ident as spacetimedb::rt::ReducerInfo>::NAME);)*
@@ -1006,6 +1071,11 @@ fn table_impl(mut args: TableArgs, mut item: MutItem<syn::DeriveInput>) -> syn::
 
             #tabletype_impl
 
+            #[allow(non_camel_case_types)]
+            mod __indices {
+                #(#index_marker_types)*
+            }
+
             #describe_table_func
         };
 
diff --git a/crates/bindings/src/table.rs b/crates/bindings/src/table.rs
@@ -4,8 +4,9 @@ use std::marker::PhantomData;
 use std::{fmt, ops};
 
 use spacetimedb_lib::buffer::{BufReader, Cursor};
+
 pub use spacetimedb_lib::db::raw_def::v9::TableAccess;
-use spacetimedb_primitives::ColId;
+pub use spacetimedb_primitives::{ColId, IndexId};
 
 use crate::{bsatn, sys, DeserializeOwned, IterBuf, Serialize, SpacetimeType, TableId};
 
@@ -132,6 +133,7 @@ pub trait TableInternal: Sized {
 /// Describe a named index with an index type over a set of columns identified by their IDs.
 #[derive(Clone, Copy)]
 pub struct IndexDesc<'a> {
+    pub name: &'a str,
     pub accessor_name: &'a str,
     pub algo: IndexAlgo<'a>,
 }
@@ -337,11 +339,15 @@ where
     }
 }
 
-pub struct BTreeIndex<Tbl: Table, IndexType, const INDEX_INDEX: u32> {
-    _marker: PhantomData<(Tbl, IndexType)>,
+pub trait Index {
+    fn index_id() -> IndexId;
+}
+
+pub struct BTreeIndex<Tbl: Table, IndexType, Idx: Index> {
+    _marker: PhantomData<(Tbl, IndexType, Idx)>,
 }
 
-impl<Tbl: Table, IndexType, const INDEX_INDEX: u32> BTreeIndex<Tbl, IndexType, INDEX_INDEX> {
+impl<Tbl: Table, IndexType, Idx: Index> BTreeIndex<Tbl, IndexType, Idx> {
     #[doc(hidden)]
     pub fn __new() -> Self {
         Self { _marker: PhantomData }
@@ -353,13 +359,16 @@ impl<Tbl: Table, IndexType, const INDEX_INDEX: u32> BTreeIndex<Tbl, IndexType, I
     /// - A value for the first indexed column.
     /// - A range of values for the first indexed column.
     /// - A tuple of values for any prefix of the indexed columns, optionally terminated by a range for the next.
-    pub fn filter<B: BTreeIndexBounds<IndexType, K>, K>(&self, b: B) -> impl Iterator<Item = Tbl::Row> {
+    pub fn filter<B, K>(&self, b: B) -> impl Iterator<Item = Tbl::Row>
+    where
+        B: BTreeIndexBounds<IndexType, K>,
+    {
+        let index_id = Idx::index_id();
         let args = b.get_args();
         let (prefix, prefix_elems, rstart, rend) = args.args_for_syscall();
-        #[allow(unreachable_code)]
-        TableIter::new(todo!(
-            "once implemented: datastore_btree_scan_bsatn({prefix:?}, {prefix_elems:?}, {rstart:?}, {rend:?})"
-        ))
+        let iter = sys::datastore_btree_scan_bsatn(index_id, prefix, prefix_elems, rstart, rend)
+            .unwrap_or_else(|e| panic!("unexpected error from datastore_btree_scan_bsatn: {e}"));
+        TableIter::new(iter)
     }
 
     /// Deletes all rows in the database state where the indexed column(s) match the bounds `b`.
@@ -371,10 +380,16 @@ impl<Tbl: Table, IndexType, const INDEX_INDEX: u32> BTreeIndex<Tbl, IndexType, I
     ///
     /// May panic if deleting any one of the rows would violate a constraint,
     /// though as of proposing no such constraints exist.
-    pub fn delete<B: BTreeIndexBounds<IndexType, K>, K>(&self, b: B) -> u64 {
+    pub fn delete<B, K>(&self, b: B) -> u64
+    where
+        B: BTreeIndexBounds<IndexType, K>,
+    {
+        let index_id = Idx::index_id();
         let args = b.get_args();
         let (prefix, prefix_elems, rstart, rend) = args.args_for_syscall();
-        todo!("once implemented: datastore_delete_by_btree_scan_bsatn({prefix:?}, {prefix_elems:?}, {rstart:?}, {rend:?})")
+        sys::datastore_delete_by_btree_scan_bsatn(index_id, prefix, prefix_elems, rstart, rend)
+            .unwrap_or_else(|e| panic!("unexpected error from datastore_delete_by_btree_scan_bsatn: {e}"))
+            .into()
     }
 }
 
@@ -394,13 +409,14 @@ pub struct BTreeScanArgs {
 
 impl BTreeScanArgs {
     pub(crate) fn args_for_syscall(&self) -> (&[u8], ColId, &[u8], &[u8]) {
-        let len = self.data.len();
-        (
-            &self.data[..self.rstart_idx],
-            ColId::from(self.prefix_elems),
-            &self.data[self.rstart_idx..self.rend_idx.unwrap_or(len)],
-            &self.data[self.rend_idx.unwrap_or(self.rstart_idx)..],
-        )
+        let prefix = &self.data[..self.rstart_idx];
+        let (rstart, rend) = if let Some(rend_idx) = self.rend_idx {
+            (&self.data[self.rstart_idx..rend_idx], &self.data[rend_idx..])
+        } else {
+            let elem = &self.data[self.rstart_idx..];
+            (elem, elem)
+        };
+        (prefix, ColId::from(self.prefix_elems), rstart, rend)
     }
 }
 
@@ -464,14 +480,10 @@ impl<T: Serialize> TermBound<&T> {
             TermBound::Single(elem) => (elem, None),
             TermBound::Range(start, end) => (start, Some(end)),
         };
-        let serialize_bound = |_buf: &mut Vec<u8>, _bound: &ops::Bound<&T>| {
-            // bsatn::to_writer(buf, bound).unwrap();
-            todo!();
-        };
-        serialize_bound(buf, start);
+        bsatn::to_writer(buf, start).unwrap();
         end.map(|end| {
             let rend_idx = buf.len();
-            serialize_bound(buf, end);
+            bsatn::to_writer(buf, end).unwrap();
             rend_idx
         })
     }
