derive Key

crates/bonsaidb-core/src/key.rs

@@ -18,6 +18,8 @@ use serde::{Deserialize, Serialize};
 
 use crate::{connection::Range, AnyError};
 
+pub use bonsaidb_macros::Key;
+
 /// A trait that enables a type to convert itself into a `memcmp`-compatible
 /// sequence of bytes.
 pub trait KeyEncoding<'k, K>: std::fmt::Debug + Send + Sync

crates/bonsaidb-macros/Cargo.toml

@@ -19,8 +19,9 @@ proc-macro = true
 attribute-derive = "0.2.2"
 proc-macro-crate = "1.1.0"
 proc-macro-error = "1"
-proc-macro2 = "1.0.36"
+proc-macro2 = { version = "1.0.37", features = ["nightly"] }
 quote = "1"
+quote-use = { version = "0.3.2", features = ["namespace_idents"] }
 syn = "1"
 trybuild = "1.0.54"
 

crates/bonsaidb-macros/src/lib.rs

@@ -16,10 +16,12 @@ use attribute_derive::Attribute;
 use proc_macro2::{Span, TokenStream};
 use proc_macro_crate::{crate_name, FoundCrate};
 use proc_macro_error::{abort, abort_call_site, proc_macro_error, ResultExt};
-use quote::{quote, ToTokens};
+use quote::ToTokens;
+use quote_use::{format_ident_namespaced as format_ident, quote_use as quote};
 use syn::{
-    parse_macro_input, parse_quote, punctuated::Punctuated, token::Paren, DeriveInput, Expr, Ident,
-    LitStr, Path, Type, TypeTuple,
+    parse_macro_input, parse_quote, punctuated::Punctuated, token::Paren, DataEnum, DataStruct,
+    DeriveInput, Expr, Field, Fields, FieldsNamed, FieldsUnnamed, Ident, Index, LitStr, Path,
+    Token, Type, TypeTuple, Variant,
 };
 
 fn core_path() -> Path {
@@ -179,12 +181,12 @@ pub fn collection_derive(input: proc_macro::TokenStream) -> proc_macro::TokenStr
                 if #core::ENCRYPTION_ENABLED {
                     #encryption_key
                 } else {
-                    ::core::option::Option::None
+                    None
                 }
             }
         };
         quote! {
-            fn encryption_key() -> ::core::option::Option<#core::document::KeyId> {
+            fn encryption_key() -> Option<#core::document::KeyId> {
                 #encryption
             }
         }
@@ -197,9 +199,9 @@ pub fn collection_derive(input: proc_macro::TokenStream) -> proc_macro::TokenStr
             fn collection_name() -> #core::schema::CollectionName {
                 #name
             }
-            fn define_views(schema: &mut #core::schema::Schematic) -> ::core::result::Result<(), #core::Error> {
+            fn define_views(schema: &mut #core::schema::Schematic) -> Result<(), #core::Error> {
                 #( schema.define_view(#views)?; )*
-                ::core::result::Result::Ok(())
+                Ok(())
             }
             #encryption
         }
@@ -306,7 +308,7 @@ pub fn view_derive(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
 #[derive(Attribute)]
 #[attribute(ident = "schema")]
 #[attribute(
-    invalid_field = r#"Only `name = "name""`, `authority = "authority"`, `collections = [SomeCollection, AnotherCollection]` and `core = bonsaidb::core` are supported attributes"#
+    invalid_field = r#"Only `name = "name"`, `authority = "authority"`, `collections = [SomeCollection, AnotherCollection]` and `core = bonsaidb::core` are supported attributes"#
 )]
 struct SchemaAttribute {
     #[attribute(missing = r#"You need to specify the schema name via `#[schema(name = "name")]`"#)]
@@ -357,9 +359,222 @@ pub fn schema_derive(input: proc_macro::TokenStream) -> proc_macro::TokenStream
 
             fn define_collections(
                 schema: &mut #core::schema::Schematic
-            ) -> ::core::result::Result<(), #core::Error> {
+            ) -> Result<(), #core::Error> {
                 #( schema.define_collection::<#collections>()?; )*
