refactor(transform): retag_stack use AncestorType (#3173)

Make the code for `retag_stack` in `walk_*` functions more
comprehensible, by replacing hard-coded enum discriminants as integers
with an `AncestorType` type.

Alternative solution to the problem raised in #3170.

close: #3170

crates/oxc_traverse/scripts/lib/ancestor.mjs

@@ -2,7 +2,8 @@ import {camelToSnake, snakeToCamel} from './utils.mjs';
 
 export default function generateAncestorsCode(types) {
     const variantNamesForEnums = Object.create(null);
-    let enumVariants = '',
+    let ancestorTypeEnumVariants = '',
+        ancestorEnumVariants = '',
         isFunctions = '',
         ancestorTypes = '',
         discriminant = 1;
@@ -62,8 +63,8 @@ export default function generateAncestorsCode(types) {
             const variantName = `${type.name}${fieldNameCamel}`;
             variantNames.push(variantName);
 
-            enumVariants += `${variantName}(${structName}) = ${discriminant},\n`;
-            field.ancestorDiscriminant = discriminant;
+            ancestorTypeEnumVariants += `${variantName} = ${discriminant},\n`;
+            ancestorEnumVariants += `${variantName}(${structName}) = AncestorType::${variantName} as u16,\n`;
             discriminant++;
 
             if (fieldType.kind === 'enum') {
@@ -96,8 +97,6 @@ export default function generateAncestorsCode(types) {
         `;
     }
 
-    const discriminantType = discriminant <= 256 ? 'u8' : 'u16';
-
     return `
         #![allow(
             unsafe_code,
@@ -117,20 +116,34 @@ export default function generateAncestorsCode(types) {
             AssignmentOperator, BinaryOperator, LogicalOperator, UnaryOperator, UpdateOperator,
         };
 
-        pub(crate) type AncestorDiscriminant = ${discriminantType};
+        /// Type of [\`Ancestor\`].
+        /// Used in [\`crate::TraverseCtx::retag_stack\`].
+        #[repr(u16)]
+        #[derive(Clone, Copy)]
+        #[allow(dead_code)]
+        pub(crate) enum AncestorType {
+            None = 0,
+            ${ancestorTypeEnumVariants}
+        }
 
         /// Ancestor type used in AST traversal.
         ///
         /// Encodes both the type of the parent, and child's location in the parent.
         /// i.e. variants for \`BinaryExpressionLeft\` and \`BinaryExpressionRight\`, not just \`BinaryExpression\`.
         //
-        // SAFETY: This type MUST be \`#[repr(u8)]\` or \`#[repr(u16)]\` (depending on number of variants)
-        // to maintain the safety of \`TraverseCtx::retag_stack\`.
-        #[repr(C, ${discriminantType})]
+        // SAFETY:
+        // * This type must be \`#[repr(u16)]\`.
+        // * Variant discriminants must correspond to those in \`AncestorType\`.
+        //
+        // These invariants make it possible to set the discriminant of an \`Ancestor\` without altering
+        // the "payload" pointer with:
+        // \`*(ancestor as *mut _ as *mut AncestorType) = AncestorType::Program\`.
+        // \`TraverseCtx::retag_stack\` uses this technique.
+        #[repr(C, u16)]
         #[derive(Debug)]
         pub enum Ancestor<'a> {
-            None = 0,
-            ${enumVariants}
+            None = AncestorType::None as u16,
+            ${ancestorEnumVariants}
         }
 
         impl<'a> Ancestor<'a> {

crates/oxc_traverse/scripts/lib/walk.mjs

@@ -27,7 +27,7 @@ export default function generateWalkFunctionsCode(types) {
         #[allow(clippy::wildcard_imports)]
         use oxc_ast::ast::*;
 
-        use crate::{ancestor, Ancestor, Traverse, TraverseCtx};
+        use crate::{ancestor::{self, AncestorType}, Ancestor, Traverse, TraverseCtx};
 
         ${walkMethods}
 
@@ -51,8 +51,10 @@ function generateWalkForStruct(type, types) {
     const fieldsCodes = visitedFields.map((field, index) => {
         const fieldWalkName = `walk_${camelToSnake(field.innerTypeName)}`;
 
-        const retagCode = index === 0 ? '' : `ctx.retag_stack(${field.ancestorDiscriminant});`,
-            fieldCode = `(node as *mut u8).add(ancestor::${field.offsetVarName}) as *mut ${field.typeName}`;
+        const retagCode = index === 0
+            ? ''
+            : `ctx.retag_stack(AncestorType::${type.name}${snakeToCamel(field.name)});`;
+        const fieldCode = `(node as *mut u8).add(ancestor::${field.offsetVarName}) as *mut ${field.typeName}`;
 
         if (field.wrappers[0] === 'Option') {
             let walkCode;

crates/oxc_traverse/src/context.rs

@@ -1,7 +1,7 @@
 use oxc_allocator::{Allocator, Box};
 use oxc_ast::AstBuilder;
 
-use crate::ancestor::{Ancestor, AncestorDiscriminant};
+use crate::ancestor::{Ancestor, AncestorType};
 
 const INITIAL_STACK_CAPACITY: usize = 64;
 
@@ -119,25 +119,21 @@ impl<'a> TraverseCtx<'a> {
     /// of pointer to the ancestor node.
     ///
     /// This is purely a performance optimization. If the last item on stack already contains the
-    /// correct pointer, then `ctx.retag_stack(3)` is equivalent to:
+    /// correct pointer, then `ctx.retag_stack(AncestorType::ProgramBody)` is equivalent to:
     ///
     /// ```nocompile
     /// ctx.pop_stack();
     /// ctx.push_stack(Ancestor::ProgramBody(ProgramWithoutBody(node_ptr)));
     /// ```
     ///
-    /// (3 is the discriminant for `Ancestor::ProgramBody`)
-    ///
     /// `retag_stack` is only a single 2-byte write operation.
     ///
     /// # SAFETY
     /// * Stack must not be empty.
-    /// * `discriminant` must be valid discriminant value for `Ancestor` enum.
-    /// * Last item on stack must contain pointer to type corresponding to provided discriminant.
+    /// * Last item on stack must contain pointer to type corresponding to provided `AncestorType`.
     #[inline]
     #[allow(unsafe_code, clippy::ptr_as_ptr, clippy::ref_as_ptr)]
-    pub(crate) unsafe fn retag_stack(&mut self, discriminant: AncestorDiscriminant) {
-        *(self.stack.last_mut().unwrap_unchecked() as *mut _ as *mut AncestorDiscriminant) =
-            discriminant;
+    pub(crate) unsafe fn retag_stack(&mut self, ty: AncestorType) {
+        *(self.stack.last_mut().unwrap_unchecked() as *mut _ as *mut AncestorType) = ty;
     }
 }