Q: How to manipulate nodes+parent_nodes in AST & generate the tree back to JS code?

[ShayRubach]

Hey. My project requires manipulating js files. In example, changing all "A()" function calls to "B(A())".
Currently I have a js+babel implementation that is doing the basic following:

• parsing a js file into AST
• traversing it and manipulating wtvr is needed
• generating the new manipulated ast tree back to js code
• profit!
  Unfortunately, babel seems to be performance-bottlenecked on large files, thus I started experimenting with SWC-Rust.
  I have made a small program as a POC. I managed to parse my js file, traverse it and manipulate (well, access) Identifier nodes by implementing the Fold trait. See code at the end of the post.

Now, I have 2 main questions:

1. The trait methods API only provides the current node, without its parent. How do I get the direct parent then? I've seen some example using Visitors, but I couldn't run it since I couldn't find the exact same API in the docs that is using the parent: Node.
2. How do I generate the code back from the AST? couldn't find it either.

Thanks!

main.js:

extern crate swc_common;
extern crate swc_ecma_parser;
use swc_ecma_ast::{EsVersion, Ident};
use swc_ecma_visit::{Fold};
use swc::{Compiler};
use swc_common::sync::Lrc;
use swc_common::{errors::{ColorConfig, Handler}, SourceMap};
use std::{path::Path};
use swc_ecma_parser::Syntax;


fn main() {
    let cm: Lrc<SourceMap> = Default::default();
    let handler = Handler::with_tty_emitter(ColorConfig::Auto, true, false, Some(cm.clone()));
    let fm3 = cm.load_file(Path::new("test.js")).unwrap();

    struct Manipulator {}
    impl Fold for Manipulator {
        fn fold_ident(&mut self, n: Ident) -> Ident {
            println!("{:?}", n);
            return n;
        }
    }

    let comp = Compiler::new(cm);
    let _js = comp.parse_js(fm3, &handler, EsVersion::Es2022, Syntax::Es(Default::default()), swc::config::IsModule::Bool(false), false).unwrap();
    let _trns = comp.transform(&handler, _js, false, Manipulator{});
}

Cargo.toml:

[package]
name = "rust"
version = "0.1.0"
edition = "2021"

# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html

[dependencies]
structopt = { version = "0.3.15" }
swc = { git = "https://github.com/swc-project/swc" }
swc_atoms= { git = "https://github.com/swc-project/swc" }
swc_ecma_ast= { git = "https://github.com/swc-project/swc" }
swc_ecma_codegen= { git = "https://github.com/swc-project/swc" }
swc_ecma_parser= { git = "https://github.com/swc-project/swc" }
swc_ecma_preset_env= { git = "https://github.com/swc-project/swc" }
swc_ecma_transforms= { git = "https://github.com/swc-project/swc" }
swc_ecma_visit= { git = "https://github.com/swc-project/swc" }
swc_visit= { git = "https://github.com/swc-project/swc" }

[dependencies.swc_common]
git = "https://github.com/swc-project/swc"
features = ["tty-emitter"]

[RiESAEX]

for the second question, I think you can use comp.process_js_with_custom_pass instead of parse_js and transform.

[ShayRubach]

Ok so let's look at it from the other direction. Can I get to a node (which ought to be the 'parent') and fetch its children?

If, for example, I want to wrap all arguments inside a function with another function, like modifying this:
int x = Foo(arg1, arg2, arg3)
into this
int x = Foo(WrapperFunc(arg1, arg2, arg3))
Can I look for the declaration of Foo and then dive down and get its arguments (children) and only them? So that I could create a new expression and replace it with the original one that was inside Foo()?

Thanks!

[RiESAEX]

just a example. I'm not sure how many types of node can be a parent of Expr

struct MyVisitor{
    parent: ExprStmt
}
impl swc_ecma_visit::VisitMut for MyVisitor {
   fn visit_mut_expr_stmt(&mut self, expr: &mut ExprStmt){
       self.parent = expr.clone();
       expr.visit_mut_children_with(self);
   }
   fn visit_mut_expr(&mut self, expr: &mut swc_ecma_ast::Expr) {
       // now you can get parent from `self.parent`
       expr.visit_mut_children_with(self);
   }
}

Usually I would collect the child nodes, but it depends on the usage.
btw if you just need a transform like a()->b(a()), I think you don't need parent node, just visit_mut_call_expr, and *expr = CallExpr{ args: expr.take(),..} will work

[ShayRubach]

Interesting. So the approach here was to get to the node, replace its parent with itself and dive down to the children while keeping a reference to the parent. Correct?
Another question, can you elaborate on the .take()? I've read a little about it. I'm guessing it is straight forward copying all of the args from the original expression?

Also, following that, I was wondering this - if I use visit_mut_children_with as in:

fn visit_mut_expr(&mut self, expr: &mut swc_ecma_ast::Expr) {
       // now you can get parent from `self.parent`
       expr.visit_mut_children_with(self);
}

