refactor(traverse): enter node before entering scope (#4684)

Closes #4200. Align `Traverse`'s behavior with `Visit` and `VisitMut`. For types with scopes, call `enter_*` before entering scope, and call `exit_*` after exiting scope.
oxc-project · Aug 6, 2024 · 83546d3 · 83546d3
1 parent 01d85de
commit 83546d3
Show file tree

Hide file tree

Showing 6 changed files with 90 additions and 71 deletions.
diff --git a/crates/oxc_transformer/src/typescript/annotations.rs b/crates/oxc_transformer/src/typescript/annotations.rs
@@ -7,7 +7,10 @@ use oxc_ast::ast::*;
 use oxc_semantic::SymbolFlags;
 use oxc_span::{Atom, GetSpan, Span, SPAN};
 use oxc_syntax::{
-    operator::AssignmentOperator, reference::ReferenceFlag, scope::ScopeFlags, symbol::SymbolId,
+    operator::AssignmentOperator,
+    reference::ReferenceFlag,
+    scope::{ScopeFlags, ScopeId},
+    symbol::SymbolId,
 };
 use oxc_traverse::TraverseCtx;
 use rustc_hash::FxHashSet;
@@ -419,7 +422,7 @@ impl<'a> TypeScriptAnnotations<'a> {
             }
         }
 
