use heck::SnakeCase;
use proc_macro::TokenStream;
use quote::{format_ident, quote};
use syn::Item;
pub(crate) fn walk_visitor(args: TokenStream, raw_input: TokenStream) -> TokenStream {
let input = syn::parse_macro_input!(raw_input as Item);
let (name, generics) = match &input {
Item::Enum(item) => (&item.ident, &item.generics),
Item::Struct(item) => (&item.ident, &item.generics),
Item::Type(item) => (&item.ident, &item.generics),
_ => panic!("unsupported item to derive WalkVisitor for"),
};
let funnel_name = syn::parse_macro_input!(args as Option<syn::Ident>);
let pre_visit_fn = format_ident!(
"pre_visit_{}",
name.to_string().to_snake_case(),
span = name.span()
);
let post_visit_fn = format_ident!(
"post_visit_{}",
name.to_string().to_snake_case(),
span = name.span()
);
crate::visitor::add_call(&name, &generics);
let mut impl_generics = generics.clone();
let lt = crate::first_lifetime(&mut impl_generics);
let doc = format!(
r#"
Walk a visitor over `self`.
Calling this function is equivalent to calling:
- `visitor.{}(self)`
- `self.accept(visitor)`
- `visitor.{}(self);`
"#,
pre_visit_fn, post_visit_fn,
);
let mut pre_visit_calls = vec![];
let mut post_visit_calls = vec![];
pre_visit_calls.push(quote! {
let mut visit = visitor.#pre_visit_fn(self);
});
post_visit_calls.push(quote! {
visitor.#post_visit_fn(self);
});
if let Some(funnel_name) = funnel_name {
let pre_funnel_fn = format_ident!("pre_visit_{}", funnel_name);
let post_funnel_fn = format_ident!("post_visit_{}", funnel_name);
pre_visit_calls.push(quote! {
if visit { visit = visitor.#pre_funnel_fn(self); }
});
post_visit_calls.push(quote! {
visitor.#post_funnel_fn(self);
});
}
post_visit_calls.reverse();
let output = quote! {
#input
impl #impl_generics WalkVisitor<#lt> for #name #generics {
#[doc = #doc]
fn walk(&#lt self, visitor: &mut dyn Visitor<#lt>) {
#(#pre_visit_calls)*
if visit {
self.accept(visitor);
}
#(#post_visit_calls)*
}
}
};
output.into()
}