refactor(transform): retag_stack use AncestorType (oxc-project#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 oxc-project#3170.

close: oxc-project#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;