Codegen functions to walk a node's children

Continuation from https://github.com/oxc-project/oxc/pull/12524#issuecomment-3134424218.

The problem

Currently we have 2 sets of functions for Visit. e.g. for Function:

pub trait Visit<'a>: Sized {
    fn visit_function(&mut self, it: &Function<'a>, flags: ScopeFlags) {
        walk_function(self, it, flags);
    }
}

pub fn walk_function<'a, V: Visit<'a>>(visitor: &mut V, it: &Function<'a>, flags: ScopeFlags) {
    let kind = AstKind::Function(visitor.alloc(it));
    visitor.enter_node(kind);
    visitor.enter_scope(
        {
            let mut flags = flags;
            if it.has_use_strict_directive() { flags |= ScopeFlags::StrictMode; }
            flags
        },
        &it.scope_id,
    );
    visitor.visit_span(&it.span);
    if let Some(id) = &it.id { visitor.visit_binding_identifier(id); }
    if let Some(type_parameters) = &it.type_parameters { visitor.visit_ts_type_parameter_declaration(type_parameters); }
    if let Some(this_param) = &it.this_param { visitor.visit_ts_this_parameter(this_param); }
    visitor.visit_formal_parameters(&it.params);
    if let Some(return_type) = &it.return_type { visitor.visit_ts_type_annotation(return_type); }
    if let Some(body) = &it.body { visitor.visit_function_body(body); }
    visitor.leave_scope();
    visitor.leave_node(kind);
}

When writing a Visit impl with a visit_function method, you can do:

impl<'a> Visit<'a> for MyVisitor {
    fn visit_function(&mut self, it: &Function<'a>, flags: ScopeFlags) {
        self.do_stuff_with_function(it);
        walk::walk_function(self, it, flags);
    }
}

But if you want to inject logic to run after enter_node or before exit_node, you have to reimplement all of the field visitation logic from walk_function.

impl<'a> Visit<'a> for MyVisitor {
    fn visit_function(&mut self, func: &Function<'a>, flags: ScopeFlags) {
        let kind = AstKind::Function(visitor.alloc(func));
        visitor.enter_node(kind);
        visitor.enter_scope( /* ... */ );

        self.do_stuff_with_function_before(func);

        self.visit_span(&func.span);
        // ... all the other fields ....
        if let Some(body) = &func.body { self.visit_function_body(body); }

        self.do_stuff_with_function_after(func);

        visitor.leave_scope();
        visitor.leave_node(kind);
    }
}

This is both verbose, and error-prone. When we alter the AST, there's more changes to make, or (worse) more scope for errors due to forgetting to make those changes everywhere.

An example of this problem in practice is in SemanticBuilder.

Possible solution 1

Break out logic for walking children into separate walk_children_* functions:

pub fn walk_function<'a, V: Visit<'a>>(visitor: &mut V, it: &Function<'a>, flags: ScopeFlags) {
    let kind = AstKind::Function(visitor.alloc(it));
    visitor.enter_node(kind);
    visitor.enter_scope(
        {
            let mut flags = flags;
            if it.has_use_strict_directive() { flags |= ScopeFlags::StrictMode; }
            flags
        },
        &it.scope_id,
    );

    walk_children_function(visitor, it, flags);

    visitor.leave_scope();
    visitor.leave_node(kind);
}

pub fn walk_children_function<'a, V: Visit<'a>>(visitor: &mut V, it: &Function<'a>, flags: ScopeFlags) {
    visitor.visit_span(&it.span);
    if let Some(id) = &it.id { visitor.visit_binding_identifier(id); }
    if let Some(type_parameters) = &it.type_parameters { visitor.visit_ts_type_parameter_declaration(type_parameters); }
    if let Some(this_param) = &it.this_param { visitor.visit_ts_this_parameter(this_param); }
    visitor.visit_formal_parameters(&it.params);
    if let Some(return_type) = &it.return_type { visitor.visit_ts_type_annotation(return_type); }
    if let Some(body) = &it.body { visitor.visit_function_body(body); }
}

Now your custom visitor can use walk_children_function:

impl<'a> Visit<'a> for MyVisitor {
    fn visit_function(&mut self, func: &Function<'a>, flags: ScopeFlags) {
        let kind = AstKind::Function(visitor.alloc(func));
        visitor.enter_node(kind);
        visitor.enter_scope( /* ... */ );

        self.do_stuff_with_function_before(func);
        walk_children_function(self, func, flags);
        self.do_stuff_with_function_after(func);

        visitor.leave_scope();
        visitor.leave_node(kind);
    }
}

Possible solution 2

More complicated functions, but less code in the custom visitor:

