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
80                fn to_bytes(&self) -> Self::Array {
81                    self.0.to_bytes()
82                }
83
84                fn from_bytes(bytes: &[u8]) -> ::vecdb::Result<Self> {
85                    Ok(Self(<#inner_type>::from_bytes(bytes)?))
86                }
87            }
88        }
89    } else {
90        quote! {
91            impl ::vecdb::Bytes for #struct_name {
92                type Array = <#inner_type as ::vecdb::Bytes>::Array;
93
94                fn to_bytes(&self) -> Self::Array {
95                    self.0.to_bytes()
96                }
97
98                fn from_bytes(bytes: &[u8]) -> ::vecdb::Result<Self> {
99                    Ok(Self(<#inner_type>::from_bytes(bytes)?))
100                }
101            }
102        }
103    };
104
105    TokenStream::from(expanded)
106}
107
108/// Derives the `Pco` trait for single-field tuple structs containing numeric types.
109///
110/// This macro enables custom wrapper types to work with `PcoVec` for compressed storage
111/// of numeric data using Pcodec compression.
112///
113/// # Requirements
114///
115/// - Must be a tuple struct with exactly one field
116/// - The inner type must implement `Pco` (numeric types: u16-u64, i16-i64, f32, f64)
117/// - Supports generic type parameters
118///
119/// # Generated Implementation
120///
121/// The derive generates three trait implementations:
122///
123/// 1. `Bytes` - For serialization (same as `#[derive(Bytes)]`)
124/// 2. `Pco` - Specifies the numeric type for compression
125/// 3. `TransparentPco` - Marker trait for transparent wrappers
126///
127/// ```rust,ignore
128/// impl Pco for Wrapper<T> where T: Pco + Bytes {
129///     type NumberType = <T as Pco>::NumberType;
130/// }
131///
132/// impl TransparentPco<<T as Pco>::NumberType> for Wrapper<T>
133/// where T: Pco + Bytes {}
134///
135/// impl Bytes for Wrapper<T> where T: Pco + Bytes {
136///     // ... same as Bytes derive
137/// }
138/// ```
139///
140/// The `NumberType` is automatically propagated from the inner type, ensuring the
141/// wrapper has the same compression characteristics.
142///
143/// # Example
144///
145/// ```rust,ignore
146/// use vecdb::{Pco, PcoVec};
147///
148/// #[derive(Pco)]
149/// struct Price(f64);
150///
151/// #[derive(Pco)]
152/// struct NumericWrapper<T>(T); // Generic types supported
153///
154/// // Nested generics work too
155/// #[derive(Pco)]
156/// struct Container<T>(NumericWrapper<T>);
157/// ```
158#[proc_macro_derive(Pco)]
159pub fn derive_pco(input: TokenStream) -> TokenStream {
160    let input = parse_macro_input!(input as DeriveInput);
161    let struct_name = &input.ident;
162    let generics = &input.generics;
163    let (impl_generics, ty_generics, where_clause) = generics.split_for_impl();
164
165    let inner_type = match &input.data {
166        Data::Struct(DataStruct {
167            fields: Fields::Unnamed(fields),
168            ..
169        }) if fields.unnamed.len() == 1 => &fields.unnamed[0].ty,
170        _ => {
171            return syn::Error::new_spanned(
172                &input.ident,
173                "Pco can only be derived for single-field tuple structs",
174            )
175            .to_compile_error()
176            .into();
177        }
178    };
179
180    // Check if we have generic parameters
181    let has_generics = !generics.params.is_empty();
182
183    let expanded = if has_generics {
184        // For generic types, we need both Pco and Bytes bounds because:
185        // - Pco trait requires the NumberType
186        // - We call to_bytes/from_bytes methods which require Bytes
187        let where_clause = if where_clause.is_some() {
188            quote! { #where_clause #inner_type: ::vecdb::Pco + ::vecdb::Bytes, }
189        } else {
190            quote! { where #inner_type: ::vecdb::Pco + ::vecdb::Bytes, }
191        };
192
193        quote! {
194            impl #impl_generics ::vecdb::Bytes for #struct_name #ty_generics #where_clause {
195                type Array = <#inner_type as ::vecdb::Bytes>::Array;
196
197                fn to_bytes(&self) -> Self::Array {
198                    self.0.to_bytes()
199                }
200
201                fn from_bytes(bytes: &[u8]) -> ::vecdb::Result<Self> {
202                    Ok(Self(<#inner_type>::from_bytes(bytes)?))
203                }
204            }
205
206            impl #impl_generics ::vecdb::TransparentPco<<#inner_type as ::vecdb::Pco>::NumberType> for #struct_name #ty_generics #where_clause {}
207
208            impl #impl_generics ::vecdb::Pco for #struct_name #ty_generics #where_clause {
209                type NumberType = <#inner_type as ::vecdb::Pco>::NumberType;
210            }
211        }
212    } else {
213        quote! {
214            impl ::vecdb::Bytes for #struct_name {
215                type Array = <#inner_type as ::vecdb::Bytes>::Array;
216
217                fn to_bytes(&self) -> Self::Array {
218                    self.0.to_bytes()
219                }
220
221                fn from_bytes(bytes: &[u8]) -> ::vecdb::Result<Self> {
222                    Ok(Self(<#inner_type>::from_bytes(bytes)?))
223                }
224            }
225
226            impl ::vecdb::TransparentPco<<#inner_type as ::vecdb::Pco>::NumberType> for #struct_name {}
227
228            impl ::vecdb::Pco for #struct_name {
229                type NumberType = <#inner_type as ::vecdb::Pco>::NumberType;
230            }
231        }
232    };
233
234    TokenStream::from(expanded)
235}