refactor(traverse): make Ancestor an owned type (#5269)

Make the `Ancestor` type used in `oxc_traverse` an owned type. Instead of `TraverseCtx::parent` returning a `&'t Ancestor<'a>`, it now returns an `Ancestor<'a, 't>`.

This allows `Ancestor` to be `Copy`.

The `'t` lifetime plays the same role in both cases - preventing any `Ancestor` from escaping from the `enter_*` / `exit_*` method in which it's obtained.

Same for the `*Without*` types which are `Ancestor` enum's "payloads". Any AST node references obtained from an `Ancestor` are also constrained by same `'t` lifetime - e.g. `&'t Statement<'a>`.

crates/oxc_traverse/scripts/lib/ancestor.mjs

@@ -33,7 +33,7 @@ export default function generateAncestorsCode(types) {
 
                 methodsCode += `
                     #[inline]
-                    pub fn ${otherField.rawName}(&self) -> &${otherField.rawTypeName} {
+                    pub fn ${otherField.rawName}(self) -> &'t ${otherField.rawTypeName} {
                         unsafe {
                             &*(
                                 (self.0 as *const u8).add(${otherField.offsetVarName})
@@ -45,17 +45,18 @@ export default function generateAncestorsCode(types) {
             }
 
             const fieldNameCamel = snakeToCamel(field.name),
-                lifetime = type.rawName.slice(type.name.length),
+                lifetimes = type.rawName.length > type.name.length ? `<'a, 't>` : "<'t>",
-                structName = `${type.name}Without${fieldNameCamel}${lifetime}`;
+                structName = `${type.name}Without${fieldNameCamel}${lifetimes}`;
 
             thisAncestorTypes += `
                 #[repr(transparent)]
-                #[derive(Debug)]
+                #[derive(Clone, Copy, Debug)]
                 pub struct ${structName}(
-                    pub(crate) *const ${type.rawName}
+                    pub(crate) *const ${type.rawName},
+                    pub(crate) PhantomData<&'t ()>,
                 );
 
-                impl${lifetime} ${structName} {
+                impl${lifetimes} ${structName} {
                     ${methodsCode}
                 }
             `;
@@ -81,7 +82,7 @@ export default function generateAncestorsCode(types) {
 
             isFunctions += `
                 #[inline]
-                pub fn is_${typeSnakeName}(&self) -> bool {
+                pub fn is_${typeSnakeName}(self) -> bool {
                     matches!(self, ${variantNames.map(name => `Self::${name}(_)`).join(' | ')})
                 }
             `;
@@ -91,7 +92,7 @@ export default function generateAncestorsCode(types) {
     for (const [typeName, variantNames] of Object.entries(variantNamesForEnums)) {
         isFunctions += `
             #[inline]
-            pub fn is_via_${camelToSnake(typeName)}(&self) -> bool {
+            pub fn is_via_${camelToSnake(typeName)}(self) -> bool {
                 matches!(self, ${variantNames.map(name => `Self::${name}(_)`).join(' | ')})
             }
         `;
@@ -106,7 +107,7 @@ export default function generateAncestorsCode(types) {
             clippy::cast_ptr_alignment
         )]
 
-        use std::cell::Cell;
+        use std::{cell::Cell, marker::PhantomData};
 
         use memoffset::offset_of;
 
@@ -129,6 +130,11 @@ export default function generateAncestorsCode(types) {
         ///
         /// Encodes both the type of the parent, and child's location in the parent.
         /// i.e. variants for \`BinaryExpressionLeft\` and \`BinaryExpressionRight\`, not just \`BinaryExpression\`.
+        ///
+        /// \`'a\` is lifetime of AST nodes.
+        /// \`'t\` is lifetime of the \`Ancestor\` (which inherits lifetime from \`&'t TraverseCtx'\`).
+        /// i.e. \`Ancestor\`s can only exist within the body of \`enter_*\` and \`exit_*\` methods
+        /// and cannot "escape" from them.
         //
         // SAFETY:
         // * This type must be \`#[repr(u16)]\`.
@@ -139,13 +145,13 @@ export default function generateAncestorsCode(types) {
         // \`*(ancestor as *mut _ as *mut AncestorType) = AncestorType::Program\`.
         // \`TraverseCtx::retag_stack\` uses this technique.
         #[repr(C, u16)]
-        #[derive(Debug)]
+        #[derive(Clone, Copy, Debug)]
-        pub enum Ancestor<'a> {
+        pub enum Ancestor<'a, 't> {
             None = AncestorType::None as u16,
             ${ancestorEnumVariants}
         }
 
-        impl<'a> Ancestor<'a> {
+        impl<'a, 't> Ancestor<'a, 't> {
             ${isFunctions}
         }
 

crates/oxc_traverse/scripts/lib/walk.mjs

@@ -23,7 +23,7 @@ export default function generateWalkFunctionsCode(types) {
             clippy::cast_ptr_alignment
         )]
 
-        use std::cell::Cell;
+        use std::{cell::Cell, marker::PhantomData};
 
         use oxc_allocator::Vec;
         #[allow(clippy::wildcard_imports)]
@@ -95,7 +95,7 @@ function generateWalkForStruct(type, types) {
             tagCode = `
                 ctx.push_stack(
                     Ancestor::${type.name}${fieldCamelName}(
-                        ancestor::${type.name}Without${fieldCamelName}(node)
+                        ancestor::${type.name}Without${fieldCamelName}(node, PhantomData)
                     )
                 );
             `;

crates/oxc_traverse/src/compile_fail_tests.rs

@@ -14,7 +14,7 @@ use oxc_ast::ast::IdentifierReference;
 use oxc_traverse::{Ancestor, Traverse, TraverseCtx};
 
 struct Trans<'a, 'b> {
-    ancestor: Option<&'b Ancestor<'a>>,
+    ancestor: Option<Ancestor<'a, 'b>>,
 }
 
 impl<'a, 'b> Traverse<'a> for Trans<'a, 'b> {
@@ -36,7 +36,7 @@ use oxc_ast::ast::IdentifierReference;
 use oxc_traverse::{ancestor::ProgramWithoutDirectives, Ancestor, Traverse, TraverseCtx};
 
 struct Trans<'a, 'b> {
-    program: Option<&'b ProgramWithoutDirectives<'a>>,
+    program: Option<ProgramWithoutDirectives<'a, 'b>>,
 }
 
 impl<'a, 'b> Traverse<'a> for Trans<'a, 'b> {
@@ -86,7 +86,7 @@ use oxc_ast::ast::IdentifierReference;
 use oxc_traverse::{Ancestor, Traverse, TraverseCtx};
 
 struct Trans<'a, 'b> {
-    ancestor: Option<&'b Ancestor<'a>>,
+    ancestor: Option<Ancestor<'a, 'b>>,
 }
 
 impl<'a, 'b> Traverse<'a> for Trans<'a, 'b> {
@@ -108,7 +108,7 @@ use oxc_ast::ast::IdentifierReference;
 use oxc_traverse::{ancestor::ProgramWithoutDirectives, Ancestor, Traverse, TraverseCtx};
 
 struct Trans<'a, 'b> {
-    program: Option<&'b ProgramWithoutDirectives<'a>>,
+    program: Option<ProgramWithoutDirectives<'a, 'b>>,
 }
 
 impl<'a, 'b> Traverse<'a> for Trans<'a, 'b> {

crates/oxc_traverse/src/context/ancestry.rs

@@ -1,3 +1,5 @@
+use std::mem::transmute;
+
 use crate::ancestor::{Ancestor, AncestorType};
 
 const INITIAL_STACK_CAPACITY: usize = 64; // 64 entries = 1 KiB
@@ -8,6 +10,17 @@ const INITIAL_STACK_CAPACITY: usize = 64; // 64 entries = 1 KiB
 ///
 /// `walk_*` methods push/pop `Ancestor`s to `stack` when entering/exiting nodes.
 ///
+/// `Ancestor<'a, 't>` is an owned type.
+/// * `'a` is lifetime of AST nodes.
+/// * `'t` is lifetime of the `Ancestor` (derived from `&'t TraverseAncestry`).
+///
+/// `'t` is constrained in `parent`, `ancestor` and `ancestors` methods to only live as long as
+/// the `&'t TraverseAncestry` passed to the method.
+/// i.e. `Ancestor`s can only live as long as `enter_*` or `exit_*` method in which they're obtained,
+/// and cannot "escape" those methods.
+/// This is required for soundness. If an `Ancestor` could be retained longer, the references that
+/// can be got from it could alias a `&mut` reference to the same AST node.
+///
 /// # SAFETY
 /// This type MUST NOT be mutable by consumer.
 ///
@@ -24,18 +37,21 @@ const INITIAL_STACK_CAPACITY: usize = 64; // 64 entries = 1 KiB
 ///    b. cannot obtain an owned `TraverseAncestry` from a `&TraverseAncestry`
 ///       - `TraverseAncestry` is not `Clone`.
 pub struct TraverseAncestry<'a> {
-    stack: Vec<Ancestor<'a>>,
+    stack: Vec<Ancestor<'a, 'static>>,
 }
 
 // Public methods
 impl<'a> TraverseAncestry<'a> {
     /// Get parent of current node.
     #[inline]
-
+    pub fn parent<'t>(&'t self) -> Ancestor<'a, 't> {
-    pub fn parent(&self) -> &Ancestor<'a> {
         // SAFETY: Stack contains 1 entry initially. Entries are pushed as traverse down the AST,
         // and popped as go back up. So even when visiting `Program`, the initial entry is in the stack.
-        unsafe { self.stack.last().unwrap_unchecked() }
+        let ancestor = unsafe { *self.stack.last().unwrap_unchecked() };
+        // Shrink `Ancestor`'s `'t` lifetime to lifetime of `&'t self`.
+        // SAFETY: The `Ancestor` is guaranteed valid for `'t`. It is not possible to obtain
+        // a `&mut` ref to any AST node which this `Ancestor` gives access to during `'t`.
+        unsafe { transmute::<Ancestor<'a, '_>, Ancestor<'a, 't>>(ancestor) }
     }
 
     /// Get ancestor of current node.
@@ -43,13 +59,23 @@ impl<'a> TraverseAncestry<'a> {
     /// `level` is number of levels above.
     /// `ancestor(1).unwrap()` is equivalent to `parent()`.
     #[inline]
-    pub fn ancestor(&self, level: usize) -> Option<&Ancestor<'a>> {
+    pub fn ancestor<'t>(&'t self, level: usize) -> Option<Ancestor<'a, 't>> {
-        self.stack.get(self.stack.len() - level)
+        self.stack.get(self.stack.len() - level).map(|&ancestor| {
+            // Shrink `Ancestor`'s `'t` lifetime to lifetime of `&'t self`.
+            // SAFETY: The `Ancestor` is guaranteed valid for `'t`. It is not possible to obtain
+            // a `&mut` ref to any AST node which this `Ancestor` gives access to during `'t`.
+            unsafe { transmute::<Ancestor<'a, '_>, Ancestor<'a, 't>>(ancestor) }
+        })
     }
 
     /// Get iterator over ancestors, starting with closest ancestor
-    pub fn ancestors<'b>(&'b self) -> impl Iterator<Item = &'b Ancestor<'a>> {
+    pub fn ancestors<'t>(&'t self) -> impl Iterator<Item = Ancestor<'a, 't>> {
-        self.stack.iter().rev()
+        self.stack.iter().rev().map(|&ancestor| {
+            // Shrink `Ancestor`'s `'t` lifetime to lifetime of `&'t self`.
+            // SAFETY: The `Ancestor` is guaranteed valid for `'t`. It is not possible to obtain
+            // a `&mut` ref to any AST node which this `Ancestor` gives access to during `'t`.
+            unsafe { transmute::<Ancestor<'a, '_>, Ancestor<'a, 't>>(ancestor) }
+        })
     }
 
     /// Get depth in the AST.
@@ -78,7 +104,7 @@ impl<'a> TraverseAncestry<'a> {
     /// # 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>) {
+    pub(crate) fn push_stack(&mut self, ancestor: Ancestor<'a, 'static>) {
         self.stack.push(ancestor);
     }
 
@@ -90,7 +116,6 @@ impl<'a> TraverseAncestry<'a> {
     ///
     /// This method must not be public outside this crate, or consumer could break safety invariants.
     #[inline]
-
     pub(crate) unsafe fn pop_stack(&mut self) {
         self.stack.pop().unwrap_unchecked();
     }

crates/oxc_traverse/src/context/mod.rs

@@ -133,8 +133,7 @@ impl<'a> TraverseCtx<'a> {
     ///
     /// Shortcut for `ctx.ancestry.parent`.
     #[inline]
-
+    pub fn parent<'t>(&'t self) -> Ancestor<'a, 't> {
-    pub fn parent(&self) -> &Ancestor<'a> {
         self.ancestry.parent()
     }
 
@@ -145,7 +144,7 @@ impl<'a> TraverseCtx<'a> {
     ///
     /// Shortcut for `ctx.ancestry.ancestor`.
     #[inline]
-    pub fn ancestor(&self, level: usize) -> Option<&Ancestor<'a>> {
+    pub fn ancestor<'t>(&'t self, level: usize) -> Option<Ancestor<'a, 't>> {
         self.ancestry.ancestor(level)
     }
 
@@ -153,7 +152,7 @@ impl<'a> TraverseCtx<'a> {
     ///
     /// Shortcut for `ctx.ancestry.ancestors`.
     #[inline]
-    pub fn ancestors<'b>(&'b self) -> impl Iterator<Item = &'b Ancestor<'a>> {
+    pub fn ancestors<'t>(&'t self) -> impl Iterator<Item = Ancestor<'a, 't>> {
         self.ancestry.ancestors()
     }
 
@@ -441,7 +440,7 @@ impl<'a> TraverseCtx<'a> {
     /// # 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>) {
+    pub(crate) fn push_stack(&mut self, ancestor: Ancestor<'a, 'static>) {
         self.ancestry.push_stack(ancestor);
     }
 
@@ -451,7 +450,6 @@ impl<'a> TraverseCtx<'a> {
     /// See safety constraints of `TraverseAncestry.pop_stack`.
     /// This method must not be public outside this crate, or consumer could break safety invariants.
     #[inline]
-
     pub(crate) unsafe fn pop_stack(&mut self) {
         self.ancestry.pop_stack();
     }
@@ -462,7 +460,6 @@ impl<'a> TraverseCtx<'a> {
     /// See safety constraints of `TraverseAncestry.retag_stack`.
     /// This method must not be public outside this crate, or consumer could break safety invariants.
     #[inline]
-
     pub(crate) unsafe fn retag_stack(&mut self, ty: AncestorType) {
         self.ancestry.retag_stack(ty);
     }