Search code examples
rustabstract-syntax-tree

How to insert Expr when using VisitMut?


I'm currently using syn following an example to create an AST that can be mutated. I understand that I can modify the node I'm travesing (as shown below in my current code) but I'm curious if I can add some code in between the current node and the next node. Is the syn crate capable of this?

use syn::visit_mut::{self, VisitMut};
use syn::Expr;

#[derive(Debug)]
struct MyStruct;

impl VisitMut for MyStruct {
    fn visit_expr_mut(&mut self, node: &mut Expr) {
        if let Expr::MethodCall(expr) = &node.to_owned() {
            // I can modify the existing node like so:
            *node = parse_quote!("// Hello World");
            
            // How could I add something after this node and before the next?
        }
    }
}

pub fn create() {
    let current_dir = std::env::current_dir().expect("Unable to get current directory");
    let rust_file = std::fs::read_to_string(current_dir.join("src").join("lib.rs")).expect("Unable to read rust file");
    let ast = syn::parse_file(&rust_file).expect("Unable to create AST from rust file");

    MyStruct.visit_file_mut(&mut ast);
}

Edit to show use case:

The file I'm currently parsing looks like:

#[macro_use]
extern crate foo;
mod test;
fn init(handle: foo::InitHandle) {
    handle.add_class::<Test::test>();
}

Let's say that when I read the AST, I want to add another mod and another handle for it like so:

#[macro_use]
extern crate foo;
mod test;
mod store;
fn init(handle: foo::InitHandle) {
    handle.add_class::<Test::test>();
    handle.add_class::<Store::store>();
}

Solution

  • As I commented, it highly depends on what you want to insert. Because you can't just insert anything before or after node easily.

    For your specific case, you could use parse_quote! to produce an ExprBlock.

    *node = parse_quote!(
        {
            #expr;
            handle.add_class::<Store::store>();
        }
    );
    

    Which with the following input:

    fn init(handle: foo::InitHandle) {
        handle.add_class::<Test::test>();
    }
    

    Would produce this output:

    fn init(handle: foo::InitHandle) {
        {
            handle.add_class::<Test::test>();
            handle.add_class::<Store::store>();
        };
    }
    

    (Note I have reformatted the output, to be prettier)


    Alternatively, you could override visit_block_mut() instead. That way you'd have access to stmts: Vec<Stmt>, and would be able to insert before and after a Stmt. The downside is that by doing it that way, you wouldn't be able to easily visit all Exprs, as by using visit_expr_mut().