tantivy_macro/
lib.rs

1use darling::{ast, FromDeriveInput, FromField};
2use proc_macro2::TokenStream;
3use quote::{quote, ToTokens};
4use syn::{parse_macro_input, parse_quote, DeriveInput, Type};
5
6/// #[derive(Schema)]
7/// struct Struct {
8///     #[field(name="field_name", stored, indexed, coerce, norm)]
9///     field: String
10/// }
11/// detail in test mod.
12#[proc_macro_derive(Schema, attributes(field))]
13pub fn derive_tantivy_schema(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
14    let input = parse_macro_input!(input as DeriveInput);
15    let receiver = InputReceiver::from_derive_input(&input).unwrap();
16    quote!(#receiver).into()
17}
18
19impl ToTokens for InputReceiver {
20    fn to_tokens(&self, tokens: &mut TokenStream) {
21        let (impl_generics, ty_generics, where_clause) = self.generics.split_for_impl();
22        let name = &self.ident;
23        let fields = self
24            .data
25            .as_ref()
26            .take_struct()
27            .expect("Should never be enum")
28            .fields;
29
30        let mut field_entrys = Vec::new();
31        let mut field_values = Vec::new();
32        for (id, field) in fields.iter().enumerate() {
33            let (entry, value) = field.get_field_type_and_value(id as u32);
34            field_entrys.push(entry);
35            field_values.push(value);
36        }
37
38        tokens.extend(quote! {
39            impl #impl_generics #name #ty_generics #where_clause {
40                pub fn schema() -> tantivy::schema::Schema {
41                    let mut builder = tantivy::schema::Schema::builder();
42                    #(
43                        #field_entrys
44                        builder.add_field(entry);
45                    )*
46                    builder.build()
47                }
48            }
49
50            impl #impl_generics std::convert::Into<tantivy::schema::Document> for #name #ty_generics #where_clause {
51                fn into(self) -> tantivy::schema::Document {
52                    let mut document = tantivy::schema::Document::new();
53                    #(
54                        #field_values
55                        document.add_field_value(field, value);
56                    )*
57                    document
58                }
59            }
60        });
61    }
62}
63
64#[derive(Debug, FromDeriveInput)]
65#[darling(supports(struct_named))]
66struct InputReceiver {
67    ident: syn::Ident,
68    generics: syn::Generics,
69    data: ast::Data<(), FieldReceiver>,
70}
71
72/// #[field(name="name",fast,stored,indexed,coerce,norm,token)]
73#[derive(Debug, FromField)]
74#[darling(attributes(field))]
75struct FieldReceiver {
76    ident: Option<syn::Ident>,
77    ty: syn::Type,
78    name: Option<String>,
79    #[darling(default)]
80    fast: bool,
81    #[darling(default)]
82    stored: bool,
83    #[darling(default)]
84    indexed: bool,
85    #[darling(default)]
86    coerce: bool,
87    #[darling(default)]
88    norm: bool,
89    #[darling(default)]
90    tokenized: bool,
91}
92
93impl FieldReceiver {
94    /// -> (entry: tantivy::schema::FieldEntry, field: tantivy::schema::Field + value: tantivy::schema::Value)
95    fn get_field_type_and_value(&self, id: u32) -> (TokenStream, TokenStream) {
96        let FieldReceiver {
97            ref ident,
98            ref ty,
99            ref name,
100            fast,
101            stored,
102            indexed,
103            coerce,
104            norm,
105            tokenized,
106        } = *self;
107
108        let set_stored = if stored {
109            quote! { let options = options.set_stored(); }
110        } else {
111            TokenStream::new()
112        };
113
114        let set_fast = if fast {
115            quote! { let options = options.set_fast(); }
116        } else {
117            TokenStream::new()
118        };
119
120        let set_coerce = if coerce {
121            quote! { let options = options.set_coerce(); }
122        } else {
123            TokenStream::new()
124        };
125
126        let set_indexed = if indexed {
127            quote! { let options = options.set_indexed(); }
128        } else {
129            TokenStream::new()
130        };
131
132        let set_fieldnorm = if norm {
133            quote! { let options = options.set_fieldnorm(); }
134        } else {
135            TokenStream::new()
136        };
137
138        let create_field = quote! {
139            let field = tantivy::schema::Field::from_field_id(#id);
140        };
141
142        let field_name = ident
143            .as_ref()
144            .expect("only supported named struct")
145            .to_string();
146        let field_name = name.as_ref().unwrap_or(&field_name);
147
148        let str_ty: Type = parse_quote!(String);
149        let json_ty: Type = parse_quote!(Map<String, Value>);
150        if *ty == str_ty || *ty == json_ty {
151            // indexed default
152            let set_tokenizer = if tokenized {
153                quote! { let index_options = index_options.set_tokenizer("default").set_index_option(tantivy::schema::IndexRecordOption::WithFreqsAndPositions); }
154            } else {
155                quote! { let index_options = index_options.set_tokenizer("raw").set_index_option(tantivy::schema::IndexRecordOption::Basic); }
156            };
157            let set_fast = if fast {
158                quote! { let options = options.set_fast(None); }
159            } else {
160                TokenStream::new()
161            };
162            return if *ty == str_ty {
163                (
164                    quote! {
165                        let index_options = tantivy::schema::TextFieldIndexing::default().set_fieldnorms(#norm);
166                        #set_tokenizer
167                        let options = tantivy::schema::TextOptions::default().set_indexing_options(index_options);
168                        #set_stored
169                        #set_fast
170                        #set_coerce
171                        let entry = tantivy::schema::FieldEntry::new(
172                            String::from(#field_name),
173                            tantivy::schema::FieldType::Str(options)
174                        );
175                    },
176                    quote! {
177                        #create_field
178                        let value = tantivy::schema::Value::Str(self.#ident);
179                    },
180                )
181            } else {
182                // json
183                (
184                    quote! {
185                        let index_options = tantivy::schema::TextFieldIndexing::default().set_fieldnorms(#norm);
186                        #set_tokenizer
187                        let options = tantivy::schema::JsonObjectOptions::default().set_indexing_options(index_options);
188                        #set_stored
189                        #set_fast
190                        let entry = tantivy::schema::FieldEntry::new(
191                            String::from(#field_name),
192                            tantivy::schema::FieldType::JsonObject(options)
193                        );
194                    },
195                    quote! {
196                        #create_field
197                        let value = tantivy::schema::Value::JsonObject(self.#ident);
198                    },
199                )
200            };
201        }
202        let u64_ty: Type = parse_quote!(u64);
203        let i64_ty: Type = parse_quote!(i64);
204        let f64_ty: Type = parse_quote!(f64);
205        let bool_ty: Type = parse_quote!(bool);
206        if *ty == u64_ty || *ty == i64_ty || *ty == f64_ty || *ty == bool_ty {
207            let options = quote! {
208                let options = tantivy::schema::NumericOptions::default();
209                #set_stored
210                #set_fast
211                #set_coerce
212                #set_indexed
213                #set_fieldnorm
214            };
215            return if *ty == u64_ty {
216                (
217                    quote! {
218                        #options
219                        let entry = tantivy::schema::FieldEntry::new(
220                            String::from(#field_name),
221                            tantivy::schema::FieldType::U64(options)
222                        );
223                    },
224                    quote! {
225                        #create_field
226                        let value = tantivy::schema::Value::U64(self.#ident);
227                    },
228                )
229            } else if *ty == i64_ty {
230                (
231                    quote! {
232                        #options
233                        let ty = tantivy::schema::FieldType::I64(options);
234                        let entry = tantivy::schema::FieldEntry::new(
235                            String::from(#field_name),
236                            options
237                        );
238                    },
239                    quote! {
240                        #create_field
241                        let value = tantivy::schema::Value::I64(self.#ident);
242                    },
243                )
244            } else if *ty == f64_ty {
245                (
246                    quote! {
247                        #options
248                        let entry = tantivy::schema::FieldEntry::new(
249                            String::from(#field_name),
250                            tantivy::schema::FieldType::F64(options)
251                        );
252                    },
253                    quote! {
254                        #create_field
255                        let value = tantivy::schema::Value::F64(self.#ident);
256                    },
257                )
258            } else {
259                // bool
260                (
261                    quote! {
262                        #options
263                        let entry = tantivy::schema::FieldEntry::new(
264                            String::from(#field_name),
265                            tantivy::schema::FieldType::Bool(options)
266                        );
267                    },
268                    quote! {
269                        #create_field
270                        let value = tantivy::schema::Value::Bool(self.#ident);
271                    },
272                )
273            };
274        }
275        let bytes_ty: Type = parse_quote!(Vec<u8>);
276        if *ty == bytes_ty {
277            return (
278                quote! {
279                    let options = tantivy::schema::BytesOptions::default();
280                    #set_stored
281                    #set_fast
282                    #set_indexed
283                    #set_fieldnorm
284                    let entry = tantivy::schema::FieldEntry::new(
285                        String::from(#field_name),
286                        tantivy::schema::FieldType::Bytes(options)
287                    );
288                },
289                quote! {
290                    #create_field
291                    let value = tantivy::schema::Value::Bytes(self.#ident);
292                },
293            );
294        }
295        let facet_ty: Type = parse_quote!(Facet);
296        if *ty == facet_ty {
297            return (
298                quote! {
299                    let options = tantivy::schema::FacetOptions::default();
300                    #set_stored
301                    let entry = tantivy::schema::FieldEntry::new(
302                        String::from(#field_name),
303                        tantivy::schema::FieldType::Facet(options)
304                    );
305                },
306                quote! {
307                    #create_field
308                    let value = tantivy::schema::Value::Facet(#ident);
309                },
310            );
311        }
312        let date_ty: Type = parse_quote!(DateTime);
313        if *ty == date_ty {
314            return (
315                quote! {
316                    let options = tantivy::schema::DateOptions::default();
317                    #set_stored
318                    #set_fast
319                    #set_indexed
320                    #set_fieldnorm
321                    let entry = tantivy::schema::FieldEntry::new(
322                        String::from(#field_name),
323                        tantivy::schema::FieldType::Date(options)
324                    );
325                },
326                quote! {
327                    #create_field
328                    let value = tantivy::schema::Value::Date(self.#ident);
329                },
330            );
331        }
332        let ip_ty: Type = parse_quote!(Ipv6Addr);
333        if *ty == ip_ty {
334            return (
335                quote! {
336                    let options = tantivy::schema::IpAddrOptions::default();
337                    #set_stored
338                    #set_fast
339                    #set_indexed
340                    #set_fieldnorm
341                    let entry = tantivy::schema::FieldEntry::new(
342                        String::from(#field_name),
343                        tantivy::schema::FieldType::IpAddr(options)
344                    );
345                },
346                quote! {
347                    #create_field
348                    let value = tantivy::schema::Value::IpAddr(self.#ident);
349                },
350            );
351        }
352
353        match ty {
354            _ => panic!("unsupported field type"),
355        }
356    }
357}
358
359mod test {
360    #[test]
361    fn should_work() {
362        let input = r#"#[derive(Schema)]
363pub struct Doc {
364    #[field(name = "str", stored, tokenized)]
365    text: String,
366    #[field(fast)]
367    id: String,
368    #[field(fast, norm, coerce)]
369    num: u64,
370    #[field(stored, fast)]
371    date: DateTime,
372    #[field(stored, indexed)]
373    facet: Facet,
374    #[field(stored, indexed)]
375    bytes: Vec<u8>,
376    #[field(stored, indexed)]
377    json: Map<String, Value>,
378    #[field(fast)]
379    ip: Ipv6Addr
380}"#;
381
382        let parsed = syn::parse_str(input).unwrap();
383        use darling::FromDeriveInput;
384        let receiver = crate::InputReceiver::from_derive_input(&parsed).unwrap();
385        let tokens = quote::quote!(#receiver);
386
387        println!(
388            r#"
389INPUT:
390
391{}
392
393PARSED AS:
394
395{:?}
396
397EMITS:
398
399{}
400    "#,
401            input, receiver, tokens
402        );
403
404        let result = r#"impl Doc { pub fn schema () -> tantivy :: schema :: Schema { let mut builder = tantivy :: schema :: Schema :: builder () ; let index_options = tantivy :: schema :: TextFieldIndexing :: default () . set_fieldnorms (false) ; let index_options = index_options . set_tokenizer ("default") . set_index_option (tantivy :: schema :: IndexRecordOption :: WithFreqsAndPositions) ; let options = tantivy :: schema :: TextOptions :: default () . set_indexing_options (index_options) ; let options = options . set_stored () ; let entry = tantivy :: schema :: FieldEntry :: new (String :: from ("str") , tantivy :: schema :: FieldType :: Str (options)) ; builder . add_field (entry) ; let index_options = tantivy :: schema :: TextFieldIndexing :: default () . set_fieldnorms (false) ; let index_options = index_options . set_tokenizer ("raw") . set_index_option (tantivy :: schema :: IndexRecordOption :: Basic) ; let options = tantivy :: schema :: TextOptions :: default () . set_indexing_options (index_options) ; let options = options . set_fast (None) ; let entry = tantivy :: schema :: FieldEntry :: new (String :: from ("id") , tantivy :: schema :: FieldType :: Str (options)) ; builder . add_field (entry) ; let options = tantivy :: schema :: NumericOptions :: default () ; let options = options . set_fast () ; let options = options . set_coerce () ; let options = options . set_fieldnorm () ; let entry = tantivy :: schema :: FieldEntry :: new (String :: from ("num") , tantivy :: schema :: FieldType :: U64 (options)) ; builder . add_field (entry) ; let options = tantivy :: schema :: DateOptions :: default () ; let options = options . set_stored () ; let options = options . set_fast () ; let entry = tantivy :: schema :: FieldEntry :: new (String :: from ("date") , tantivy :: schema :: FieldType :: Date (options)) ; builder . add_field (entry) ; let options = tantivy :: schema :: FacetOptions :: default () ; let options = options . set_stored () ; let entry = tantivy :: schema :: FieldEntry :: new (String :: from ("facet") , tantivy :: schema :: FieldType :: Facet (options)) ; builder . add_field (entry) ; let options = tantivy :: schema :: BytesOptions :: default () ; let options = options . set_stored () ; let options = options . set_indexed () ; let entry = tantivy :: schema :: FieldEntry :: new (String :: from ("bytes") , tantivy :: schema :: FieldType :: Bytes (options)) ; builder . add_field (entry) ; let index_options = tantivy :: schema :: TextFieldIndexing :: default () . set_fieldnorms (false) ; let index_options = index_options . set_tokenizer ("default") . set_index_option (tantivy :: schema :: IndexRecordOption :: WithFreqsAndPositions) ; let options = tantivy :: schema :: JsonObjectOptions :: default () . set_indexing_options (index_options) ; let options = options . set_stored () ; let entry = tantivy :: schema :: FieldEntry :: new (String :: from ("json") , tantivy :: schema :: FieldType :: JsonObject (options)) ; builder . add_field (entry) ; let options = tantivy :: schema :: IpAddrOptions :: default () ; let options = options . set_fast () ; let entry = tantivy :: schema :: FieldEntry :: new (String :: from ("ip") , tantivy :: schema :: FieldType :: IpAddr (options)) ; builder . add_field (entry) ; builder . build () } } impl std :: convert :: Into < tantivy :: schema :: Document > for Doc { fn into (self) -> tantivy :: schema :: Document { let mut document = tantivy :: schema :: Document :: new () ; let field = tantivy :: schema :: Field :: from_field_id (0u32) ; let value = tantivy :: schema :: Value :: Str (self . text) ; document . add_field_value (field , value) ; let field = tantivy :: schema :: Field :: from_field_id (1u32) ; let value = tantivy :: schema :: Value :: Str (self . id) ; document . add_field_value (field , value) ; let field = tantivy :: schema :: Field :: from_field_id (2u32) ; let value = tantivy :: schema :: Value :: U64 (self . num) ; document . add_field_value (field , value) ; let field = tantivy :: schema :: Field :: from_field_id (3u32) ; let value = tantivy :: schema :: Value :: Date (self . date) ; document . add_field_value (field , value) ; let field = tantivy :: schema :: Field :: from_field_id (4u32) ; let value = tantivy :: schema :: Value :: Facet (facet) ; document . add_field_value (field , value) ; let field = tantivy :: schema :: Field :: from_field_id (5u32) ; let value = tantivy :: schema :: Value :: Bytes (self . bytes) ; document . add_field_value (field , value) ; let field = tantivy :: schema :: Field :: from_field_id (6u32) ; let value = tantivy :: schema :: Value :: JsonObject (self . json) ; document . add_field_value (field , value) ; let field = tantivy :: schema :: Field :: from_field_id (7u32) ; let value = tantivy :: schema :: Value :: IpAddr (self . ip) ; document . add_field_value (field , value) ; document } }"#;
405        assert_eq!(result, tokens.to_string())
406    }
407}