-        Self::replace_with_empty_block_if_ts(&mut stmt.consequent, ctx);
+        Self::replace_with_empty_block_if_ts(&mut stmt.consequent, ctx.current_scope_id(), ctx);
 
         if stmt.alternate.as_ref().is_some_and(Statement::is_typescript_syntax) {
             stmt.alternate = None;
@@ -442,28 +445,33 @@ impl<'a> TypeScriptAnnotations<'a> {
         stmt: &mut ForStatement<'a>,
         ctx: &mut TraverseCtx<'a>,
     ) {
-        Self::replace_with_empty_block_if_ts(&mut stmt.body, ctx);
+        let scope_id = stmt.scope_id.get().unwrap_or(ctx.current_scope_id());
+        Self::replace_with_empty_block_if_ts(&mut stmt.body, scope_id, ctx);
     }
 
     pub fn transform_while_statement(
         &mut self,
         stmt: &mut WhileStatement<'a>,
         ctx: &mut TraverseCtx<'a>,
     ) {
-        Self::replace_with_empty_block_if_ts(&mut stmt.body, ctx);
+        Self::replace_with_empty_block_if_ts(&mut stmt.body, ctx.current_scope_id(), ctx);
     }
 
     pub fn transform_do_while_statement(
         &mut self,
         stmt: &mut DoWhileStatement<'a>,
         ctx: &mut TraverseCtx<'a>,
     ) {
-        Self::replace_with_empty_block_if_ts(&mut stmt.body, ctx);
+        Self::replace_with_empty_block_if_ts(&mut stmt.body, ctx.current_scope_id(), ctx);
     }
 
-    fn replace_with_empty_block_if_ts(stmt: &mut Statement<'a>, ctx: &mut TraverseCtx<'a>) {
+    fn replace_with_empty_block_if_ts(
+        stmt: &mut Statement<'a>,
+        parent_scope_id: ScopeId,
+        ctx: &mut TraverseCtx<'a>,
+    ) {
         if stmt.is_typescript_syntax() {
-            let scope_id = ctx.create_scope_child_of_current(ScopeFlags::empty());
+            let scope_id = ctx.create_child_scope(parent_scope_id, ScopeFlags::empty());
             let block = BlockStatement {
                 span: stmt.span(),
                 body: ctx.ast.vec(),

diff --git a/crates/oxc_transformer/src/typescript/namespace.rs b/crates/oxc_transformer/src/typescript/namespace.rs
@@ -254,10 +254,9 @@ impl<'a> TypeScript<'a> {
         }
 
         if new_stmts.is_empty() {
-            // Delete the scope binding that `ctx.generate_uid_in_current_scope` created above,
+            // Delete the scope binding that `ctx.generate_uid` created above,
             // as no binding is actually being created
-            let current_scope_id = ctx.current_scope_id();
-            ctx.scopes_mut().remove_binding(current_scope_id, &CompactStr::from(name.as_str()));
+            ctx.scopes_mut().remove_binding(scope_id, &CompactStr::from(name.as_str()));
 
             return None;
         }

diff --git a/crates/oxc_traverse/scripts/lib/walk.mjs b/crates/oxc_traverse/scripts/lib/walk.mjs
@@ -69,7 +69,8 @@ function generateWalkForStruct(type, types) {
                 `\`ast\` attr says to enter scope before field '${enterFieldName}' `
                 + `in '${type.name}', but that field is not visited`
             );
-            if (scopeEnterField === visitedFields[0]) scopeEnterField = undefined;
+        } else {
+            scopeEnterField = visitedFields[0];
         }
 
         // TODO: Maybe this isn't quite right. `scope_id` fields are `Cell<Option<ScopeId>>`,
@@ -93,18 +94,25 @@ function generateWalkForStruct(type, types) {
     }
 
     const fieldsCodes = visitedFields.map((field, index) => {
-        const fieldWalkName = `walk_${camelToSnake(field.innerTypeName)}`;
+        const fieldWalkName = `walk_${camelToSnake(field.innerTypeName)}`,
+            fieldCamelName = snakeToCamel(field.name);
+        const scopeCode = field === scopeEnterField ? enterScopeCode : '';
 
-        const retagCode = index === 0
-            ? ''
-            : `ctx.retag_stack(AncestorType::${type.name}${snakeToCamel(field.name)});`;
-        const fieldCode = makeFieldCode(field);
-        let scopeCode = '';
-        if (field === scopeEnterField) {
-            scopeCode = enterScopeCode;
-            enterScopeCode = '';
+        let tagCode = '', retagCode = '';
+        if (index === 0) {
+            tagCode = `
+                ctx.push_stack(
+                    Ancestor::${type.name}${fieldCamelName}(
+                        ancestor::${type.name}Without${fieldCamelName}(node)
+                    )
+                );
+            `;
+        } else {
+            retagCode = `ctx.retag_stack(AncestorType::${type.name}${fieldCamelName});`;
         }
 
+        const fieldCode = makeFieldCode(field);
+
         if (field.wrappers[0] === 'Option') {
             let walkCode;
             if (field.wrappers.length === 2 && field.wrappers[1] === 'Vec') {
@@ -127,6 +135,7 @@ function generateWalkForStruct(type, types) {
 
             return `
                 ${scopeCode}
+                ${tagCode}
                 if let Some(field) = &mut *(${fieldCode}) {
                     ${retagCode}
                     ${walkCode}
@@ -159,15 +168,15 @@ function generateWalkForStruct(type, types) {
 
             return `
                 ${scopeCode}
-                ${retagCode}
+                ${tagCode || retagCode}
                 ${walkVecCode}
             `;
         }
 
         if (field.wrappers.length === 1 && field.wrappers[0] === 'Box') {
             return `
                 ${scopeCode}
-                ${retagCode}
+                ${tagCode || retagCode}
                 ${fieldWalkName}(traverser, (&mut **(${fieldCode})) as *mut _, ctx);
             `;
         }
@@ -176,23 +185,12 @@ function generateWalkForStruct(type, types) {
 
         return `
             ${scopeCode}
-            ${retagCode}
+            ${tagCode || retagCode}
             ${fieldWalkName}(traverser, ${fieldCode}, ctx);
         `;
     });
 
-    if (visitedFields.length > 0) {
-        const field = visitedFields[0],
-            fieldCamelName = snakeToCamel(field.name);
-        fieldsCodes.unshift(`
-            ctx.push_stack(
-                Ancestor::${type.name}${fieldCamelName}(
-                    ancestor::${type.name}Without${fieldCamelName}(node)
-                )
-            );
-        `);
-        fieldsCodes.push('ctx.pop_stack();');
-    }
+    if (visitedFields.length > 0) fieldsCodes.push('ctx.pop_stack();');
 
     const typeSnakeName = camelToSnake(type.name);
     return `
@@ -201,12 +199,10 @@ function generateWalkForStruct(type, types) {
             node: *mut ${type.rawName},
             ctx: &mut TraverseCtx<'a>
         ) {
-            ${enterScopeCode}
             traverser.enter_${typeSnakeName}(&mut *node, ctx);
             ${fieldsCodes.join('\n')}
-            ${enterScopeCode ? '' : exitScopeCode}
+            ${exitScopeCode}
             traverser.exit_${typeSnakeName}(&mut *node, ctx);
-            ${enterScopeCode ? exitScopeCode : ''}
         }
     `.replace(/\n\s*\n+/g, '\n');
 }

diff --git a/crates/oxc_traverse/src/context/mod.rs b/crates/oxc_traverse/src/context/mod.rs
@@ -286,13 +286,22 @@ impl<'a> TraverseCtx<'a> {
         self.scoping.find_scope_by_flags(finder)
     }
 
+    /// Create new scope as child of provided scope.
+    ///
+    /// `flags` provided are amended to inherit from parent scope's flags.
+    ///
+    /// This is a shortcut for `ctx.scoping.create_child_scope`.
+    pub fn create_child_scope(&mut self, parent_id: ScopeId, flags: ScopeFlags) -> ScopeId {
+        self.scoping.create_child_scope(parent_id, flags)
+    }
+
     /// Create new scope as child of current scope.
     ///
     /// `flags` provided are amended to inherit from parent scope's flags.
     ///
-    /// This is a shortcut for `ctx.scoping.create_scope_child_of_current`.
-    pub fn create_scope_child_of_current(&mut self, flags: ScopeFlags) -> ScopeId {
-        self.scoping.create_scope_child_of_current(flags)
+    /// This is a shortcut for `ctx.scoping.create_child_scope_of_current`.
+    pub fn create_child_scope_of_current(&mut self, flags: ScopeFlags) -> ScopeId {
+        self.scoping.create_child_scope_of_current(flags)
     }
 
     /// Insert a scope into scope tree below a statement.

diff --git a/crates/oxc_traverse/src/context/scoping.rs b/crates/oxc_traverse/src/context/scoping.rs
@@ -116,12 +116,19 @@ impl TraverseScoping {
         })
     }
 
+    /// Create new scope as child of provided scope.
+    ///
+    /// `flags` provided are amended to inherit from parent scope's flags.
+    pub fn create_child_scope(&mut self, parent_id: ScopeId, flags: ScopeFlags) -> ScopeId {
+        let flags = self.scopes.get_new_scope_flags(flags, parent_id);
+        self.scopes.add_scope(parent_id, AstNodeId::DUMMY, flags)
+    }
+
     /// Create new scope as child of current scope.
     ///
     /// `flags` provided are amended to inherit from parent scope's flags.
-    pub fn create_scope_child_of_current(&mut self, flags: ScopeFlags) -> ScopeId {
-        let flags = self.scopes.get_new_scope_flags(flags, self.current_scope_id);
-        self.scopes.add_scope(self.current_scope_id, AstNodeId::DUMMY, flags)
+    pub fn create_child_scope_of_current(&mut self, flags: ScopeFlags) -> ScopeId {
+        self.create_child_scope(self.current_scope_id, flags)
     }
 
     /// Insert a scope into scope tree below a statement.
@@ -162,7 +169,7 @@ impl TraverseScoping {
         }
 
         // Create new scope as child of parent
-        let new_scope_id = self.create_scope_child_of_current(flags);
+        let new_scope_id = self.create_child_scope_of_current(flags);
 
         // Set scopes as children of new scope instead
         for &child_id in child_scope_ids {