Suggested changes to arrow functions transform

[overlookmotel]

Continuing from conversation on #8335 (comment):

I said:

The principle is this: You must make sure that every pop follows a push so the stack and traversal stay in sync. So if you push/pop conditionally e.g.:

fn enter_some_node(...) {
  if node.whatever {
    stack.push(data);
  }
}

fn exit_some_node(...) {
  if node.whatever {
    let data = stack.pop();
  }
}

Then you need to be absolutely sure that the value of node.whatever is the same in both enter_* and exit_*. That includes what any other transform may do to the node in between.

You can only rely on 2 things:

1. Every enter_* visitor call has a corresponding exit_* call.
2. The ancestors of the node are unchanged between enter_* and exit_* (because you cannot mutate upwards).

But all the properties of the node itself are fluid - any other transform can change them in any way.

As much as possible, I think it's ideal if individual transforms are not tightly coupled. When working on one transform, you shouldn't need to understand what other transforms may do to AST. None of us knows all the transforms inside out, so that'd be impossible! So I have tended to code defensively, and assume other transforms could do anything in between enter_* and exit_*.

I try to always push to a stack when entering a node, and always pop when exiting it, so it's impossible for the stack and the traversal to get out of sync. Then use the saved value to decide what to do.

SparseStack is ideal for this because pushing None to it is very cheap indeed (internally it just writes a single bool). Then in exit_* you pop from the stack unconditionally, and decide whether you need to make a modification to the AST or not depending on if the value is Some or None.

@Dunqing replied:

Another thing I want to avoid is passing too many arguments to insert_variable_statement_at_the_top_of_statements, we have to pass all stack values in call site, because only program we need to use last_mut, otherwise we use pop.

2 suggestions to improve on #8335 based on that conversation:

Defensive coding

As I said above, you cannot rely on the node not being changed by other transforms between enter_* and exit_* visitors.

In this case, self.is_async_only() && Self::is_class_method_like_ancestor(ctx.parent()) is guaranteed to be the same in enter_function and exit_function. BUT it's possible that func.r#async is false in enter_function but has been changed to true by another transform before exit_function. That would cause a panic because stack goes out of sync with traversal.

So I suggest instead of:

oxc/crates/oxc_transformer/src/common/arrow_function_converter.rs

Lines 216 to 242 in 335065d

	if self.is_async_only()
	&& (func.r#async || self.super_methods_stack.len() > 1)
	&& Self::is_class_method_like_ancestor(ctx.parent())
	{
	// `self.super_methods_stack.len() > 1` means we are in a nested class method
	//
	// Only `super` that inside async methods need to be transformed, if it is a
	// nested class method and it is not async, we still need to push a `None` to
	// `self.super_methods_stack`, because if we don't get a `FxIndexMap` from
	// `self.super_methods_stack.last_mut()`, that means we don't need to transform.
	// See how to transform `super` in `self.transform_member_expression_for_super`
	//
	// ```js
	// class Outer {
	// async method() {
	// class Inner extends Outer {
	// normal() {
	// // `super.value` should not be transformed, because it is not in an async method
	// super.value
	// }
	// }
	// }
	// }
	// ```
	let super_methods = if func.r#async { Some(FxIndexMap::default()) } else { None };
	self.super_methods_stack.push(super_methods);
	}

if self.is_async_only() && Self::is_class_method_like_ancestor(ctx.parent()) {
    let super_methods = if func.r#async { Some(FxIndexMap::default()) } else { None };
    self.super_methods_stack.push(super_methods);
}

Then in exit_function, do the same.

i.e. Don't rely on func.r#async being the same in enter_function and exit_function to ensure stack stays in sync with traversal.

This does have a small perf cost, because you push to the stack for every function now. But that cost is very small because pushing None to a SparseStack is only a 1-byte write, and SparseStack::push is very optimized. So personally I think the small perf cost is worth it for the greater safety.

However... of course it's a problem if some other transform has converted a non-async class method to an async one between enter_function and exit_function because our transform won't work correctly. But we can catch that in exit_function with a debug assertion:

// Check that another transform hasn't altered value of `func.r#async` since our `enter_function` visitor
debug_assert!(func.r#async == super_methods.is_some());

Don't pass super_methods to insert_variable_statement_at_the_top_of_statements

I think you can achieve your goal of not having to pass super_methods to insert_variable_statement_at_the_top_of_statements by changing:

oxc/crates/oxc_transformer/src/common/arrow_function_converter.rs

Lines 1110 to 1114 in 335065d

	if let Some(super_methods) = super_methods {
	declarations.extend(super_methods.into_iter().map(|(key, super_method)| {
	Self::generate_super_method(target_scope_id, super_method, key.is_assignment, ctx)
	}));
	}

to:

if let Some(super_methods) = self.super_methods_stack.take_last() {
  // same as before
}

Then call self.super_methods_stack.pop() in exit_function after calling insert_variable_statement_at_the_top_of_statements. At that point the popped value will always be None, and you can discard it.

You can also do the same with this_var and arguments_var.