pleingres_macros/
lib.rs

1#![recursion_limit = "128"]
2extern crate proc_macro;
3extern crate proc_macro2;
4#[macro_use]
5extern crate syn;
6#[macro_use]
7extern crate quote;
8extern crate regex;
9
10use proc_macro2::*;
11use std::collections::HashMap;
12use regex::Captures;
13use syn::DeriveInput;
14use proc_macro::TokenStream;
15
16#[proc_macro_derive(Request)]
17pub fn request(input: TokenStream) -> TokenStream {
18    let input = parse_macro_input!(input as DeriveInput);
19    let ref name = input.ident;
20    let enum_ = match input.data {
21        syn::Data::Enum(ref enum_) => enum_,
22        _ => panic!("only enums are supported")
23    };
24    let request_clauses = enum_.variants.iter().map(|var| {
25        let var_name = &var.ident;
26        quote! {
27            #name::#var_name(ref mut a) => a.request(buf)
28        }
29    });
30    let from_instances = enum_.variants.iter().map(|var| {
31        let ty = &var.fields.iter().next().unwrap().ty;
32        let var_name = &var.ident;
33        quote! {
34            impl std::convert::From<#ty> for #name {
35                fn from(s: #ty) -> #name {
36                    #name::#var_name(s)
37                }
38            }
39            impl std::convert::From<#name> for #ty {
40                fn from(s: #name) -> #ty {
41                    match s {
42                        #name::#var_name(s) => s,
43                        _ => panic!(concat!("Wrong conversion into ", stringify!(#ty), " for type ", stringify!(#name)))
44                    }
45                }
46            }
47        }
48    });
49    TokenStream::from(quote!{
50        impl pleingres::Request for #name {
51            fn request(&mut self, buf: &mut pleingres::Buffer) {
52                match *self {
53                    #(#request_clauses),*
54                }
55            }
56        }
57        #(#from_instances)*
58    })
59}
60
61#[proc_macro_derive(HandleRowJoin)]
62pub fn handle_row_join(input: TokenStream) -> TokenStream {
63    let input = parse_macro_input!(input as DeriveInput);
64    let ref name = input.ident;
65    let enum_ = match input.data {
66        syn::Data::Enum(ref enum_) => enum_,
67        _ => panic!("only enums are supported")
68    };
69    let row_clauses = enum_.variants.iter().map(|var| {
70        let var_name = &var.ident;
71        quote! {
72            #name::#var_name(ref mut a) => a.row(row)
73        }
74    });
75    let complete_clauses = enum_.variants.iter().map(|var| {
76        let var_name = &var.ident;
77        quote! {
78            #name::#var_name(ref mut a) => a.complete(n)
79        }
80    });
81    let ready_clauses = enum_.variants.iter().map(|var| {
82        let var_name = &var.ident;
83        quote! {
84            #name::#var_name(ref mut a) => a.ready_for_query()
85        }
86    });
87    let err_clauses = enum_.variants.iter().map(|var| {
88        let var_name = &var.ident;
89        quote! {
90            #name::#var_name(ref mut a) => a.err(err)
91        }
92    });
93    TokenStream::from(quote!{
94        impl pleingres::HandleRow for #name {
95            fn row(&mut self, row: pleingres::Row) -> bool {
96                match *self {
97                    #(#row_clauses),*
98                }
99            }
100            fn complete(&mut self, n: u32) {
101                match *self {
102                    #(#complete_clauses),*
103                }
104            }
105            fn ready_for_query(&mut self) {
106                match *self {
107                    #(#ready_clauses),*
108                }
109            }
110            fn err(&mut self, err: &str) {
111                match *self {
112                    #(#err_clauses),*
113                }
114            }
115        }
116    })
117}
118
119#[proc_macro_attribute]
120pub fn sql(attr: proc_macro::TokenStream, item: proc_macro::TokenStream) -> proc_macro::TokenStream {
121    let attr = proc_macro2::TokenStream::from(attr);
122    let item = proc_macro2::TokenStream::from(item);
123    let re = regex::Regex::new(r"\$([a-zA-Z0-9_\.]+)").unwrap();
124    let binds = attr.into_iter().filter_map(|s| {
125        match s {
126            TokenTree::Literal(_) => {},
127            _ => return None,
128        }
129        let s = format!("{}", s);
130        let s = s.trim_matches('"');
131        let mut variable_numbers = HashMap::new();
132        let mut binds = Vec::new();
133        let replaced = re.replace_all(&s, |cap: &Captures| {
134            let var = cap.get(1).unwrap().as_str();
135            if variable_numbers.get(var).is_none() {
136                // Mapping the variable name to a number.
137                let len = format!("${}", variable_numbers.len() + 1);
138                variable_numbers.insert(var.to_string(), len);
139
140                // Making an expression
141                let fields = var.split('.').map(|field| {
142                    syn::Ident::new(&field, Span::call_site())
143                });
144                binds.push(quote!(
145                    &self.#(#fields).*
146                ));
147            }
148            variable_numbers.get(var).unwrap().to_string()
149        });
150        Some(quote! {
151            buf.bind(#replaced, &[#(#binds),*]).execute(0);
152        })
153    });
154    let mut name = None;
155    let mut item = item.into_iter();
156    let mut item_ = Vec::new();
157    loop {
158        match item.next() {
159            Some(TokenTree::Ident(id)) => {
160                if id.to_string() == "struct" {
161                    let it = item.next().unwrap();
162                    name = Some(syn::Ident::new(&format!("{}", it), it.span()));
163                    item_.push(TokenTree::Ident(id));
164                    item_.push(it);
165                } else {
166                    item_.push(TokenTree::Ident(id));
167                }
168            }
169            None => break,
170            Some(it) => {
171                item_.push(it)
172            }
173        }
174    }
175    let name = name.unwrap();
176    use std::iter::FromIterator;
177    let item = proc_macro2::TokenStream::from_iter(item_);
178    let tokens = quote! {
179        impl pleingres::Request for #name {
180            fn request(&mut self, buf: &mut pleingres::Buffer) {
181                #(#binds)*
182                debug!("request sent");
183            }
184        }
185        #item
186    };
187    proc_macro::TokenStream::from(tokens)
188}