pub fn walk_function_hooked<'a, V: Visit<'a>>(
    visitor: &mut V,
    it: &Function<'a>,
    flags: ScopeFlags,
    before_hook: impl FnOnce(&Function<'a>),
    after_hook: impl FnOnce(&Function<'a>),
) {
    let kind = AstKind::Function(visitor.alloc(it));
    visitor.enter_node(kind);
    visitor.enter_scope(
        {
            let mut flags = flags;
            if it.has_use_strict_directive() { flags |= ScopeFlags::StrictMode; }
            flags
        },
        &it.scope_id,
    );

    before_hook(it);

    visitor.visit_span(&it.span);
    if let Some(id) = &it.id { visitor.visit_binding_identifier(id); }
    if let Some(type_parameters) = &it.type_parameters { visitor.visit_ts_type_parameter_declaration(type_parameters); }
    if let Some(this_param) = &it.this_param { visitor.visit_ts_this_parameter(this_param); }
    visitor.visit_formal_parameters(&it.params);
    if let Some(return_type) = &it.return_type { visitor.visit_ts_type_annotation(return_type); }
    if let Some(body) = &it.body { visitor.visit_function_body(body); }

    after_hook(it);

    visitor.leave_scope();
    visitor.leave_node(kind);
}

Now your visitor doesn't have to set the right scope flags etc either:

impl<'a> Visit<'a> for MyVisitor {
    fn visit_function(&mut self, func: &Function<'a>, flags: ScopeFlags) {
        walk_function_hooked(
            self, func, flags,
            |_| self.do_stuff_with_function_before(func),
            |_| self.do_stuff_with_function_after(func),
        );
    }
}

Possible solution 3

Even more complicated, but more flexible:

pub trait VisitFunctionHooks {
    fn before_enter_scope<'a, V: Visit<'a>>(&mut self, visitor: &mut V, it: &Function<'a>) {}
    fn before_span<'a, V: Visit<'a>>(&mut self, visitor: &mut V, it: &Function<'a>) {}
    fn before_id<'a, V: Visit<'a>>(&mut self, visitor: &mut V, it: &Function<'a>) {}
    fn before_type_parameters<'a, V: Visit<'a>>(&mut self, visitor: &mut V, it: &Function<'a>) {}
    fn before_this_param<'a, V: Visit<'a>>(&mut self, visitor: &mut V, it: &Function<'a>) {}
    fn before_params<'a, V: Visit<'a>>(&mut self, visitor: &mut V, it: &Function<'a>) {}
    fn before_return_type<'a, V: Visit<'a>>(&mut self, visitor: &mut V, it: &Function<'a>) {}
    fn before_body<'a, V: Visit<'a>>(&mut self, visitor: &mut V, it: &Function<'a>) {}
    fn before_leave_scope<'a, V: Visit<'a>>(&mut self, visitor: &mut V, it: &Function<'a>) {}
    fn before_leave_node<'a, V: Visit<'a>>(&mut self, visitor: &mut V, it: &Function<'a>) {}
}

pub fn walk_function_hooked<'a, V, H>(
    visitor: &mut V,
    it: &Function<'a>,
    flags: ScopeFlags,
    hooks: &mut H,
) where
    V: Visit<'a>,
    H: VisitFunctionHooks,
{
    let kind = AstKind::Function(visitor.alloc(it));
    visitor.enter_node(kind);

    hooks.before_enter_scope(visitor, it);
    visitor.enter_scope(
        {
            let mut flags = flags;
            if it.has_use_strict_directive() { flags |= ScopeFlags::StrictMode; }
            flags
        },
        &it.scope_id,
    );

    hooks.before_span(visitor, it);
    visitor.visit_span(&it.span);

    hooks.before_id(visitor, it);
    if let Some(id) = &it.id { visitor.visit_binding_identifier(id); }

    hooks.before_type_parameters(visitor, it);
    if let Some(type_parameters) = &it.type_parameters { visitor.visit_ts_type_parameter_declaration(type_parameters); }

    hooks.before_this_param(visitor, it);
    if let Some(this_param) = &it.this_param { visitor.visit_ts_this_parameter(this_param); }

    hooks.before_params(visitor, it);
    visitor.visit_formal_parameters(&it.params);

    hooks.before_return_type(visitor, it);
    if let Some(return_type) = &it.return_type { visitor.visit_ts_type_annotation(return_type); }

    hooks.before_body(visitor, it);
    if let Some(body) = &it.body { visitor.visit_function_body(body); }

    hooks.before_leave_scope(visitor, it);
    visitor.leave_scope();

    hooks.before_leave_node(visitor, it);
    visitor.leave_node(kind);
}

impl<'a> Visit<'a> for MyVisitor {
    fn visit_function(&mut self, func: &Function<'a>, flags: ScopeFlags) {
        struct Hooks;
        impl VisitFunctionHooks for Hooks {
            fn before_span(&mut self, visitor: &mut MyVisitor, func: &Function<'_>) {
                visitor.do_stuff_with_function_before(func);
            }
            fn before_body(&mut self, visitor: &mut MyVisitor, func: &Function<'_>) {
                // This is what we couldn't do before.
                // Runs after visiting params, but before visiting body.
                visitor.do_stuff_with_function_before_body(func);
            }
            fn before_leave_scope(&mut self, visitor: &mut MyVisitor, func: &Function<'_>) {
                visitor.do_stuff_with_function_after(func);
            }
        }

        let mut hooks = Hooks;
        walk_function_hooked(self, func, flags, &mut hooks);
    }
}

Maybe there's a nicer API which would allow using closures for the hooks. But point is to allow a way to hook into every stage of the visitation.

What to do?

I suspect the 3rd solution is too complicated. But probably worthwhile implementing at least Solution 1, and possibly Solution 2 (before and after hooks).

Any thoughts anyone?

Start with the Visit trait, walk_function, the proposed walk_children_function and hooked variants, and the SemanticBuilder example described in the issue. First determine which hook API is wanted; done would require an agreed design and corresponding visitor-walking changes, but the issue does not identify files or tests.