-                ::core::result::Result::Ok(())
+                Ok(())
+            }
+        }
+    }
+    .into()
+}
+
+#[derive(Attribute)]
+#[attribute(ident = "key")]
+#[attribute(invalid_field = r#"Only `core = bonsaidb::core` is supported"#)]
+struct KeyAttribute {
+    #[attribute(expected = r#"Specify the the path to `core` like so: `core = bosaidb::core`"#)]
+    core: Option<Path>,
+    allow_null_bytes: bool,
+}
+
+/// Derives the `bonsaidb::core::key::Key` trait.
+///
+/// `#[schema(core = bonsaidb::core]`, `core` is optional
+#[proc_macro_error]
+#[proc_macro_derive(Key, attributes(schema))]
+pub fn key_derive(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
+    let DeriveInput {
+        attrs,
+        ident,
+        generics,
+        data,
+        ..
+    } = parse_macro_input!(input as DeriveInput);
+
+    let KeyAttribute {
+        core,
+        allow_null_bytes,
+    } = KeyAttribute::from_attributes(attrs).unwrap_or_abort();
+
+    let allow_null_bytes = if allow_null_bytes {
+        quote!($encoder.allow_null_bytes_in_variable_fields();)
+    } else {
+        quote!()
+    };
+
+    let core = core.unwrap_or_else(core_path);
+    let (_, ty_generics, _) = generics.split_for_impl();
+    let mut generics = generics.clone();
+    generics
+        .params
+        .push(syn::GenericParam::Lifetime(parse_quote!('key)));
+    let (impl_generics, _, where_clause) = generics.split_for_impl();
+
+    let (encode_fields, decode_fields): (TokenStream, TokenStream) = match data {
+        syn::Data::Struct(DataStruct {
+            fields: Fields::Named(FieldsNamed { named, .. }),
+            ..
+        }) => {
+            let (encode_fields, decode_fields): (TokenStream, TokenStream) = named
+                .into_iter()
+                .map(|Field { ident, .. }| {
+                    let ident = ident.expect("named fields have idents");
+                    (
+                        quote!($encoder.encode(&self.#ident)?;),
+                        quote!(#ident: $decoder.decode()?,),
+                    )
+                })
+                .unzip();
+            (encode_fields, quote!( Self { #decode_fields }))
+        }
+
+        syn::Data::Struct(DataStruct {
+            fields: Fields::Unnamed(FieldsUnnamed { unnamed, .. }),
+            ..
+        }) => {
+            let (encode_fields, decode_fields): (TokenStream, TokenStream) = unnamed
+                .into_iter()
+                .enumerate()
+                .map(|(idx, _)| {
+                    let idx = Index::from(idx);
+                    (
+                        quote!($encoder.encode(&self.#idx)?;),
+                        quote!($decoder.decode()?,),
+                    )
+                })
+                .unzip();
+            (encode_fields, quote!(Self(#decode_fields)))
+        }
+        syn::Data::Struct(DataStruct {
+            fields: Fields::Unit,
+            ..
+        }) => abort_call_site!("unit structs are not supported"),
+        syn::Data::Enum(DataEnum { variants, .. }) => {
+            let (encode_variants, decode_variants): (TokenStream, TokenStream) = variants
+                .into_iter()
+                .enumerate()
+                .map(|(idx, Variant { fields, ident, .. })| {
+                    let idx = idx as u64;
+                    match fields {
+                        Fields::Named(FieldsNamed { named, .. }) => {
+                            let (idents, (encode_fields, decode_fields)): (
+                                Punctuated<_, Token![,]>,
+                                (TokenStream, TokenStream),
+                            ) = named
+                                .into_iter()
+                                .map(|Field { ident, .. }| {
+                                    let ident = ident.expect("named fields have idents");
+                                    (
+                                        ident.clone(),
+                                        (
+                                            quote!($encoder.encode(#ident)?;),
+                                            quote!(#ident: $decoder.decode()?,),
+                                        ),
+                                    )
+                                })
+                                .unzip();
+                            (
+                                quote! {
+                                    Self::#ident{#idents} => {
+                                        $encoder.encode(&#idx)?;
+                                        #encode_fields
+                                    },
+                                },
+                                quote! {
+                                    #idx => Self::#ident{#decode_fields},
+                                },
+                            )
+                        }
+                        Fields::Unnamed(FieldsUnnamed { unnamed, .. }) => {
+                            let (idents, (encode_fields, decode_fields)): (
+                                Punctuated<_, Token![,]>,
+                                (TokenStream, TokenStream),
+                            ) = unnamed
+                                .into_iter()
+                                .enumerate()
+                                .map(|(idx, _)| {
+                                    let ident = format_ident!("$field_{idx}");
+                                    (
+                                        ident.clone(),
+                                        (
+                                            quote!($encoder.encode(#ident)?;),
+                                            quote!($decoder.decode()?,),
+                                        ),
+                                    )
+                                })
+                                .unzip();
+                            (
+                                quote! {
+                                    Self::#ident(#idents) => {
+                                        $encoder.encode(&#idx)?;
+                                        #encode_fields
+                                    },
+                                },
+                                quote! {
+                                    #idx => Self::#ident(#decode_fields),
+                                },
+                            )
+                        }
+                        Fields::Unit => (
+                            quote!(Self::#ident => $encoder.encode(&#idx)?,),
+                            quote!(#idx => Self::#ident,),
+                        ),
+                    }
+                })
+                .unzip();
+            (
+                quote! {
+                    match self{
+                        #encode_variants
+                    }
+                },
+                quote! {
+                    use std::io::{self, ErrorKind};
+
+                    match $decoder.decode::<u64>()? {
+                        #decode_variants
+                        _ => return Err(#core::key::CompositeKeyError::from(io::Error::from(
+                                ErrorKind::InvalidData,
+                        )))
+                    }
+                },
+            )
+        }
+        syn::Data::Union(_) => abort_call_site!("unions are not supported"),
+    };
+
+    quote! {
+        use std::borrow::Cow;
+        use std::io::{self, ErrorKind};
+
+        impl #impl_generics #core::key::Key<'key> for #ident #ty_generics #where_clause {
+
+            fn from_ord_bytes(mut $bytes: &'key [u8]) -> Result<Self, Self::Error> {
+
+                let mut $decoder = #core::key::CompositeKeyDecoder::new($bytes);
+
+                let $self_ = #decode_fields;
+
+                $decoder.finish()?;
+
+                Ok($self_)
+            }
+        }
+
+        impl #impl_generics #core::key::KeyEncoding<'key, Self> for #ident #ty_generics #where_clause {
+            type Error = #core::key::CompositeKeyError;
+
+            // TODO fixed width if possible
+            const LENGTH: Option<usize> = None;
+
+            fn as_ord_bytes(&'key self) -> Result<Cow<'key, [u8]>, Self::Error> {
+                let mut $encoder = #core::key::CompositeKeyEncoder::default();
+
+                #allow_null_bytes
+
+                #encode_fields
+
+                Ok(Cow::Owned($encoder.finish()))
             }
         }
     }

crates/bonsaidb-macros/tests/key.rs

@@ -0,0 +1,59 @@
+use bonsaidb::core::key::{Key, KeyEncoding};
+
+#[test]
+fn tuple_struct() {
+    #[derive(Clone, Debug, Key)]
+    struct Test(i32, i32, String);
+
+    assert_eq!(
+        &[0, 0, 0, 1, 0, 0, 0, 2, 116, 101, 115, 116, 0, 4],
+        Test(1, 2, "test".into())
+            .as_ord_bytes()
+            .unwrap()
+            .as_ref()
+    )
+}
+
+#[test]
+fn struct_struct() {
+    #[derive(Clone, Debug, Key)]
+    struct Test {
+        a: i32,
+        b: String,
+    }
+    assert_eq!(
+        &[255, 255, 255, 214, 109, 101, 97, 110, 105, 110, 103, 0, 7],
+        Test {
+            a: -42,
+            b: "meaning".into()
+        }
+        .as_ord_bytes()
+        .unwrap()
+        .as_ref()
+    )
+}
+
+#[test]
+fn r#enum() {
+    #[derive(Clone, Debug, Key)]
+    enum Test {
+        A,
+        B(i32, String),
+        C { a: String, b: i32 },
+    }
+
+    assert_eq!(
+        &[0, 0, 0, 0, 0, 0, 0, 0],
+        Test::A.as_ord_bytes().unwrap().as_ref()
+    );
+
+    assert_eq!(
+        &[0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 2, 97, 0, 1],
+        Test::B(2, "a".into()).as_ord_bytes().unwrap().as_ref()
+    );
+
+    assert_eq!(
+        &[0, 0, 0, 0, 0, 0, 0, 2, 98, 0, 0, 0, 0, 3, 1],
+        Test::C{a: "b".into(), b: 3}.as_ord_bytes().unwrap().as_ref()
+    )
+}

crates/bonsaidb-macros/tests/ui/schema/invalid_attribute.stderr

@@ -4,7 +4,7 @@ error: expected identifier
 5 | #[schema(name = "name", "hi")]
   |                         ^^^^
 
-error: Only `name = \"name\"\"`, `authority = \"authority\"`, `collections = [SomeCollection, AnotherCollection]` and `core = bonsaidb::core` are supported attributes
+error: Only `name = \"name\"`, `authority = \"authority\"`, `collections = [SomeCollection, AnotherCollection]` and `core = bonsaidb::core` are supported attributes
  --> tests/ui/schema/invalid_attribute.rs:9:25
   |
 9 | #[schema(name = "name", test = "hi")]