Writing proc macros with syn, is there a way to visit parts of the AST that match a given format?

Let's say I want to build a #[proc_macro_attribute] that takes the following code:

#[say_hello_if_blue]
fn foobar() {
  do_something(1, 2, 3);
  do_something_blue(1, 2, 3);
  if some_condition() {
    if other_condition() {
      let a = get_value();
      let b = get_value_blue(a);
    }
  }
}

and transforms it into:

fn foobar() {
  do_something(1, 2, 3);
  do_something_blue(1, 2, 3, "Hello");
  if some_condition() {
    if other_condition() {
      let a = get_value();
      let b = get_value_blue(a, "Hello");
    }
  }
}

In other words, I'd like to be able to take an arbitrary syntax tree (in that case a function body), select parts of the tree that match simple conditions (eg "is a function calls whose name ends with _blue"), and perform some transformations on them while ignoring the rest.

Is this implemented in syn somewhere? Is there a crate that implements that feature? Or is it something I need to code from the ground up?

Aaaand... I just noticed syn::visit.

It's not exactly what I had in mind, but it looks like the direction I need to go in.

1 Like

Yes, there is the visit / visit_mut / fold APIs, I'd suggest you use the visit_mut for your use case:

struct AppendHelloToBlues;
impl VisitMut for AppendHelloToBlues {
    fn visit_expr_call_mut (
        self: &'_ mut Self,
        call: &'_ mut ExprCall,
    )
    {
        // 1 - subrecurse
        visit_expr_call_mut(self, call);
        // 2 - special case functions whose name ends in `_blue`
        if matches!(
            *call.func,
            Expr::Path(ExprPath { ref path, .. })
            if path.segments.last().unwrap().ident.to_string().ends_with("_blue")
        )
        {
            call.args.push(parse_quote!( "hello" ));
        }
    }

    fn visit_expr_method_call_mut (
        self: &'_ mut Self,
        call: &'_ mut ExprMethodCall,
    )
    {
        // 1 - subrecurse
        visit_expr_method_call_mut(self, call);
        // 2 - special case functions whose name ends in `_blue`
        if call.method.to_string().ends_with("_blue") {
            call.args.push(parse_quote!( "hello" ));
        }
    }
}
AppendHelloToBlues.visit_item_fn_mut(&mut code);