feat(traverse): pass &mut TraverseCtx to visitors (#3312)

Pass `&mut TraverseCtx` to `Traverse::enter_*` and `Traverse::exit_*` visitor methods. Previously was immutable `&TraverseCtx`.

This is a step towards exposing mutable scope tree + symbol table to visitors.

crates/oxc_transformer/src/lib.rs

@@ -92,11 +92,11 @@ impl<'a> Transformer<'a> {
 }
 
 impl<'a> Traverse<'a> for Transformer<'a> {
-    fn enter_program(&mut self, program: &mut Program<'a>, _ctx: &TraverseCtx<'a>) {
+    fn enter_program(&mut self, program: &mut Program<'a>, _ctx: &mut TraverseCtx<'a>) {
         self.x0_typescript.transform_program(program);
     }
 
-    fn exit_program(&mut self, program: &mut Program<'a>, _ctx: &TraverseCtx<'a>) {
+    fn exit_program(&mut self, program: &mut Program<'a>, _ctx: &mut TraverseCtx<'a>) {
         self.x1_react.transform_program_on_exit(program);
         self.x0_typescript.transform_program_on_exit(program);
     }
@@ -106,102 +106,114 @@ impl<'a> Traverse<'a> for Transformer<'a> {
     fn enter_arrow_function_expression(
         &mut self,
         expr: &mut ArrowFunctionExpression<'a>,
-        _ctx: &TraverseCtx<'a>,
+        _ctx: &mut TraverseCtx<'a>,
     ) {
         self.x0_typescript.transform_arrow_expression(expr);
     }
 
-    fn enter_binding_pattern(&mut self, pat: &mut BindingPattern<'a>, _ctx: &TraverseCtx<'a>) {
+    fn enter_binding_pattern(&mut self, pat: &mut BindingPattern<'a>, _ctx: &mut TraverseCtx<'a>) {
         self.x0_typescript.transform_binding_pattern(pat);
     }
 
-    fn enter_call_expression(&mut self, expr: &mut CallExpression<'a>, ctx: &TraverseCtx<'a>) {
+    fn enter_call_expression(&mut self, expr: &mut CallExpression<'a>, ctx: &mut TraverseCtx<'a>) {
         self.x0_typescript.transform_call_expression(expr);
         self.x1_react.transform_call_expression(expr, ctx);
     }
 
-    fn enter_class(&mut self, class: &mut Class<'a>, _ctx: &TraverseCtx<'a>) {
+    fn enter_class(&mut self, class: &mut Class<'a>, _ctx: &mut TraverseCtx<'a>) {
         self.x0_typescript.transform_class(class);
         self.x3_es2015.transform_class(class);
     }
 
-    fn exit_class(&mut self, class: &mut Class<'a>, _ctx: &TraverseCtx<'a>) {
+    fn exit_class(&mut self, class: &mut Class<'a>, _ctx: &mut TraverseCtx<'a>) {
         self.x3_es2015.transform_class_on_exit(class);
     }
 
-    fn enter_class_body(&mut self, body: &mut ClassBody<'a>, _ctx: &TraverseCtx<'a>) {
+    fn enter_class_body(&mut self, body: &mut ClassBody<'a>, _ctx: &mut TraverseCtx<'a>) {
         self.x0_typescript.transform_class_body(body);
     }
 
     fn enter_export_named_declaration(
         &mut self,
         decl: &mut ExportNamedDeclaration<'a>,
-        _ctx: &TraverseCtx<'a>,
+        _ctx: &mut TraverseCtx<'a>,
     ) {
         self.x0_typescript.transform_export_named_declaration(decl);
     }
 
-    fn enter_expression(&mut self, expr: &mut Expression<'a>, ctx: &TraverseCtx<'a>) {
+    fn enter_expression(&mut self, expr: &mut Expression<'a>, ctx: &mut TraverseCtx<'a>) {
         self.x0_typescript.transform_expression(expr);
         self.x1_react.transform_expression(expr, ctx);
         self.x3_es2015.transform_expression(expr);
     }
 
-    fn exit_expression(&mut self, expr: &mut Expression<'a>, _ctx: &TraverseCtx<'a>) {
+    fn exit_expression(&mut self, expr: &mut Expression<'a>, _ctx: &mut TraverseCtx<'a>) {
         self.x3_es2015.transform_expression_on_exit(expr);
     }
 
-    fn enter_formal_parameter(&mut self, param: &mut FormalParameter<'a>, _ctx: &TraverseCtx<'a>) {
+    fn enter_formal_parameter(
+        &mut self,
+        param: &mut FormalParameter<'a>,
+        _ctx: &mut TraverseCtx<'a>,
+    ) {
         self.x0_typescript.transform_formal_parameter(param);
     }
 
-    fn enter_function(&mut self, func: &mut Function<'a>, _ctx: &TraverseCtx<'a>) {
+    fn enter_function(&mut self, func: &mut Function<'a>, _ctx: &mut TraverseCtx<'a>) {
         self.x0_typescript.transform_function(func);
     }
 
-    fn enter_jsx_element(&mut self, node: &mut JSXElement<'a>, _ctx: &TraverseCtx<'a>) {
+    fn enter_jsx_element(&mut self, node: &mut JSXElement<'a>, _ctx: &mut TraverseCtx<'a>) {
         self.x0_typescript.transform_jsx_element(node);
     }
 
-    fn enter_jsx_fragment(&mut self, node: &mut JSXFragment<'a>, _ctx: &TraverseCtx<'a>) {
+    fn enter_jsx_fragment(&mut self, node: &mut JSXFragment<'a>, _ctx: &mut TraverseCtx<'a>) {
         self.x0_typescript.transform_jsx_fragment(node);
     }
 
     fn enter_jsx_opening_element(
         &mut self,
         elem: &mut JSXOpeningElement<'a>,
-        ctx: &TraverseCtx<'a>,
+        ctx: &mut TraverseCtx<'a>,
     ) {
         self.x0_typescript.transform_jsx_opening_element(elem);
         self.x1_react.transform_jsx_opening_element(elem, ctx);
         self.x3_es2015.transform_jsx_opening_element(elem);
     }
 
-    fn enter_method_definition(&mut self, def: &mut MethodDefinition<'a>, _ctx: &TraverseCtx<'a>) {
+    fn enter_method_definition(
+        &mut self,
+        def: &mut MethodDefinition<'a>,
+        _ctx: &mut TraverseCtx<'a>,
+    ) {
         self.x0_typescript.transform_method_definition(def);
     }
 
-    fn exit_method_definition(&mut self, def: &mut MethodDefinition<'a>, _ctx: &TraverseCtx<'a>) {
+    fn exit_method_definition(
+        &mut self,
+        def: &mut MethodDefinition<'a>,
+        _ctx: &mut TraverseCtx<'a>,
+    ) {
         self.x0_typescript.transform_method_definition_on_exit(def);
     }
 
-    fn enter_new_expression(&mut self, expr: &mut NewExpression<'a>, _ctx: &TraverseCtx<'a>) {
+    fn enter_new_expression(&mut self, expr: &mut NewExpression<'a>, _ctx: &mut TraverseCtx<'a>) {
         self.x0_typescript.transform_new_expression(expr);
     }
 
     fn enter_property_definition(
         &mut self,
         def: &mut PropertyDefinition<'a>,
-        _ctx: &TraverseCtx<'a>,
+        _ctx: &mut TraverseCtx<'a>,
     ) {
         self.x0_typescript.transform_property_definition(def);
     }
 
-    fn enter_statements(&mut self, stmts: &mut Vec<'a, Statement<'a>>, _ctx: &TraverseCtx<'a>) {
+    fn enter_statements(&mut self, stmts: &mut Vec<'a, Statement<'a>>, _ctx: &mut TraverseCtx<'a>) {
         self.x3_es2015.enter_statements(stmts);
     }
 
-    fn exit_statements(&mut self, stmts: &mut Vec<'a, Statement<'a>>, _ctx: &TraverseCtx<'a>) {
+    fn exit_statements(&mut self, stmts: &mut Vec<'a, Statement<'a>>, _ctx: &mut TraverseCtx<'a>) {
         self.x0_typescript.transform_statements_on_exit(stmts);
         self.x3_es2015.exit_statements(stmts);
     }
@@ -209,7 +221,7 @@ impl<'a> Traverse<'a> for Transformer<'a> {
     fn enter_tagged_template_expression(
         &mut self,
         expr: &mut TaggedTemplateExpression<'a>,
-        _ctx: &TraverseCtx<'a>,
+        _ctx: &mut TraverseCtx<'a>,
     ) {
         self.x0_typescript.transform_tagged_template_expression(expr);
     }
@@ -217,32 +229,32 @@ impl<'a> Traverse<'a> for Transformer<'a> {
     fn enter_identifier_reference(
         &mut self,
         ident: &mut IdentifierReference<'a>,
-        ctx: &TraverseCtx<'a>,
+        ctx: &mut TraverseCtx<'a>,
     ) {
         self.x0_typescript.transform_identifier_reference(ident, ctx);
     }
 
-    fn enter_statement(&mut self, stmt: &mut Statement<'a>, _ctx: &TraverseCtx<'a>) {
+    fn enter_statement(&mut self, stmt: &mut Statement<'a>, _ctx: &mut TraverseCtx<'a>) {
         self.x0_typescript.transform_statement(stmt);
     }
 
-    fn enter_declaration(&mut self, decl: &mut Declaration<'a>, _ctx: &TraverseCtx<'a>) {
+    fn enter_declaration(&mut self, decl: &mut Declaration<'a>, _ctx: &mut TraverseCtx<'a>) {
         self.x0_typescript.transform_declaration(decl);
         self.x3_es2015.transform_declaration(decl);
     }
 
-    fn exit_declaration(&mut self, decl: &mut Declaration<'a>, _ctx: &TraverseCtx<'a>) {
+    fn exit_declaration(&mut self, decl: &mut Declaration<'a>, _ctx: &mut TraverseCtx<'a>) {
         self.x3_es2015.transform_declaration_on_exit(decl);
     }
 
-    fn enter_if_statement(&mut self, stmt: &mut IfStatement<'a>, _ctx: &TraverseCtx<'a>) {
+    fn enter_if_statement(&mut self, stmt: &mut IfStatement<'a>, _ctx: &mut TraverseCtx<'a>) {
         self.x0_typescript.transform_if_statement(stmt);
     }
 
     fn enter_module_declaration(
         &mut self,
         decl: &mut ModuleDeclaration<'a>,
-        _ctx: &TraverseCtx<'a>,
+        _ctx: &mut TraverseCtx<'a>,
     ) {
         self.x0_typescript.transform_module_declaration(decl);
     }

crates/oxc_traverse/scripts/lib/traverse.mjs

@@ -9,9 +9,9 @@ export default function generateTraverseTraitCode(types) {
         const snakeName = camelToSnake(type.name);
         traverseMethods += `
             #[inline]
-            fn enter_${snakeName}(&mut self, node: &mut ${type.rawName}, ctx: &TraverseCtx<'a>) {}
+            fn enter_${snakeName}(&mut self, node: &mut ${type.rawName}, ctx: &mut TraverseCtx<'a>) {}
             #[inline]
-            fn exit_${snakeName}(&mut self, node: &mut ${type.rawName}, ctx: &TraverseCtx<'a>) {}
+            fn exit_${snakeName}(&mut self, node: &mut ${type.rawName}, ctx: &mut TraverseCtx<'a>) {}
         `;
     }
 

crates/oxc_traverse/src/context.rs

@@ -49,6 +49,22 @@ pub struct TraverseCtx<'a> {
 /// Contains a stack of `Ancestor`s, and provides methods to get parent/ancestor of current node.
 ///
 /// `walk_*` methods push/pop `Ancestor`s to `stack` when entering/exiting nodes.
+///
+/// # SAFETY
+/// This type MUST NOT be mutable by consumer.
+///
+/// The safety scheme is entirely reliant on `stack` being in sync with the traversal,
+/// to prevent consumer from accessing fields of nodes which traversal has passed through,
+/// so as to not violate Rust's aliasing rules.
+/// If consumer could alter `stack` in any way, they could break the safety invariants and cause UB.
+///
+/// We prevent this in 3 ways:
+/// 1. `TraverseAncestry`'s `stack` field is private.
+/// 2. Public methods of `TraverseAncestry` provide no means for mutating `stack`.
+/// 3. Visitors receive a `&mut TraverseCtx`, but cannot overwrite its `ancestry` field because they:
+///    a. cannot create a new `TraverseAncestry` - `TraverseAncestry::new` is private.
+///    b. cannot obtain an owned `TraverseAncestry` from a `&TraverseAncestry`
+///       - `TraverseAncestry` is not `Clone`.
 pub struct TraverseAncestry<'a> {
     stack: Vec<Ancestor<'a>>,
 }
@@ -127,7 +143,7 @@ impl<'a> TraverseCtx<'a> {
     ///
     /// struct MyTraverse;
     /// impl<'a> Traverse<'a> for MyTraverse {
-    ///     fn enter_this_expression(&mut self, this_expr: &mut ThisExpression, ctx: &TraverseCtx<'a>) {
+    ///     fn enter_this_expression(&mut self, this_expr: &mut ThisExpression, ctx: &mut TraverseCtx<'a>) {
     ///         // Get name of function where `this` is bound.
     ///         // NB: This example doesn't handle `this` in class fields or static blocks.
     ///         let fn_id = ctx.find_ancestor(|ancestor| {
@@ -221,6 +237,9 @@ impl<'a> TraverseCtx<'a> {
 // Methods used internally within crate
 impl<'a> TraverseCtx<'a> {
     /// Shortcut for `self.ancestry.push_stack`, to make `walk_*` methods less verbose.
+    ///
+    /// # SAFETY
+    /// This method must not be public outside this crate, or consumer could break safety invariants.
     #[inline]
     pub(crate) fn push_stack(&mut self, ancestor: Ancestor<'a>) {
         self.ancestry.push_stack(ancestor);
@@ -230,6 +249,7 @@ impl<'a> TraverseCtx<'a> {
     ///
     /// # SAFETY
     /// See safety constraints of `TraverseAncestry.pop_stack`.
+    /// This method must not be public outside this crate, or consumer could break safety invariants.
     #[inline]
     #[allow(unsafe_code)]
     pub(crate) unsafe fn pop_stack(&mut self) {
@@ -240,6 +260,7 @@ impl<'a> TraverseCtx<'a> {
     ///
     /// # SAFETY
     /// See safety constraints of `TraverseAncestry.retag_stack`.
+    /// This method must not be public outside this crate, or consumer could break safety invariants.
     #[inline]
     #[allow(unsafe_code, clippy::ptr_as_ptr, clippy::ref_as_ptr)]
     pub(crate) unsafe fn retag_stack(&mut self, ty: AncestorType) {
@@ -288,7 +309,7 @@ impl<'a> TraverseAncestry<'a> {
     ///
     /// struct MyTraverse;
     /// impl<'a> Traverse<'a> for MyTraverse {
-    ///     fn enter_this_expression(&mut self, this_expr: &mut ThisExpression, ctx: &TraverseCtx<'a>) {
+    ///     fn enter_this_expression(&mut self, this_expr: &mut ThisExpression, ctx: &mut TraverseCtx<'a>) {
     ///         // Get name of function where `this` is bound.
     ///         // NB: This example doesn't handle `this` in class fields or static blocks.
     ///         let fn_id = ctx.ancestry.find_ancestor(|ancestor| {
@@ -327,9 +348,12 @@ impl<'a> TraverseAncestry<'a> {
     }
 }
 
-// Methods used internally within crate
+// Methods used internally within crate.
 impl<'a> TraverseAncestry<'a> {
     /// Create new `TraverseAncestry`.
+    ///
+    /// # SAFETY
+    /// This method must not be public outside this crate, or consumer could break safety invariants.
     fn new() -> Self {
         let mut stack = Vec::with_capacity(INITIAL_STACK_CAPACITY);
         stack.push(Ancestor::None);
@@ -337,15 +361,21 @@ impl<'a> TraverseAncestry<'a> {
     }
 
     /// Push item onto ancestry stack.
+    ///
+    /// # SAFETY
+    /// This method must not be public outside this crate, or consumer could break safety invariants.
     #[inline]
     pub(crate) fn push_stack(&mut self, ancestor: Ancestor<'a>) {
         self.stack.push(ancestor);
     }
 
     /// Pop last item off ancestry stack.
+    ///
     /// # SAFETY
     /// * Stack must not be empty.
     /// * Each `pop_stack` call must correspond to a `push_stack` call for same type.
+    ///
+    /// This method must not be public outside this crate, or consumer could break safety invariants.
     #[inline]
     #[allow(unsafe_code)]
     pub(crate) unsafe fn pop_stack(&mut self) {
@@ -370,6 +400,8 @@ impl<'a> TraverseAncestry<'a> {
     /// # SAFETY
     /// * Stack must not be empty.
     /// * Last item on stack must contain pointer to type corresponding to provided `AncestorType`.
+    ///
+    /// This method must not be public outside this crate, or consumer could break safety invariants.
     #[inline]
     #[allow(unsafe_code, clippy::ptr_as_ptr, clippy::ref_as_ptr)]
     pub(crate) unsafe fn retag_stack(&mut self, ty: AncestorType) {

crates/oxc_traverse/src/lib.rs

@@ -114,7 +114,7 @@ mod walk;
 /// struct MyTransform;
 ///
 /// impl<'a> Traverse<'a> for MyTransform {
-///     fn enter_numeric_literal(&mut self, node: &mut NumericLiteral<'a>, ctx: &TraverseCtx<'a>) {
+///     fn enter_numeric_literal(&mut self, node: &mut NumericLiteral<'a>, ctx: &mut TraverseCtx<'a>) {
 ///         // Read parent
 ///         if let Ancestor::BinaryExpressionRight(bin_expr_ref) = ctx.parent() {
 ///             // This is legal

crates/oxc_traverse/tests/compile_fail/ancestor_lifetime1.rs

@@ -9,7 +9,7 @@ impl<'a, 'b> Traverse<'a> for Trans<'a, 'b> {
     fn enter_identifier_reference(
         &mut self,
         _node: &mut IdentifierReference<'a>,
-        ctx: &TraverseCtx<'a>,
+        ctx: &mut TraverseCtx<'a>,
     ) {
         self.ancestor = Some(ctx.parent());
     }

crates/oxc_traverse/tests/compile_fail/ancestor_lifetime1.stderr

@@ -4,7 +4,7 @@ error: lifetime may not live long enough
 8  | impl<'a, 'b> Traverse<'a> for Trans<'a, 'b> {
    |          -- lifetime `'b` defined here
 ...
-12 |         ctx: &TraverseCtx<'a>,
+12 |         ctx: &mut TraverseCtx<'a>,
    |              - let's call the lifetime of this reference `'1`
 13 |     ) {
 14 |         self.ancestor = Some(ctx.parent());

crates/oxc_traverse/tests/compile_fail/ancestor_lifetime2.rs

@@ -9,7 +9,7 @@ impl<'a, 'b> Traverse<'a> for Trans<'a, 'b> {
     fn enter_identifier_reference(
         &mut self,
         _node: &mut IdentifierReference<'a>,
-        ctx: &TraverseCtx<'a>,
+        ctx: &mut TraverseCtx<'a>,
     ) {
         if let Ancestor::ProgramDirectives(program) = ctx.parent() {
             self.program = Some(program);

crates/oxc_traverse/tests/compile_fail/ancestor_lifetime2.stderr

@@ -4,7 +4,7 @@ error: lifetime may not live long enough
 8  | impl<'a, 'b> Traverse<'a> for Trans<'a, 'b> {
    |          -- lifetime `'b` defined here
 ...
-12 |         ctx: &TraverseCtx<'a>,
+12 |         ctx: &mut TraverseCtx<'a>,
    |              - let's call the lifetime of this reference `'1`
 ...
 15 |             self.program = Some(program);

crates/oxc_traverse/tests/compile_fail/ancestor_lifetime3.rs

@@ -10,7 +10,7 @@ impl<'a, 'b> Traverse<'a> for Trans<'a, 'b> {
     fn enter_identifier_reference(
         &mut self,
         _node: &mut IdentifierReference<'a>,
-        ctx: &TraverseCtx<'a>,
+        ctx: &mut TraverseCtx<'a>,
     ) {
         if let Ancestor::ProgramDirectives(program) = ctx.parent() {
             self.program_body = Some(program.body());

crates/oxc_traverse/tests/compile_fail/ancestor_lifetime3.stderr

@@ -4,7 +4,7 @@ error: lifetime may not live long enough
 9  | impl<'a, 'b> Traverse<'a> for Trans<'a, 'b> {
    |          -- lifetime `'b` defined here
 ...
-13 |         ctx: &TraverseCtx<'a>,
+13 |         ctx: &mut TraverseCtx<'a>,
    |              - let's call the lifetime of this reference `'1`
 ...
 16 |             self.program_body = Some(program.body());