Skip to main content

anu_macros/
lib.rs

1#![recursion_limit = "256"]
2extern crate proc_macro;
3extern crate proc_macro2;
4#[macro_use]
5extern crate quote;
6
7use proc_macro::TokenStream;
8use proc_macro2::*;
9
10use std::iter::FromIterator;
11
12fn name_capital(name: &str) -> String {
13    name.chars()
14        .enumerate()
15        .map(|(i, s)| {
16            if i == 0 {
17                s.to_uppercase().nth(0).unwrap()
18            } else {
19                s
20            }
21        })
22        .collect()
23}
24
25#[proc_macro]
26pub fn table(input: proc_macro::TokenStream) -> TokenStream {
27    let input = proc_macro2::TokenStream::from(input);
28    let mut input_iter = input.into_iter();
29    let name = match input_iter.next() {
30        Some(TokenTree::Ident(id)) => id.to_string(),
31        _ => panic!("txn_table: first argument not an identifier"),
32    };
33    assert!(input_iter.next().is_none());
34    let name_capital = syn::Ident::new(&name_capital(&name), Span::call_site());
35    proc_macro::TokenStream::from(quote! {
36        #[doc(hidden)]
37        type #name_capital;
38    })
39}
40
41#[proc_macro]
42pub fn sanakirja_table_get(input: proc_macro::TokenStream) -> TokenStream {
43    let input = proc_macro2::TokenStream::from(input);
44    let mut input_iter = input.into_iter();
45    let name = match input_iter.next() {
46        Some(TokenTree::Ident(id)) => id.to_string(),
47        _ => panic!("txn_table: first argument not an identifier"),
48    };
49    let name_get = syn::Ident::new(&format!("get_{}", name), Span::call_site());
50    let name = syn::Ident::new(&name, Span::call_site());
51    let key = proc_macro2::TokenStream::from_iter(next(&mut input_iter).into_iter());
52    let value = proc_macro2::TokenStream::from_iter(next(&mut input_iter).into_iter());
53
54    let pre_ = proc_macro2::TokenStream::from_iter(next(&mut input_iter).into_iter());
55    let pre = if !pre_.is_empty() {
56        quote! {
57            let (key, value) = #pre_;
58        }
59    } else {
60        quote! {}
61    };
62    let post_ = proc_macro2::TokenStream::from_iter(next(&mut input_iter).into_iter());
63    let post = if post_.is_empty() {
64        quote! { self.txn.get(&self.#name, key, value) }
65    } else {
66        quote! { self.txn.get(&self.#name, key, value) . #post_ }
67    };
68    proc_macro::TokenStream::from(quote! {
69        #[doc(hidden)]
70        fn #name_get <'txn> (&'txn self, key: #key, value: Option<#value>) -> Option<#value> {
71            use ::sanakirja::Transaction;
72            #pre
73            #post
74        }
75    })
76}
77
78#[proc_macro]
79pub fn sanakirja_get(input: proc_macro::TokenStream) -> TokenStream {
80    let input = proc_macro2::TokenStream::from(input);
81    let mut input_iter = input.into_iter();
82    let name = match input_iter.next() {
83        Some(TokenTree::Ident(id)) => id.to_string(),
84        _ => panic!("txn_table: first argument not an identifier"),
85    };
86    let name_capital = syn::Ident::new(&name_capital(&name), Span::call_site());
87    let name_get = syn::Ident::new(&format!("get_{}", name), Span::call_site());
88    let key = proc_macro2::TokenStream::from_iter(next(&mut input_iter).into_iter());
89    let value = proc_macro2::TokenStream::from_iter(next(&mut input_iter).into_iter());
90    assert!(input_iter.next().is_none());
91    proc_macro::TokenStream::from(quote! {
92        #[doc(hidden)]
93        fn #name_get(&self, db: &Self::#name_capital, key: #key, value: Option<#value>) -> Option<#value> {
94            use ::sanakirja::Transaction;
95            self.txn.get(db, key, value)
96        }
97    })
98}
99
100#[proc_macro]
101pub fn table_get(input: proc_macro::TokenStream) -> TokenStream {
102    let input = proc_macro2::TokenStream::from(input);
103    let mut input_iter = input.into_iter();
104    let name = match input_iter.next() {
105        Some(TokenTree::Ident(id)) => id.to_string(),
106        _ => panic!("txn_table: first argument not an identifier"),
107    };
108    let name_get = syn::Ident::new(&format!("get_{}", name), Span::call_site());
109    let key = proc_macro2::TokenStream::from_iter(next(&mut input_iter).into_iter());
110    let value = proc_macro2::TokenStream::from_iter(next(&mut input_iter).into_iter());
111    assert!(input_iter.next().is_none());
112    proc_macro::TokenStream::from(quote! {
113        #[doc(hidden)]
114        fn #name_get<'txn>(&'txn self, key: #key, value: Option<#value>) -> Option<#value>;
115    })
116}
117
118#[proc_macro]
119pub fn get(input: proc_macro::TokenStream) -> TokenStream {
120    let input = proc_macro2::TokenStream::from(input);
121    let mut input_iter = input.into_iter();
122    let name = match input_iter.next() {
123        Some(TokenTree::Ident(id)) => id.to_string(),
124        _ => panic!("txn_table: first argument not an identifier"),
125    };
126    let name_capital = syn::Ident::new(&name_capital(&name), Span::call_site());
127    let name_get = syn::Ident::new(&format!("get_{}", name), Span::call_site());
128    let key = proc_macro2::TokenStream::from_iter(next(&mut input_iter).into_iter());
129    let value = proc_macro2::TokenStream::from_iter(next(&mut input_iter).into_iter());
130    assert!(input_iter.next().is_none());
131    proc_macro::TokenStream::from(quote! {
132        #[doc(hidden)]
133        fn #name_get<'txn>(&'txn self, db: &Self::#name_capital, key: #key, value: Option<#value>) -> Option<#value>;
134    })
135}
136
137fn next(input_iter: &mut proc_macro2::token_stream::IntoIter) -> Vec<TokenTree> {
138    let mut result = Vec::new();
139    let mut is_first = true;
140    loop {
141        match input_iter.next() {
142            Some(TokenTree::Punct(p)) => {
143                if p.as_char() == ',' {
144                    if !is_first {
145                        return result;
146                    }
147                } else {
148                    result.push(TokenTree::Punct(p))
149                }
150            }
151            Some(e) => result.push(e),
152            None => return result,
153        }
154        is_first = false
155    }
156}
157
158#[proc_macro]
159pub fn cursor(input: proc_macro::TokenStream) -> TokenStream {
160    cursor_(input, false, false, false)
161}
162
163#[proc_macro]
164pub fn cursor_ref(input: proc_macro::TokenStream) -> TokenStream {
165    cursor_(input, false, false, true)
166}
167
168#[proc_macro]
169pub fn iter(input: proc_macro::TokenStream) -> TokenStream {
170    cursor_(input, false, true, false)
171}
172
173#[proc_macro]
174pub fn rev_cursor(input: proc_macro::TokenStream) -> TokenStream {
175    cursor_(input, true, false, false)
176}
177
178fn cursor_(input: proc_macro::TokenStream, rev: bool, iter: bool, borrow: bool) -> TokenStream {
179    let input = proc_macro2::TokenStream::from(input);
180    let mut input_iter = input.into_iter();
181    let name = match input_iter.next() {
182        Some(TokenTree::Ident(id)) => id.to_string(),
183        _ => panic!("txn_table: first argument not an identifier"),
184    };
185    let capital = name_capital(&name);
186    let cursor_name = syn::Ident::new(&format!("{}Cursor", capital,), Span::call_site());
187    let name_capital = syn::Ident::new(&name_capital(&name), Span::call_site());
188    let name_iter = syn::Ident::new(&format!("iter_{}", name), Span::call_site());
189    let name_next = syn::Ident::new(&format!("cursor_{}_next", name), Span::call_site());
190    let name_prev = syn::Ident::new(&format!("cursor_{}_prev", name), Span::call_site());
191    let name_cursor = syn::Ident::new(
192        &format!("{}cursor_{}", if rev { "rev_" } else { "" }, name),
193        Span::call_site(),
194    );
195    let name_cursor_ref = syn::Ident::new(
196        &format!("{}cursor_{}_ref", if rev { "rev_" } else { "" }, name),
197        Span::call_site(),
198    );
199
200    let key = proc_macro2::TokenStream::from_iter(next(&mut input_iter).into_iter());
201    let value = proc_macro2::TokenStream::from_iter(next(&mut input_iter).into_iter());
202    let cursor_type = if rev {
203        quote! {
204            crate::pristine::RevCursor<Self, &'txn Self, Self::#cursor_name, #key, #value>
205        }
206    } else {
207        quote! {
208            crate::pristine::Cursor<Self, &'txn Self, Self::#cursor_name, #key, #value>
209        }
210    };
211    let def = if rev {
212        quote! {}
213    } else {
214        quote! {
215            #[doc(hidden)]
216            type #cursor_name;
217            #[doc(hidden)]
218            fn #name_next <'txn> (
219                &'txn self,
220                cursor: &mut Self::#cursor_name,
221            ) -> Option<(#key, #value)>;
222            #[doc(hidden)]
223            fn #name_prev <'txn> (
224                &'txn self,
225                cursor: &mut Self::#cursor_name,
226            ) -> Option<(#key, #value)>;
227        }
228    };
229    let borrow = if borrow {
230        quote! {
231        #[doc(hidden)]
232        fn #name_cursor_ref<RT: std::ops::Deref<Target = Self>>(
233            txn: RT,
234            db: &Self::#name_capital,
235            pos: Option<(#key, Option<#value>)>,
236        ) -> crate::pristine::Cursor<Self, RT, Self::#cursor_name, #key, #value>;
237        }
238    } else {
239        quote! {}
240    };
241    let iter = if !iter {
242        quote! {}
243    } else {
244        quote! {
245            #[doc(hidden)]
246            fn #name_iter <'txn> (
247                &'txn self,
248                k: #key,
249                v: Option<#value>
250            ) -> #cursor_type;
251        }
252    };
253    assert!(input_iter.next().is_none());
254    proc_macro::TokenStream::from(quote! {
255        #def
256        #[doc(hidden)]
257        fn #name_cursor<'txn>(
258            &'txn self,
259            db: &Self::#name_capital,
260            pos: Option<(#key, Option<#value>)>,
261        ) -> #cursor_type;
262        #borrow
263        #iter
264    })
265}
266
267#[proc_macro]
268pub fn sanakirja_cursor(input: proc_macro::TokenStream) -> TokenStream {
269    sanakirja_cursor_(input, false, false, false)
270}
271
272#[proc_macro]
273pub fn sanakirja_cursor_ref(input: proc_macro::TokenStream) -> TokenStream {
274    sanakirja_cursor_(input, false, false, true)
275}
276
277#[proc_macro]
278pub fn sanakirja_iter(input: proc_macro::TokenStream) -> TokenStream {
279    sanakirja_cursor_(input, false, true, false)
280}
281
282#[proc_macro]
283pub fn sanakirja_rev_cursor(input: proc_macro::TokenStream) -> TokenStream {
284    sanakirja_cursor_(input, true, false, false)
285}
286
287fn sanakirja_cursor_(
288    input: proc_macro::TokenStream,
289    rev: bool,
290    iter: bool,
291    borrow: bool,
292) -> TokenStream {
293    let input = proc_macro2::TokenStream::from(input);
294    let mut input_iter = input.into_iter();
295    let name = match input_iter.next() {
296        Some(TokenTree::Ident(id)) => id.to_string(),
297        _ => panic!("txn_table: first argument not an identifier"),
298    };
299    let cursor_name = syn::Ident::new(
300        &format!("{}Cursor", name_capital(&name),),
301        Span::call_site(),
302    );
303
304    let name_capital = syn::Ident::new(&name_capital(&name), Span::call_site());
305    let name_next = syn::Ident::new(&format!("cursor_{}_next", name), Span::call_site());
306    let name_prev = syn::Ident::new(&format!("cursor_{}_prev", name), Span::call_site());
307    let name_cursor = syn::Ident::new(
308        &format!("{}cursor_{}", if rev { "rev_" } else { "" }, name),
309        Span::call_site(),
310    );
311    let name_cursor_ref = syn::Ident::new(
312        &format!("{}cursor_{}_ref", if rev { "rev_" } else { "" }, name),
313        Span::call_site(),
314    );
315    let name_iter = syn::Ident::new(
316        &format!("{}iter_{}", if rev { "rev_" } else { "" }, name),
317        Span::call_site(),
318    );
319
320    let name = syn::Ident::new(&name, Span::call_site());
321    let key = proc_macro2::TokenStream::from_iter(next(&mut input_iter).into_iter());
322    let value = proc_macro2::TokenStream::from_iter(next(&mut input_iter).into_iter());
323
324    let pre = proc_macro2::TokenStream::from_iter(next(&mut input_iter).into_iter());
325    let post = proc_macro2::TokenStream::from_iter(next(&mut input_iter).into_iter());
326    let pre_init = if !pre.is_empty() {
327        quote! { let pos = #pre; }
328    } else {
329        quote! {}
330    };
331    let post = if !post.is_empty() {
332        quote! { . #post }
333    } else {
334        quote! {}
335    };
336    let iter = if iter {
337        quote! {
338            #[doc(hidden)]
339            fn #name_iter <'txn> (
340                &'txn self,
341                k: #key,
342                v: Option<#value>
343            ) -> super::Cursor<Self, &'txn Self, Self::#cursor_name, #key, #value> {
344                self.#name_cursor(&self.#name, Some((k, v)))
345            }
346        }
347    } else {
348        quote! {}
349    };
350
351    let borrow = if borrow {
352        quote! {
353            #[doc(hidden)]
354            fn #name_cursor_ref <RT: std::ops::Deref<Target = Self>> (
355                txn: RT,
356                db: &Self::#name_capital,
357                pos: Option<(#key, Option<#value>)>,
358            ) -> super::Cursor<Self, RT, Self::#cursor_name, #key, #value> {
359                #pre_init
360                let mut cursor = txn.txn.set_cursors(&db, pos).0;
361                super::Cursor {
362                    cursor,
363                    txn,
364                    marker: std::marker::PhantomData,
365                }
366            }
367        }
368    } else {
369        quote! {}
370    };
371
372    let result = proc_macro::TokenStream::from(if rev {
373        quote! {
374            #[doc(hidden)]
375            fn #name_cursor<'txn>(
376                &'txn self,
377                db: &Self::#name_capital,
378                pos: Option<(#key, Option<#value>)>,
379            ) -> super::RevCursor<Self, &'txn Self, Self::#cursor_name, #key, #value> {
380                #pre_init
381                let mut cursor = if pos.is_some() {
382                    self.txn.set_cursors(&db, pos).0
383                } else {
384                    self.txn.set_cursors_last(&db)
385                };
386                super::RevCursor {
387                    cursor,
388                    txn: self,
389                    marker: std::marker::PhantomData,
390                }
391            }
392        }
393    } else {
394        quote! {
395            #[doc(hidden)]
396            type #cursor_name = ::sanakirja::Cursor;
397            #[doc(hidden)]
398            fn #name_cursor<'txn>(
399                &'txn self,
400                db: &Self::#name_capital,
401                pos: Option<(#key, Option<#value>)>,
402            ) -> super::Cursor<Self, &'txn Self, Self::#cursor_name, #key, #value> {
403                #pre_init
404                let mut cursor = self.txn.set_cursors(&db, pos).0;
405                super::Cursor {
406                    cursor,
407                    txn: self,
408                    marker: std::marker::PhantomData,
409                }
410            }
411            #borrow
412            #[doc(hidden)]
413            fn #name_next <'txn> (
414                &'txn self,
415                cursor: &mut Self::#cursor_name,
416            ) -> Option<(#key, #value)> {
417                (unsafe { ::sanakirja::next(&self.txn, cursor) })
418                    #post
419            }
420            #[doc(hidden)]
421            fn #name_prev <'txn> (
422                &'txn self,
423                cursor: &mut Self::#cursor_name,
424            ) -> Option<(#key, #value)> {
425                (unsafe { ::sanakirja::prev(&self.txn, cursor) })
426                    #post
427            }
428            #iter
429        }
430    });
431    result
432}
433
434#[proc_macro]
435pub fn initialized_cursor(input: proc_macro::TokenStream) -> TokenStream {
436    initialized_cursor_(input, false)
437}
438
439#[proc_macro]
440pub fn initialized_rev_cursor(input: proc_macro::TokenStream) -> TokenStream {
441    initialized_cursor_(input, true)
442}
443
444fn initialized_cursor_(input: proc_macro::TokenStream, rev: bool) -> TokenStream {
445    let input = proc_macro2::TokenStream::from(input);
446    let mut input_iter = input.into_iter();
447    let name = match input_iter.next() {
448        Some(TokenTree::Ident(id)) => id.to_string(),
449        _ => panic!("txn_table: first argument not an identifier"),
450    };
451    let cursor_name = syn::Ident::new(
452        &format!("{}Cursor", name_capital(&name),),
453        Span::call_site(),
454    );
455    let name_next = syn::Ident::new(&format!("cursor_{}_next", name), Span::call_site());
456    let name_prev = syn::Ident::new(&format!("cursor_{}_prev", name), Span::call_site());
457    let key = proc_macro2::TokenStream::from_iter(next(&mut input_iter).into_iter());
458    let value = proc_macro2::TokenStream::from_iter(next(&mut input_iter).into_iter());
459
460    assert!(input_iter.next().is_none());
461    if rev {
462        proc_macro::TokenStream::from(quote! {
463            impl<T: TxnT, RT: std::ops::Deref<Target = T>> Iterator for crate::pristine::RevCursor<T, RT, T::#cursor_name, #key, #value>
464            {
465                type Item = (#key, #value);
466                fn next(&mut self) -> Option<(#key, #value)> {
467                    self.txn.#name_prev(&mut self.cursor)
468                }
469            }
470        })
471    } else {
472        proc_macro::TokenStream::from(quote! {
473
474            impl<T: TxnT, RT: std::ops::Deref<Target = T>>
475                crate::pristine::Cursor<T, RT, T::#cursor_name, #key, #value>
476            {
477                pub fn prev(&mut self) -> Option<(#key, #value)> {
478                    self.txn.#name_prev(&mut self.cursor)
479                }
480            }
481            impl<T: TxnT, RT: std::ops::Deref<Target = T>> Iterator for crate::pristine::Cursor<T, RT, T::#cursor_name, #key, #value>
482            {
483                type Item = (#key, #value);
484                fn next(&mut self) -> Option<(#key, #value)> {
485                    self.txn.#name_next(&mut self.cursor)
486                }
487            }
488        })
489    }
490}
491
492#[proc_macro]
493pub fn put_del(input: proc_macro::TokenStream) -> TokenStream {
494    let input = proc_macro2::TokenStream::from(input);
495    let mut input_iter = input.into_iter();
496    let name = match input_iter.next() {
497        Some(TokenTree::Ident(id)) => id.to_string(),
498        _ => panic!("txn_table: first argument not an identifier"),
499    };
500    let put = syn::Ident::new(&format!("put_{}", name), Span::call_site());
501    let del = syn::Ident::new(&format!("del_{}", name), Span::call_site());
502
503    let key = proc_macro2::TokenStream::from_iter(next(&mut input_iter).into_iter());
504    let value = proc_macro2::TokenStream::from_iter(next(&mut input_iter).into_iter());
505    assert!(input_iter.next().is_none());
506    proc_macro::TokenStream::from(quote! {
507        #[doc(hidden)]
508        fn #put(
509            &mut self,
510            k: #key,
511            e: #value,
512        ) -> Result<bool, anyhow::Error>;
513        #[doc(hidden)]
514        fn #del(
515            &mut self,
516            k: #key,
517            e: Option<#value>,
518        ) -> Result<bool, anyhow::Error>;
519    })
520}
521
522#[proc_macro]
523pub fn sanakirja_put_del(input: proc_macro::TokenStream) -> TokenStream {
524    let input = proc_macro2::TokenStream::from(input);
525    let mut input_iter = input.into_iter();
526    let name = match input_iter.next() {
527        Some(TokenTree::Ident(id)) => id.to_string(),
528        _ => panic!("txn_table: first argument not an identifier"),
529    };
530    let put = syn::Ident::new(&format!("put_{}", name), Span::call_site());
531    let del = syn::Ident::new(&format!("del_{}", name), Span::call_site());
532    let name = syn::Ident::new(&name, Span::call_site());
533
534    let key = proc_macro2::TokenStream::from_iter(next(&mut input_iter).into_iter());
535    let value = proc_macro2::TokenStream::from_iter(next(&mut input_iter).into_iter());
536
537    let pre_key = proc_macro2::TokenStream::from_iter(next(&mut input_iter).into_iter());
538    let pre_value = proc_macro2::TokenStream::from_iter(next(&mut input_iter).into_iter());
539    assert!(input_iter.next().is_none());
540    if pre_key.is_empty() {
541        proc_macro::TokenStream::from(quote! {
542            #[doc(hidden)]
543            fn #put(
544                &mut self,
545                k: #key,
546                v: #value,
547            ) -> Result<bool, anyhow::Error> {
548                Ok(self.txn.put(&mut self.rng, &mut self.#name, k, v)?)
549            }
550            #[doc(hidden)]
551            fn #del(
552                &mut self,
553                k: #key,
554                v: Option<#value>,
555            ) -> Result<bool, anyhow::Error> {
556                Ok(self.txn.del(&mut self.rng, &mut self.#name, k, v)?)
557            }
558        })
559    } else {
560        proc_macro::TokenStream::from(quote! {
561            #[doc(hidden)]
562            fn #put(
563                &mut self,
564                k: #key,
565                v: #value,
566            ) -> Result<bool, anyhow::Error> {
567                let k = #pre_key;
568                let v = #pre_value;
569                Ok(self.txn.put(&mut self.rng, &mut self.#name, k, v)?)
570            }
571            #[doc(hidden)]
572            fn #del(
573                &mut self,
574                k: #key,
575                v: Option<#value>,
576            ) -> Result<bool, anyhow::Error> {
577                let k = #pre_key;
578                let v = v.map(|v| #pre_value);
579                Ok(self.txn.del(&mut self.rng, &mut self.#name, k, v)?)
580            }
581        })
582    }
583}