Skip to main content

arc_wrapper/
lib.rs

1use quote::{format_ident, ToTokens};
2use syn::{parse::Parse, punctuated::Punctuated, Token};
3
4#[derive(Default)]
5struct Method {
6    vis: Option<syn::Visibility>,
7    name: Option<syn::Ident>,
8    metas: Vec<syn::Meta>,
9}
10
11impl Parse for Method {
12    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
13        let mut r = Method::default();
14        let metas = Punctuated::<syn::Meta, Token![,]>::parse_terminated(input)?;
15        for meta in metas {
16            match meta {
17                syn::Meta::NameValue(pair) => {
18                    let key = path_to_ident(&pair.path);
19                    match key.as_str() {
20                        "vis" => {
21                            let vis = parse_vis_from_expr(&pair.value)?;
22                            assert!(r.vis.is_none(), "duplicate key: {}", key);
23                            r.vis = Some(vis);
24                        }
25                        "rename" | "method" => {
26                            let value_tokens = trim_string_lit(&pair.value)?;
27                            let name = syn::parse2::<syn::Ident>(value_tokens)?;
28                            assert!(r.name.is_none(), "duplicate key: {}", key);
29                            r.name = Some(name);
30                        }
31                        _ => {
32                            r.metas.push(syn::Meta::from(pair));
33                        }
34                    }
35                }
36                meta => {
37                    r.metas.push(meta);
38                }
39            }
40        }
41        Ok(r)
42    }
43}
44
45#[derive(Default)]
46struct Rw {
47    read: Method,
48    write: Method,
49}
50
51impl Parse for Rw {
52    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
53        let mut rw = Rw::default();
54        let metas = Punctuated::<syn::Meta, Token![,]>::parse_terminated(input)?;
55        for meta in metas {
56            match meta {
57                syn::Meta::NameValue(pair) => {
58                    let key = path_to_ident(&pair.path);
59                    match key.as_str() {
60                        "write_vis" => {
61                            let vis = parse_vis_from_expr(&pair.value)?;
62                            assert!(rw.write.vis.is_none(), "duplicate key: {}", key);
63                            rw.write.vis = Some(vis);
64                        }
65                        "write" => {
66                            let value_tokens = trim_string_lit(&pair.value)?;
67                            let name = syn::parse2::<syn::Ident>(value_tokens)?;
68                            assert!(rw.write.name.is_none(), "duplicate key: {}", key);
69                            rw.write.name = Some(name);
70                        }
71                        "read_vis" => {
72                            let vis = parse_vis_from_expr(&pair.value)?;
73                            assert!(rw.read.vis.is_none(), "duplicate key: {}", key);
74                            rw.read.vis = Some(vis);
75                        }
76                        "read" => {
77                            let value_tokens = trim_string_lit(&pair.value)?;
78                            let name = syn::parse2::<syn::Ident>(value_tokens)?;
79                            assert!(rw.read.name.is_none(), "duplicate key: {}", key);
80                            rw.read.name = Some(name);
81                        }
82                        _ => panic!("unexpected key: {}", key),
83                    }
84                }
85                syn::Meta::List(list) => {
86                    let list_name = path_to_ident(&list.path);
87                    match list_name.as_str() {
88                        "read" => {
89                            rw.read = syn::parse2::<Method>(list.tokens)?;
90                        }
91                        "write" => {
92                            rw.write = syn::parse2::<Method>(list.tokens)?;
93                        }
94                        _ => panic!("unexpected key: {}", list_name),
95                    }
96                }
97                _ => panic!("unexpected value"),
98            }
99        }
100        Ok(rw)
101    }
102}
103
104type Mutex = Method;
105
106#[derive(Default)]
107enum Lock {
108    Mutex(Mutex),
109    Rw(Rw),
110    #[default]
111    Arc,
112}
113
114#[derive(Default)]
115struct Config {
116    struct_vis: Option<syn::Visibility>,
117    lock: Option<Lock>,
118    name: Option<syn::Ident>,
119    metas: Vec<syn::Meta>,
120}
121
122fn path_to_ident(path: &syn::Path) -> String {
123    assert!(path.leading_colon.is_none(), "expect ident");
124    assert_eq!(path.segments.len(), 1, "expect ident");
125    path.segments.first().unwrap().ident.to_string()
126}
127
128impl Parse for Config {
129    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
130        let mut config = Config::default();
131
132        let metas = Punctuated::<syn::Meta, Token![,]>::parse_terminated(input)?;
133        for meta in metas {
134            match meta {
135                syn::Meta::Path(path) => {
136                    let ident = path_to_ident(&path);
137
138                    match ident.as_str() {
139                        "mutex" => {
140                            assert!(config.lock.is_none(), "duplicate key: {}", ident);
141                            config.lock = Some(Lock::Mutex(Mutex::default()));
142                        }
143                        "rwlock" => {
144                            assert!(config.lock.is_none(), "duplicate key: {}", ident);
145                            config.lock = Some(Lock::Rw(Rw::default()));
146                        }
147                        _ => {
148                            config.metas.push(syn::Meta::from(path));
149                        }
150                    }
151                }
152                syn::Meta::NameValue(pair) => {
153                    let key = path_to_ident(&pair.path);
154
155                    match key.as_str() {
156                        "vis" => {
157                            let vis = parse_vis_from_expr(&pair.value)?;
158                            assert!(config.struct_vis.is_none(), "duplicate key: {}", key);
159                            config.struct_vis = Some(vis);
160                        }
161                        "lock" => {
162                            let value_tokens = trim_string_lit(&pair.value)?;
163                            let lock = match value_tokens.to_string().as_str() {
164                                "mutex" => Lock::Mutex(Mutex::default()),
165                                "rwlock" => Lock::Rw(Rw::default()),
166                                "none" => Lock::Arc,
167                                _ => panic!(
168                                    "unexpected value: {}; expect one of [mutex, rwloock, none]",
169                                    value_tokens
170                                ),
171                            };
172                            assert!(config.lock.is_none(), "duplicate key: {}", key);
173                            config.lock = Some(lock);
174                        }
175                        "rename" => {
176                            let value_tokens = trim_string_lit(&pair.value)?;
177                            let name = syn::parse2::<syn::Ident>(value_tokens)?;
178                            assert!(config.name.is_none(), "duplicate key: {}", key);
179                            config.name = Some(name);
180                        }
181                        _ => {
182                            config.metas.push(syn::Meta::from(pair));
183                        }
184                    }
185                }
186                syn::Meta::List(list) => {
187                    let list_name = path_to_ident(&list.path);
188                    match list_name.as_str() {
189                        "mutex" => {
190                            assert!(config.lock.is_none(), "duplicate key: {}", list_name);
191                            config.lock = Some(Lock::Mutex(syn::parse2::<Mutex>(list.tokens)?));
192                        }
193                        "rwlock" => {
194                            assert!(config.lock.is_none(), "duplicate key: {}", list_name);
195                            config.lock = Some(Lock::Rw(syn::parse2::<Rw>(list.tokens)?));
196                        }
197                        _ => {
198                            config.metas.push(syn::Meta::from(list));
199                        }
200                    }
201                }
202            }
203        }
204        Ok(config)
205    }
206}
207
208fn trim_string_lit(expr: &syn::Expr) -> Result<proc_macro2::TokenStream, proc_macro2::LexError> {
209    expr.into_token_stream()
210        .to_string()
211        .trim_start_matches("\"")
212        .trim_end_matches("\"")
213        .parse()
214}
215
216fn parse_vis_from_expr(expr: &syn::Expr) -> syn::Result<syn::Visibility> {
217    let string = expr.into_token_stream().to_string();
218    let string = string.trim_start_matches("\"").trim_end_matches("\"");
219    match string {
220        "hidden" => Ok(syn::Visibility::Inherited),
221        vis => syn::parse_str::<syn::Visibility>(vis),
222    }
223}
224
225#[proc_macro_attribute]
226pub fn arc_wrapper(
227    attr: proc_macro::TokenStream,
228    item: proc_macro::TokenStream,
229) -> proc_macro::TokenStream {
230    let derive_input = item.clone();
231    let raw = proc_macro2::TokenStream::from(item);
232    let derive_input = syn::parse_macro_input!(derive_input as syn::DeriveInput);
233    let Config {
234        struct_vis: new_struct_vis,
235        lock: lock_kind,
236        name: new_struct_name,
237        metas: new_struct_metas,
238    } = syn::parse_macro_input!(attr as Config);
239    let lock_kind = lock_kind.unwrap_or_default();
240
241    let raw_struct_name = derive_input.ident.clone();
242    let new_struct_vis = new_struct_vis.unwrap_or(derive_input.vis);
243    let new_struct_name = new_struct_name.unwrap_or(format_ident!("Arc{}", derive_input.ident));
244
245    let generices_without_bounds = 'gwob: {
246        let generics = derive_input.generics.clone();
247        if generics.lt_token.is_none() {
248            break 'gwob quote::quote! {};
249        }
250        let generics = generics.params.into_iter().map(|param| match param {
251            syn::GenericParam::Lifetime(lifetime) => lifetime.lifetime.into_token_stream(),
252            syn::GenericParam::Type(r#type) => r#type.ident.into_token_stream(),
253            syn::GenericParam::Const(r#const) => r#const.ident.into_token_stream(),
254        });
255        quote::quote! {
256            <#(#generics,)*>
257        }
258    };
259
260    let raw_struct_generic_without_where = {
261        let mut raw_struct_generic = derive_input.generics.clone();
262        raw_struct_generic.where_clause = None;
263        raw_struct_generic
264    };
265    let raw_struct_type = quote::quote! {
266        #raw_struct_name #generices_without_bounds
267    };
268    let where_clause = derive_input.generics.where_clause.clone();
269
270    let new_struct = {
271        let inner = match &lock_kind {
272            Lock::Mutex(_) => quote::quote! {
273                ::std::sync::Arc<::std::sync::Mutex<#raw_struct_type>>
274            },
275            Lock::Rw(_) => quote::quote! {
276                ::std::sync::Arc<::std::sync::RwLock<#raw_struct_type>>
277            },
278            Lock::Arc => quote::quote! {
279                ::std::sync::Arc<#raw_struct_type>
280            },
281        };
282
283        quote::quote! {
284            #(#[#new_struct_metas])*
285            #new_struct_vis struct #new_struct_name #raw_struct_generic_without_where #where_clause {
286                inner: #inner
287            }
288        }
289    };
290
291    let from_impl = quote::quote! {
292        impl #raw_struct_generic_without_where From< #raw_struct_type > for #new_struct_name #generices_without_bounds #where_clause {
293            fn from(inner: #raw_struct_type ) -> Self {
294                Self {
295                    inner: ::std::sync::Arc::new(inner.into())
296                }
297            }
298        }
299    };
300
301    let new_struct_methods = match lock_kind {
302        Lock::Mutex(mutex) => {
303            let method = mutex.name.unwrap_or_else(|| format_ident!("lock_guard"));
304            let vis = mutex.vis.unwrap_or(new_struct_vis);
305            let metas = &mutex.metas;
306            quote::quote! {
307                #(#[#metas])*
308                #vis fn #method(&self) -> ::std::sync::MutexGuard<'_, #raw_struct_type > {
309                    ::std::result::Result::unwrap(::std::sync::Mutex::lock(&self.inner))
310                }
311            }
312        }
313        Lock::Rw(rw_lock) => {
314            let read_method = rw_lock
315                .read
316                .name
317                .unwrap_or_else(|| format_ident!("read_guard"));
318            let read_vis = rw_lock.read.vis.unwrap_or(new_struct_vis.clone());
319            let read_metas = &rw_lock.read.metas;
320            let write_method = rw_lock
321                .write
322                .name
323                .unwrap_or_else(|| format_ident!("write_guard"));
324            let write_vis = rw_lock.write.vis.unwrap_or(new_struct_vis);
325            let write_metas = &rw_lock.write.metas;
326            quote::quote! {
327                #(#[#read_metas])*
328                #read_vis fn #read_method(&self) -> ::std::sync::RwLockReadGuard<'_, #raw_struct_type > {
329                    ::std::result::Result::unwrap(::std::sync::RwLock::read(&self.inner))
330                }
331
332                #(#[#write_metas])*
333                #write_vis fn #write_method(&self) -> ::std::sync::RwLockWriteGuard<'_, #raw_struct_type > {
334                    ::std::result::Result::unwrap(::std::sync::RwLock::write(&self.inner))
335                }
336            }
337        }
338        Lock::Arc => quote::quote! {},
339    };
340
341    let new_struct_impl = quote::quote! {
342        impl #raw_struct_generic_without_where #new_struct_name #generices_without_bounds #where_clause  {
343            #new_struct_methods
344        }
345    };
346
347    quote::quote! {
348        #raw
349
350        #new_struct
351
352        #from_impl
353
354        #new_struct_impl
355
356    }
357    .into()
358}