refactor(traverse): improve safety via type system (#5277)

Guard against mistakes in `oxc_traverse` codegen by making it impossible to pop from ancestors stack more times than have pushed.

crates/oxc_traverse/scripts/lib/walk.mjs

@@ -93,7 +93,7 @@ function generateWalkForStruct(type, types) {
         let tagCode = '', retagCode = '';
         if (index === 0) {
             tagCode = `
-                ctx.push_stack(
+                let pop_token = ctx.push_stack(
                     Ancestor::${type.name}${fieldCamelName}(
                         ancestor::${type.name}Without${fieldCamelName}(node, PhantomData)
                     )
@@ -182,7 +182,7 @@ function generateWalkForStruct(type, types) {
         `;
     });
 
-    if (visitedFields.length > 0) fieldsCodes.push('ctx.pop_stack();');
+    if (visitedFields.length > 0) fieldsCodes.push('ctx.pop_stack(pop_token);');
 
     const typeSnakeName = camelToSnake(type.name);
     return `

crates/oxc_traverse/src/context/ancestry.rs

@@ -105,21 +105,27 @@ 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, 'static>) {
+    pub(crate) fn push_stack(&mut self, ancestor: Ancestor<'a, 'static>) -> PopToken {
         self.stack.push(ancestor);
+
+        // Return `PopToken` which can be used to pop this entry off again
+        PopToken(())
     }
 
     /// Pop last item off ancestry stack.
     ///
     /// # SAFETY
-    /// * Stack must have length of at least 2 (so length is minimum 1 after pop).
-    /// * Each `pop_stack` call must correspond to a `push_stack` call for same type.
-    ///
     /// This method must not be public outside this crate, or consumer could break safety invariants.
     #[inline]
-    pub(crate) unsafe fn pop_stack(&mut self) {
+    #[allow(unused_variables, clippy::needless_pass_by_value)]
+    pub(crate) fn pop_stack(&mut self, token: PopToken) {
         debug_assert!(self.stack.len() >= 2);
-        self.stack.pop().unwrap_unchecked();
+        // SAFETY: `PopToken`s are only created in `push_stack`, so the fact that caller provides one
+        // guarantees that a push has happened. This method consumes the token which guarantees another
+        // pop hasn't occurred already corresponding to that push.
+        // Therefore the stack cannot by empty.
+        // The stack starts with 1 entry, so also it cannot be left empty after this pop.
+        unsafe { self.stack.pop().unwrap_unchecked() };
     }
 
     /// Retag last item on ancestry stack.
@@ -149,3 +155,10 @@ impl<'a> TraverseAncestry<'a> {
         *(self.stack.last_mut().unwrap_unchecked() as *mut _ as *mut AncestorType) = ty;
     }
 }
+
+/// Zero sized token which allows popping from stack. Used to ensure push and pop always correspond.
+/// Inner field is private to this module so can only be created by methods in this file.
+/// It is not `Clone` or `Copy`, so no way to obtain one except in this file.
+/// Only method which generates a `PopToken` is `push_stack`, and `pop_stack` consumes one,
+/// which guarantees you can't have more pops than pushes.
+pub(crate) struct PopToken(());

crates/oxc_traverse/src/context/mod.rs

@@ -14,6 +14,7 @@ use oxc_syntax::{
 use crate::ancestor::{Ancestor, AncestorType};
 mod ancestry;
 mod ast_operations;
+use ancestry::PopToken;
 pub use ancestry::TraverseAncestry;
 mod scoping;
 pub use scoping::TraverseScoping;
@@ -440,8 +441,8 @@ 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, 'static>) {
+    pub(crate) fn push_stack(&mut self, ancestor: Ancestor<'a, 'static>) -> PopToken {
-        self.ancestry.push_stack(ancestor);
+        self.ancestry.push_stack(ancestor)
     }
 
     /// Shortcut for `self.ancestry.pop_stack`, to make `walk_*` methods less verbose.
@@ -450,8 +451,8 @@ 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) {
+    pub(crate) unsafe fn pop_stack(&mut self, token: PopToken) {
-        self.ancestry.pop_stack();
+        self.ancestry.pop_stack(token);
     }
 
     /// Shortcut for `self.ancestry.retag_stack`, to make `walk_*` methods less verbose.