refactor(traverse): harden soundness of Traverse and document safety invariants better (#8507)

Harden soundness of `Traverse` by:

1. Not exposing `walk_*` methods outside of `walk.rs`.
2. Adding more debug assertions.
3. Adding `#[must_use]` to `TraverseAncestry::push_stack`, so we'll get a lint error if codegen-ed code pushes to `TraverseAncestry`'s stack and doesn't pop it off again.

Document the safety invariants better.

crates/oxc_traverse/scripts/lib/walk.mjs

@@ -42,9 +42,24 @@ export default function generateWalkFunctionsCode(types) {
 
     use crate::{ancestor::{self, AncestorType}, Ancestor, Traverse, TraverseCtx};
 
+    /// Walk AST with \`Traverse\` impl.
+    ///
+    /// SAFETY:
+    /// * \`program\` must be a pointer to a valid \`Program\` which has lifetime \`'a\`
+    ///   (\`Program<'a>\`).
+    /// * \`ctx\` must contain a \`TraverseAncestry<'a>\` with single \`Ancestor::None\` on its stack.
+    #[inline]
+    pub(crate) unsafe fn walk_ast<'a, Tr: Traverse<'a>>(
+      traverser: &mut Tr,
+      program: *mut Program<'a>,
+      ctx: &mut TraverseCtx<'a>,
+    ) {
+      walk_program(traverser, program, ctx);
+    }
+
     ${walkMethods}
 
-    pub(crate) unsafe fn walk_statements<'a, Tr: Traverse<'a>>(
+    unsafe fn walk_statements<'a, Tr: Traverse<'a>>(
       traverser: &mut Tr,
       stmts: *mut Vec<'a, Statement<'a>>,
       ctx: &mut TraverseCtx<'a>
@@ -254,7 +269,7 @@ function generateWalkForStruct(type, types) {
 
   const typeSnakeName = camelToSnake(type.name);
   return `
-    pub(crate) unsafe fn walk_${typeSnakeName}<'a, Tr: Traverse<'a>>(
+    unsafe fn walk_${typeSnakeName}<'a, Tr: Traverse<'a>>(
       traverser: &mut Tr,
       node: *mut ${type.rawName},
       ctx: &mut TraverseCtx<'a>
@@ -319,7 +334,7 @@ function generateWalkForEnum(type, types) {
 
   const typeSnakeName = camelToSnake(type.name);
   return `
-    pub(crate) unsafe fn walk_${typeSnakeName}<'a, Tr: Traverse<'a>>(
+    unsafe fn walk_${typeSnakeName}<'a, Tr: Traverse<'a>>(
       traverser: &mut Tr,
       node: *mut ${type.rawName},
       ctx: &mut TraverseCtx<'a>

crates/oxc_traverse/src/context/ancestry.rs

@@ -147,6 +147,7 @@ impl<'a> TraverseAncestry<'a> {
     /// # SAFETY
     /// This method must not be public outside this crate, or consumer could break safety invariants.
     #[inline]
+    #[must_use] // `PopToken` must be passed to `pop_stack` to pop this entry off the stack again
     pub(crate) fn push_stack(&mut self, ancestor: Ancestor<'a, 'static>) -> PopToken {
         self.stack.push(ancestor);
 

crates/oxc_traverse/src/lib.rs

@@ -80,7 +80,7 @@ mod generated {
     pub(super) mod walk;
 }
 pub use generated::{ancestor, ancestor::Ancestor, traverse::Traverse};
-use generated::{scopes_collector, walk::walk_program};
+use generated::{scopes_collector, walk::walk_ast};
 
 mod compile_fail_tests;
 
@@ -170,8 +170,33 @@ pub fn traverse_mut_with_ctx<'a, Tr: Traverse<'a>>(
     program: &mut Program<'a>,
     ctx: &mut ReusableTraverseCtx<'a>,
 ) {
+    let program = ptr::from_mut(program);
+
     let ctx = ctx.get_mut();
-    // SAFETY: Walk functions are constructed to avoid unsoundness
-    unsafe { walk_program(traverser, ptr::from_mut(program), ctx) };
+
+    // Check that `TraverseAncestry`'s stack is in correct state
     debug_assert!(ctx.ancestors_depth() == 1);
+    debug_assert!(matches!(ctx.parent(), Ancestor::None));
+
+    // SAFETY:
+    // `program` is a valid pointer to a `Program<'a>` - it was created from a `&mut Program<'a>`.
+    //
+    // `TraverseCtx`s are always created with only `Ancestor::None` on `TraverseAncestry` stack.
+    // `TraverseCtx` provides no external interfaces for caller to mutate `TraverseAncestry`,
+    // so that cannot be changed by the caller.
+    //
+    // `walk_*` methods never alter the initial entry on `TraverseAncestry`'s stack.
+    // `walk_*` methods always follow a `push` to `TraverseAncestry`'s stack with a corresponding `pop`.
+    // So at end of traversal, `TraverseAncestry`'s stack is always back in the state it was at start
+    // of traversal (single `Ancestor::None` entry).
+    //
+    // `ctx` was contained in a `ReusableTraverseCtx<'a>`. `ReusableTraverseCtx<'a>` provides no external
+    // interfaces for caller to mutate the `TraverseCtx` it contains.
+    //
+    // Therefore `TraverseAncestry`'s stack is guaranteed to only contain single `Ancestor::None` entry.
+    unsafe { walk_ast(traverser, program, ctx) };
+
+    // Check that `TraverseAncestry`'s stack is in correct state
+    debug_assert!(ctx.ancestors_depth() == 1);
+    debug_assert!(matches!(ctx.parent(), Ancestor::None));
 }