Skip to main content

vecdb_derive/
lib.rs

1use proc_macro::TokenStream;
2use quote::quote;
3use syn::{Data, DataStruct, DeriveInput, Fields, parse_macro_input};
4
5/// Derives the `Bytes` trait for single-field tuple structs.
6///
7/// This macro enables custom wrapper types to work with `BytesVec`, `LZ4Vec`, `ZstdVec`,
8/// and other vecdb vector types that require the `Bytes` trait.
9///
10/// # Requirements
11///
12/// - Must be a tuple struct with exactly one field
13/// - The inner type must implement `Bytes`
14/// - Supports generic type parameters
15///
16/// # Generated Implementation
17///
18/// The derive generates a `Bytes` implementation that delegates to the inner type:
19///
20/// ```rust,ignore
21/// impl Bytes for Wrapper<T> where T: Bytes {
22///     type Array = <T as Bytes>::Array;
23///
24///     fn to_bytes(&self) -> Self::Array {
25///         self.0.to_bytes()
26///     }
27///     fn from_bytes(bytes: &[u8]) -> Result<Self> {
28///         Ok(Self(<T>::from_bytes(bytes)?))
29///     }
30/// }
31/// ```
32///
33/// # Example
34///
35/// ```rust,ignore
36/// use vecdb::{Bytes, BytesVec};
37///
38/// #[derive(Bytes)]
39/// struct UserId(u64);
40///
41/// #[derive(Bytes)]
42/// struct Timestamp<T>(T); // Generic types supported
43/// ```
44#[proc_macro_derive(Bytes)]
45pub fn derive_bytes(input: TokenStream) -> TokenStream {
46    let input = parse_macro_input!(input as DeriveInput);
47    let struct_name = &input.ident;
48    let generics = &input.generics;
49    let (impl_generics, ty_generics, where_clause) = generics.split_for_impl();
50
51    let inner_type = match &input.data {
52        Data::Struct(DataStruct {
53            fields: Fields::Unnamed(fields),
54            ..
55        }) if fields.unnamed.len() == 1 => &fields.unnamed[0].ty,
56        _ => {
57            return syn::Error::new_spanned(
58                &input.ident,
59                "Bytes can only be derived for single-field tuple structs",
60            )
61            .to_compile_error()
62            .into();
63        }
64    };
65
66    // Check if we have generic parameters
67    let has_generics = !generics.params.is_empty();
68
69    let expanded = if has_generics {
70        let where_clause = if where_clause.is_some() {
71            quote! { #where_clause #inner_type: ::vecdb::Bytes, }
72        } else {
73            quote! { where #inner_type: ::vecdb::Bytes, }
74        };
75
76        quote! {
77            impl #impl_generics ::vecdb::Bytes for #struct_name #ty_generics #where_clause {
78                type Array = <#inner_type as ::vecdb::Bytes>::Array;
79                const IS_NATIVE_LAYOUT: bool =
80                    <#inner_type as ::vecdb::Bytes>::IS_NATIVE_LAYOUT
81                    && ::core::mem::size_of::<Self>() == ::core::mem::size_of::<#inner_type>()
82                    && ::core::mem::align_of::<Self>() == ::core::mem::align_of::<#inner_type>();
83
84                fn to_bytes(&self) -> Self::Array {
85                    self.0.to_bytes()
86                }
87
88                fn from_bytes(bytes: &[u8]) -> ::vecdb::Result<Self> {
89                    Ok(Self(<#inner_type>::from_bytes(bytes)?))
90                }
91            }
92        }
93    } else {
94        quote! {
95            impl ::vecdb::Bytes for #struct_name {
96                type Array = <#inner_type as ::vecdb::Bytes>::Array;
97                const IS_NATIVE_LAYOUT: bool =
98                    <#inner_type as ::vecdb::Bytes>::IS_NATIVE_LAYOUT
99                    && ::core::mem::size_of::<Self>() == ::core::mem::size_of::<#inner_type>()
100                    && ::core::mem::align_of::<Self>() == ::core::mem::align_of::<#inner_type>();
101
102                fn to_bytes(&self) -> Self::Array {
103                    self.0.to_bytes()
104                }
105
106                fn from_bytes(bytes: &[u8]) -> ::vecdb::Result<Self> {
107                    Ok(Self(<#inner_type>::from_bytes(bytes)?))
108                }
109            }
110        }
111    };
112
113    TokenStream::from(expanded)
114}
115
116/// Derives the `Pco` trait for single-field tuple structs containing numeric types.
117///
118/// This macro enables custom wrapper types to work with `PcoVec` for compressed storage
119/// of numeric data using Pcodec compression.
120///
121/// # Requirements
122///
123/// - Must be a tuple struct with exactly one field
124/// - The inner type must implement `Pco` (numeric types: u16-u64, i16-i64, f32, f64)
125/// - Supports generic type parameters
126///
127/// # Generated Implementation
128///
129/// The derive generates three trait implementations:
130///
131/// 1. `Bytes` - For serialization (same as `#[derive(Bytes)]`)
132/// 2. `Pco` - Specifies the numeric type for compression
133/// 3. `TransparentPco` - Marker trait for transparent wrappers
134///
135/// ```rust,ignore
136/// impl Pco for Wrapper<T> where T: Pco + Bytes {
137///     type NumberType = <T as Pco>::NumberType;
138/// }
139///
140/// impl TransparentPco<<T as Pco>::NumberType> for Wrapper<T>
141/// where T: Pco + Bytes {}
142///
143/// impl Bytes for Wrapper<T> where T: Pco + Bytes {
144///     // ... same as Bytes derive
145/// }
146/// ```
147///
148/// The `NumberType` is automatically propagated from the inner type, ensuring the
149/// wrapper has the same compression characteristics.
150///
151/// # Example
152///
153/// ```rust,ignore
154/// use vecdb::{Pco, PcoVec};
155///
156/// #[derive(Pco)]
157/// struct Price(f64);
158///
159/// #[derive(Pco)]
160/// struct NumericWrapper<T>(T); // Generic types supported
161///
162/// // Nested generics work too
163/// #[derive(Pco)]
164/// struct Container<T>(NumericWrapper<T>);
165/// ```
166#[proc_macro_derive(Pco)]
167pub fn derive_pco(input: TokenStream) -> TokenStream {
168    let input = parse_macro_input!(input as DeriveInput);
169    let struct_name = &input.ident;
170    let generics = &input.generics;
171    let (impl_generics, ty_generics, where_clause) = generics.split_for_impl();
172
173    let inner_type = match &input.data {
174        Data::Struct(DataStruct {
175            fields: Fields::Unnamed(fields),
176            ..
177        }) if fields.unnamed.len() == 1 => &fields.unnamed[0].ty,
178        _ => {
179            return syn::Error::new_spanned(
180                &input.ident,
181                "Pco can only be derived for single-field tuple structs",
182            )
183            .to_compile_error()
184            .into();
185        }
186    };
187
188    // Check if we have generic parameters
189    let has_generics = !generics.params.is_empty();
190
191    let expanded = if has_generics {
192        // For generic types, we need both Pco and Bytes bounds because:
193        // - Pco trait requires the NumberType
194        // - We call to_bytes/from_bytes methods which require Bytes
195        let where_clause = if where_clause.is_some() {
196            quote! { #where_clause #inner_type: ::vecdb::Pco + ::vecdb::Bytes, }
197        } else {
198            quote! { where #inner_type: ::vecdb::Pco + ::vecdb::Bytes, }
199        };
200
201        quote! {
202            impl #impl_generics ::vecdb::Bytes for #struct_name #ty_generics #where_clause {
203                type Array = <#inner_type as ::vecdb::Bytes>::Array;
204                const IS_NATIVE_LAYOUT: bool =
205                    <#inner_type as ::vecdb::Bytes>::IS_NATIVE_LAYOUT
206                    && ::core::mem::size_of::<Self>() == ::core::mem::size_of::<#inner_type>()
207                    && ::core::mem::align_of::<Self>() == ::core::mem::align_of::<#inner_type>();
208
209                fn to_bytes(&self) -> Self::Array {
210                    self.0.to_bytes()
211                }
212
213                fn from_bytes(bytes: &[u8]) -> ::vecdb::Result<Self> {
214                    Ok(Self(<#inner_type>::from_bytes(bytes)?))
215                }
216            }
217
218            impl #impl_generics ::vecdb::TransparentPco<<#inner_type as ::vecdb::Pco>::NumberType> for #struct_name #ty_generics #where_clause {}
219
220            impl #impl_generics ::vecdb::Pco for #struct_name #ty_generics #where_clause {
221                type NumberType = <#inner_type as ::vecdb::Pco>::NumberType;
222            }
223        }
224    } else {
225        quote! {
226            impl ::vecdb::Bytes for #struct_name {
227                type Array = <#inner_type as ::vecdb::Bytes>::Array;
228                const IS_NATIVE_LAYOUT: bool =
229                    <#inner_type as ::vecdb::Bytes>::IS_NATIVE_LAYOUT
230                    && ::core::mem::size_of::<Self>() == ::core::mem::size_of::<#inner_type>()
231                    && ::core::mem::align_of::<Self>() == ::core::mem::align_of::<#inner_type>();
232
233                fn to_bytes(&self) -> Self::Array {
234                    self.0.to_bytes()
235                }
236
237                fn from_bytes(bytes: &[u8]) -> ::vecdb::Result<Self> {
238                    Ok(Self(<#inner_type>::from_bytes(bytes)?))
239                }
240            }
241
242            impl ::vecdb::TransparentPco<<#inner_type as ::vecdb::Pco>::NumberType> for #struct_name {}
243
244            impl ::vecdb::Pco for #struct_name {
245                type NumberType = <#inner_type as ::vecdb::Pco>::NumberType;
246            }
247        }
248    };
249
250    TokenStream::from(expanded)
251}