feat(ast_tools): support #[scope(exit_before)]

crates/oxc_traverse/scripts/lib/parse.mjs

@@ -57,6 +57,7 @@ const FILENAMES = ['js.rs', 'jsx.rs', 'literal.rs', 'ts.rs'];
  * @property {string} flags
  * @property {string | null} strictIf
  * @property {string | null} enterScopeBefore
+ * @property {string | null} exitScopeBefore
  */
 
 /**
@@ -87,9 +88,18 @@ class Position {
     this.index = index;
   }
 
+  /**
+   * @param {unknown} condition
+   * @param {string} [message]
+   *
+   * @returns {asserts condition}
+   */
   assert(condition, message) {
     if (!condition) this.throw(message);
   }
+  /**
+   * @param {string} [message]
+   */
   throw(message) {
     throw new Error(`${message || 'Unknown error'} (at ${this.filename}:${this.index + 1})`);
   }
@@ -198,12 +208,14 @@ function parseStruct(name, rawName, lines, scopeArgs) {
   const fields = [];
 
   while (!lines.isEnd()) {
-    let isScopeEntry = false, line;
+    let isScopeEntry = false, isScopeExit = false, line;
     while (!lines.isEnd()) {
       line = lines.next();
       if (line === '') continue;
       if (line === '#[scope(enter_before)]') {
         isScopeEntry = true;
+      } else if (line === '#[scope(exit_before)]') {
+        isScopeExit = true;
       } else if (line.startsWith('#[')) {
         while (!line.endsWith(']')) {
           line = lines.next();
@@ -222,6 +234,7 @@ function parseStruct(name, rawName, lines, scopeArgs) {
     fields.push({ name, typeName, rawName, rawTypeName, innerTypeName, wrappers });
 
     if (isScopeEntry) scopeArgs.enterScopeBefore = name;
+    if (isScopeExit) scopeArgs.exitScopeBefore = name;
   }
   return { kind: 'struct', name, rawName, fields, scopeArgs };
 }
@@ -284,13 +297,25 @@ function parseScopeArgs(lines, scopeArgs) {
 const SCOPE_ARGS_KEYS = { flags: 'flags', strict_if: 'strictIf' };
 
 /**
+ * @param {string} argsStr
+ * @param {ScopeArgs| null} args
+ * @param {Position} position
+ *
  * @returns {ScopeArgs}
  */
 function parseScopeArgsStr(argsStr, args, position) {
-  if (!args) args = { flags: 'ScopeFlags::empty()', strictIf: null, enterScopeBefore: null };
+  if (!args) {
+    args = {
+      flags: 'ScopeFlags::empty()',
+      strictIf: null,
+      enterScopeBefore: null,
+      exitScopeBefore: null,
+    };
+  }
 
   if (!argsStr) return args;
 
+  /** @param {RegExp} regex */
   const matchAndConsume = (regex) => {
     const match = argsStr.match(regex);
     position.assert(match);

crates/oxc_traverse/scripts/lib/walk.mjs

@@ -75,7 +75,9 @@ function generateWalkForStruct(type, types) {
 
   const { scopeArgs } = type;
   /** @type {Field | undefined} */
-  let scopeEnterField;
+  let scopeEnterField,
+    /** @type {Field | undefined} */
+    scopeExitField;
   let enterScopeCode = '', exitScopeCode = '';
 
   if (scopeArgs && scopeIdField) {
@@ -92,6 +94,17 @@ function generateWalkForStruct(type, types) {
       scopeEnterField = visitedFields[0];
     }
 
+    // Get field to exit scope before
+    const exitFieldName = scopeArgs.exitScopeBefore;
+    if (exitFieldName) {
+      scopeExitField = visitedFields.find(field => field.name === exitFieldName);
+      assert(
+        scopeExitField,
+        `\`ast\` attr says to exit scope before field '${exitFieldName}' ` +
+          `in '${type.name}', but that field is not visited`,
+      );
+    }
+
     // TODO: Maybe this isn't quite right. `scope_id` fields are `Cell<Option<ScopeId>>`,
     // so visitor is able to alter the `scope_id` of a node from higher up the tree,
     // but we don't take that into account.
@@ -107,7 +120,11 @@ function generateWalkForStruct(type, types) {
   const fieldsCodes = visitedFields.map((field, index) => {
     const fieldWalkName = `walk_${camelToSnake(field.innerTypeName)}`,
       fieldCamelName = snakeToCamel(field.name);
-    const scopeCode = field === scopeEnterField ? enterScopeCode : '';
+    const scopeCode = field === scopeEnterField
+      ? enterScopeCode
+      : field === scopeExitField
+      ? exitScopeCode
+      : '';
 
     let tagCode = '', retagCode = '';
     if (index === 0) {
@@ -212,7 +229,7 @@ function generateWalkForStruct(type, types) {
     ) {
       traverser.enter_${typeSnakeName}(&mut *node, ctx);
       ${fieldsCodes.join('\n')}
-      ${exitScopeCode}
+      ${scopeExitField ? '' : exitScopeCode}
       traverser.exit_${typeSnakeName}(&mut *node, ctx);
     }
   `.replace(/\n\s*\n+/g, '\n');

tasks/ast_tools/src/generators/visit.rs

@@ -464,6 +464,7 @@ impl<'a> VisitBuilder<'a> {
         };
 
         let mut enter_scope_at = 0;
+        let mut exit_scope_at: Option<usize> = None;
         let mut enter_node_at = 0;
         let fields_visits: Vec<TokenStream> = struct_
             .fields
@@ -481,6 +482,7 @@ impl<'a> VisitBuilder<'a> {
                 let visit_args = markers.visit.visit_args.clone();
 
                 let have_enter_scope = markers.scope.enter_before;
+                let have_exit_scope = markers.scope.exit_before;
                 let have_enter_node = markers.visit.enter_before;
 
                 let (args_def, args) = visit_args
@@ -525,6 +527,18 @@ impl<'a> VisitBuilder<'a> {
                     };
                     enter_scope_at = ix;
                 }
+                if have_exit_scope {
+                    assert!(
+                        exit_scope_at.is_none(),
+                        "Scopes cannot be exited more than once. Remove the extra `#[scope(exit_before)]` attribute(s)."
+                    );
+                    let scope_exit = &scope_events.1;
+                    result = quote! {
+                        #scope_exit
+                        #result
+                    };
+                    exit_scope_at = Some(ix);
+                }
 
                 #[expect(unreachable_code)]
                 if have_enter_node {
@@ -563,17 +577,25 @@ impl<'a> VisitBuilder<'a> {
             },
         };
 
-        let with_scope_events = |body: TokenStream| match (scope_events, enter_scope_at) {
-            ((enter, leave), 0) => quote! {
-                #enter
-                #body
-                #leave
-            },
-            ((_, leave), _) => quote! {
-                #body
-                #leave
-            },
-        };
+        let with_scope_events =
+            |body: TokenStream| match (scope_events, enter_scope_at, exit_scope_at) {
+                ((enter, leave), 0, None) => quote! {
+                    #enter
+                    #body
+                    #leave
+                },
+                ((_, leave), _, None) => quote! {
+                    #body
+                    #leave
+                },
+                ((enter, _), 0, Some(_)) => quote! {
+                    #enter
+                    #body
+                },
+                ((_, _), _, Some(_)) => quote! {
+                    #body
+                },
+            };
 
         let body = with_node_events(with_scope_events(quote!(#(#fields_visits)*)));
 

tasks/ast_tools/src/markers.rs

@@ -64,7 +64,10 @@ pub struct VisitMarkers {
 /// A struct representing `#[scope(...)]` markers
 #[derive(Default, Debug)]
 pub struct ScopeMarkers {
+    /// `#[scope(enter_before)]`
     pub enter_before: bool,
+    /// `#[scope(exit_before)]`
+    pub exit_before: bool,
 }
 
 /// A struct representing all the helper attributes that might be used with `#[generate_derive(...)]`
@@ -204,7 +207,10 @@ where
         || Ok(ScopeMarkers::default()),
         |attr| {
             attr.parse_args_with(Ident::parse)
-                .map(|id| ScopeMarkers { enter_before: id == "enter_before" })
+                .map(|id| ScopeMarkers {
+                    enter_before: id == "enter_before",
+                    exit_before: id == "exit_before",
+                })
                 .normalize()
         },
     )