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 let len = format!("${}", variable_numbers.len() + 1);
138 variable_numbers.insert(var.to_string(), len);
139
140 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}