and one of its children is also of type swc_ecma_ast::Expr, wouldn't it cause a double (or more) iteration over the same node? or there's a bool that makes sure to 'mark' nodes that were visited?

Thanks a lot!

[RiESAEX]

no, and no need.
Although visit_mut_expr will be called several times, the node will not be visit again.
It's like DFS, and call visit_mut_children_with is to make visitor go deeper in ast.
and I think I write something wrong in the example about parent node.

fn visit_mut_expr_stmt(&mut self, expr: &mut ExprStmt){
       let old = self.parent; 
       self.parent = expr.clone();
       expr.visit_mut_children_with(self);
       self.parent = old; // otherwise parent will be the deepest exprstmt, just a pseudo code.
}

[ShayRubach]

Thanks 👍

[RiESAEX]

struct MyVisitor;
impl swc_ecma_visit::VisitMut for MyVisitor {
    fn visit_mut_call_expr(&mut self, expr: &mut CallExpr) {
        expr.visit_mut_children_with(self);
        match &mut expr.callee { 
            Callee::Expr(callee) => match &mut **callee { // callee is &mut Box<Expr>
                Expr::Ident(Ident { sym, .. }) => {
                    if sym.to_string() == "Foo" {
                        *expr = CallExpr {
                            span: DUMMY_SP,
                            callee: swc_ecma_ast::Callee::Expr(callee.take()), // here we want a Box<Expr>, and callee will not use in other place, so we take() it, we got a Box<Expr>,and the origin memory of callee will be replace with dummy value
                            args: vec![CallExpr {
                                span: DUMMY_SP,
                                callee: private_ident!("WrapperFunc").as_callee(),
                                args: expr.args.take(),
                                type_args: Default::default(),
                            }
                            .as_arg()],
                            type_args: Default::default(),
                        }
                    }
                }
                _ => {}
            },
            _ => {}
        }
    }
}
#[cfg(test)]
mod tests{
test!(
        ::swc_ecma_parser::Syntax::default(),
        |_| {
            as_folder(MyVisitor)
        },
        my,
        "let x = Foo(arg1, arg2, arg3)",
        r#"let x = Foo(WrapperFunc(arg1, arg2, arg3))"#
    );
}

[ShayRubach]

Thanks @RiESAEX.
I'm only a day into Rust so I'll have to explore all of the bits in your code example to make sure I'm covered language-wise and mark it as the answer afterwards 🙇

[ShayRubach]

I'm back with another question. Since there's still no AST Explorer and I have to figure out what every visit_mut_* type is related to what code expression (CallExp, VarDeclaration etc.), is there a way I can do something naive like this (psuedo):

visit_mut_every_node_in_tree(self, node) {
    if node.sym == "my keyword name goes here" {    
        if node.type == "the type I am looking for" {   //var declaration, right value of an expression, call expression etc...
            //do stuff
        }
    }
}

The reason I am asking this is because I find the AST representation with the API much more complex than babel (e.g).
Example of an input could be something like this:

eval(`1+1`)
var evalHook = eval
evalHook = eval
let url1 = location.href
let url2 = window.location.href
var obj1 = { eval: eval, evalExp: eval(1+1)}

where I wish to manipulate it into this:

eval(WrapperFunc(`1+1`))
var evalHook = myOwnEval
evalHook = myOwnEval
let url1 = myOwnLocation.href
let url2 = GetCallerFunc(window).location.href
var obj1 = { eval: myOwnEval, evalExp: eval(WrapperFunc(1+1))}

There are just so many types that I'm getting lost inside the API. I realize it might cost in performance loss as I run an extra if on every node. But I need some progress to be made for my POC.

Thanks.

[RiESAEX]

For every struct that has an impl VisitMut, every visit_mut_* has a default implementation which will visit its children.
If you define a function like visit_mut_decl then the default implementation will be overridden.
so if you want to make it go deeper then you have to call visit_mut_children_with manually.
The default implementation is generated by macro in swc_visit_macros and swc_ecma_visit, you can use cargo expand in swc_ecma_visit to get the result of marco expand.

[ShayRubach]

I see. So there's no way to naively get all of the "CallExpr" nodes by just implementing visit_mut_call_expr.
This means I'll have to think about all the possible cases where eval may appear (as a hook, func call, rvalue, etc.) and impl the parent's visit_mut_* expression->explicitly visit its children in order to get my desired CallExp.

Did I get it right?

[kdy1]

No, you will get all eval call even with one visit_mut_call_expr

[ShayRubach]

So I mustn't impl other types that may include CallExpr to avoid losing my traces OR do impl them with the macro's impl + my own added code. I get it now. My problem was that I implemented (empty impl) another type that had a child CallExpr.
Thanks!

[kdy1]

You can override methods, just call node.visit_mut_children(self) from the overrided handler.