fold_impls/
lib.rs

1// The general idea is to extract all the child `impl` blocks,
2// verify that they have the same form ("shell"),
3// then coalesce all the substance ("meat"),
4// into a single `impl` block with that same form.
5
6use proc_macro_error::{abort_call_site, proc_macro_error};
7use std::{cell::RefCell, mem::take};
8use syn::{parse_macro_input, visit_mut::VisitMut};
9
10thread_local! {
11 static IMPLS: RefCell<Vec<syn::ItemImpl>> = RefCell::new(vec![]);
12}
13
14struct ImplExtractor;
15
16impl VisitMut for ImplExtractor {
17 fn visit_stmt_mut(&mut self, node: &mut syn::Stmt) {
18  if let syn::Stmt::Item(syn::Item::Impl(item_impl)) = node {
19   IMPLS.with(|impls| impls.borrow_mut().push(item_impl.to_owned()));
20   *node = syn::parse_quote!(();); // replace the impl with a very boring stmt
21  }
22  syn::visit_mut::visit_stmt_mut(self, node); // visit node's children
23 }
24}
25
26// @mark fold_impls
27#[proc_macro_attribute]
28#[proc_macro_error]
29pub fn fold_impls(
30 _args: proc_macro::TokenStream,
31 stmt: proc_macro::TokenStream,
32) -> proc_macro::TokenStream {
33 let mut outer = parse_macro_input!(stmt as syn::Stmt);
34
35 // Extract (find and remove) impls
36 ImplExtractor.visit_stmt_mut(&mut outer);
37
38 let mut f = if let syn::Stmt::Item(syn::Item::Fn(f)) = outer {
39  f
40 } else {
41  abort_call_site!("#[fold_impls] needs to be declared on a function.")
42 };
43
44 // Separate meat (impl bodies) from shells (impl declarations)
45 let (meat, mut shells) = IMPLS.with(|i| -> (Vec<_>, Vec<_>) {
46  i.borrow_mut().drain(..).map(|mut s| (take(&mut s.items), s)).unzip()
47 });
48
49 // Verify impl shells match
50 shells.dedup();
51 let mut shell = match shells.len() {
52  0 => abort_call_site!("No impl blocks found."),
53  1 => shells.drain(..).next().unwrap(),
54  _ => abort_call_site!("All impl blocks must have the same outer definition."),
55 };
56
57 // Put bodies into one shell
58 shell.items = meat.into_iter().flatten().collect();
59
60 f.block.stmts.push(syn::Stmt::Item(syn::Item::Impl(shell)));
61
62 quote::quote!(#f).into()
63}