refactor(traverse)!: replace find_ancestor with ancestors returning iterator (#4692)

Replace `find_ancestor` with `ancestors` method which returns an iterator. This does the same thing, but is more idiomatic.

crates/oxc_transformer/src/react/display_name.rs

@@ -1,7 +1,7 @@
 use oxc_allocator::Box;
 use oxc_ast::ast::*;
 use oxc_span::{Atom, SPAN};
-use oxc_traverse::{Ancestor, FinderRet, TraverseCtx};
+use oxc_traverse::{Ancestor, TraverseCtx};
 
 use crate::context::Ctx;
 
@@ -34,58 +34,60 @@ impl<'a> ReactDisplayName<'a> {
             return;
         };
 
-        let name = ctx.find_ancestor(|ancestor| {
+        let mut ancestors = ctx.ancestors();
+        let name = loop {
+            let Some(ancestor) = ancestors.next() else {
+                return;
+            };
+
             match ancestor {
                 // `foo = React.createClass({})`
                 Ancestor::AssignmentExpressionRight(assign_expr) => match assign_expr.left() {
                     AssignmentTarget::AssignmentTargetIdentifier(ident) => {
-                        FinderRet::Found(ident.name.clone())
+                        break ident.name.clone();
                     }
                     AssignmentTarget::StaticMemberExpression(expr) => {
-                        FinderRet::Found(expr.property.name.clone())
+                        break expr.property.name.clone();
                     }
                     // Babel does not handle computed member expressions e.g. `foo["bar"]`,
                     // so we diverge from Babel here, but that's probably an improvement
                     AssignmentTarget::ComputedMemberExpression(expr) => {
-                        match expr.static_property_name() {
+                        if let Some(name) = expr.static_property_name() {
-                            Some(name) => FinderRet::Found(name),
+                            break name;
-                            None => FinderRet::Stop,
                         }
+                        return;
                     }
-                    _ => FinderRet::Stop,
+                    _ => return,
                 },
                 // `let foo = React.createClass({})`
-                Ancestor::VariableDeclaratorInit(declarator) => match &declarator.id().kind {
+                Ancestor::VariableDeclaratorInit(declarator) => {
-                    BindingPatternKind::BindingIdentifier(ident) => {
+                    if let BindingPatternKind::BindingIdentifier(ident) = &declarator.id().kind {
-                        FinderRet::Found(ident.name.clone())
+                        break ident.name.clone();
                     }
-                    _ => FinderRet::Stop,
+                    return;
-                },
+                }
                 // `{foo: React.createClass({})}`
                 Ancestor::ObjectPropertyValue(prop) => {
                     // Babel only handles static identifiers e.g. `{foo: React.createClass({})}`,
                     // whereas we also handle e.g. `{"foo-bar": React.createClass({})}`,
                     // so we diverge from Babel here, but that's probably an improvement
                     if let Some(name) = prop.key().static_name() {
-                        FinderRet::Found(ctx.ast.atom(&name))
+                        break ctx.ast.atom(&name);
-                    } else {
-                        FinderRet::Stop
                     }
+                    return;
                 }
                 // `export default React.createClass({})`
                 // Uses the current file name as the display name.
                 Ancestor::ExportDefaultDeclarationDeclaration(_) => {
-                    FinderRet::Found(ctx.ast.atom(&self.ctx.filename))
+                    break ctx.ast.atom(&self.ctx.filename);
                 }
                 // Stop crawling up when hit a statement
-                _ if ancestor.is_via_statement() => FinderRet::Stop,
+                _ if ancestor.is_via_statement() => return,
-                _ => FinderRet::Continue,
+                _ => {}
             }
-        });
+        };
 
-        if let Some(name) = name {
+        self.add_display_name(obj_expr, name);
-            self.add_display_name(obj_expr, name);
-        }
     }
 }
 

crates/oxc_transformer/src/react/jsx_self.rs

@@ -54,11 +54,12 @@ impl<'a> ReactJsxSelf<'a> {
     }
 
     fn has_no_super_class(ctx: &TraverseCtx<'a>) -> bool {
-        ctx.find_ancestor(|ancestor| match ancestor {
+        for ancestor in ctx.ancestors() {
-            Ancestor::ClassBody(class) => FinderRet::Found(class.super_class().is_none()),
+            if let Ancestor::ClassBody(class) = ancestor {
-            _ => FinderRet::Continue,
+                return class.super_class().is_none();
-        })
+            }
-        .unwrap_or(true)
+        }
+        true
     }
 
     pub fn can_add_self_attribute(&self, ctx: &TraverseCtx<'a>) -> bool {

crates/oxc_traverse/src/compile_fail_tests.rs

@@ -79,3 +79,77 @@ impl<'a, 'b> Traverse<'a> for Trans<'a, 'b> {
 ```
 */
 const CANNOT_HOLD_ONTO_AST_NODE: () = ();
+
+/**
+```compile_fail
+use oxc_ast::ast::IdentifierReference;
+use oxc_traverse::{Ancestor, Traverse, TraverseCtx};
+
+struct Trans<'a, 'b> {
+    ancestor: Option<&'b Ancestor<'a>>,
+}
+
+impl<'a, 'b> Traverse<'a> for Trans<'a, 'b> {
+    fn enter_identifier_reference(
+        &mut self,
+        _node: &mut IdentifierReference<'a>,
+        ctx: &mut TraverseCtx<'a>,
+    ) {
+        self.ancestor = ctx.ancestors().next();
+    }
+}
+```
+*/
+const CANNOT_HOLD_ONTO_ANCESTOR_FROM_ANCESTORS_ITERATOR: () = ();
+
+/**
+```compile_fail
+use oxc_ast::ast::IdentifierReference;
+use oxc_traverse::{ancestor::ProgramWithoutDirectives, Ancestor, Traverse, TraverseCtx};
+
+struct Trans<'a, 'b> {
+    program: Option<&'b ProgramWithoutDirectives<'a>>,
+}
+
+impl<'a, 'b> Traverse<'a> for Trans<'a, 'b> {
+    fn enter_identifier_reference(
+        &mut self,
+        _node: &mut IdentifierReference<'a>,
+        ctx: &mut TraverseCtx<'a>,
+    ) {
+        let parent = ctx.ancestors().next().unwrap();
+        if let Ancestor::ProgramDirectives(program) = parent {
+            self.program = Some(program);
+        }
+    }
+}
+```
+*/
+const CANNOT_HOLD_ONTO_ANCESTOR_NODE_FROM_ANCESTORS_ITERATOR: () = ();
+
+/**
+```compile_fail
+use oxc_ast::ast::{IdentifierReference, Statement};
+use oxc_traverse::{Ancestor, Traverse, TraverseCtx};
+
+struct Trans<'a, 'b> {
+    stmt: Option<&'b Statement<'a>>,
+}
+
+impl<'a, 'b> Traverse<'a> for Trans<'a, 'b> {
+    fn enter_identifier_reference(
+        &mut self,
+        _node: &mut IdentifierReference<'a>,
+        ctx: &mut TraverseCtx<'a>,
+    ) {
+        let parent = ctx.ancestors().next().unwrap();
+        if let Ancestor::ProgramDirectives(program) = parent {
+            let body = program.body();
+            let stmt = &body[0];
+            self.stmt = Some(stmt);
+        }
+    }
+}
+```
+*/
+const CANNOT_HOLD_ONTO_AST_NODE_FROM_ANCESTORS_ITERATOR: () = ();

crates/oxc_traverse/src/context/ancestry.rs

@@ -1,4 +1,3 @@
-use super::FinderRet;
 use crate::ancestor::{Ancestor, AncestorType};
 
 const INITIAL_STACK_CAPACITY: usize = 64; // 64 entries = 1 KiB
@@ -48,49 +47,9 @@ impl<'a> TraverseAncestry<'a> {
         self.stack.get(self.stack.len() - level)
     }
 
-    /// Walk up trail of ancestors to find a node.
+    /// Get iterator over ancestors, starting with closest ancestor
-    ///
+    pub fn ancestors<'b>(&'b self) -> impl Iterator<Item = &'b Ancestor<'a>> {
-    /// `finder` should return:
+        self.stack.iter().rev()
-    /// * `FinderRet::Found(value)` to stop walking and return `Some(value)`.
-    /// * `FinderRet::Stop` to stop walking and return `None`.
-    /// * `FinderRet::Continue` to continue walking up.
-    ///
-    /// # Example
-    ///
-    /// ```
-    /// use oxc_ast::ast::ThisExpression;
-    /// use oxc_traverse::{Ancestor, FinderRet, Traverse, TraverseCtx};
-    ///
-    /// struct MyTraverse;
-    /// impl<'a> Traverse<'a> for MyTraverse {
-    ///     fn enter_this_expression(&mut self, this_expr: &mut ThisExpression, ctx: &mut TraverseCtx<'a>) {
-    ///         // Get name of function where `this` is bound.
-    ///         // NB: This example doesn't handle `this` in class fields or static blocks.
-    ///         let fn_id = ctx.ancestry.find_ancestor(|ancestor| {
-    ///             match ancestor {
-    ///                 Ancestor::FunctionBody(func) => FinderRet::Found(func.id()),
-    ///                 Ancestor::FunctionParams(func) => FinderRet::Found(func.id()),
-    ///                 _ => FinderRet::Continue
-    ///             }
-    ///         });
-    ///     }
-    /// }
-    /// ```
-    //
-    // `'c` lifetime on `&'c self` and `&'c Ancestor` passed into the closure
-    // allows an `Ancestor` or AST node to be returned from the closure.
-    pub fn find_ancestor<'c, F, O>(&'c self, finder: F) -> Option<O>
-    where
-        F: Fn(&'c Ancestor<'a>) -> FinderRet<O>,
-    {
-        for ancestor in self.stack.iter().rev() {
-            match finder(ancestor) {
-                FinderRet::Found(res) => return Some(res),
-                FinderRet::Stop => return None,
-                FinderRet::Continue => {}
-            }
-        }
-        None
     }
 
     /// Get depth in the AST.

crates/oxc_traverse/src/context/mod.rs

@@ -23,7 +23,7 @@ pub use scoping::TraverseScoping;
 /// Passed to all AST visitor functions.
 ///
 /// Provides ability to:
-/// * Query parent/ancestor of current node via [`parent`], [`ancestor`], [`find_ancestor`].
+/// * Query parent/ancestor of current node via [`parent`], [`ancestor`], [`ancestors`].
 /// * Get scopes tree and symbols table via [`scopes`], [`symbols`], [`scopes_mut`], [`symbols_mut`],
 ///   [`find_scope`], [`find_scope_by_flags`].
 /// * Create AST nodes via AST builder [`ast`].
@@ -95,7 +95,7 @@ pub use scoping::TraverseScoping;
 ///
 /// [`parent`]: `TraverseCtx::parent`
 /// [`ancestor`]: `TraverseCtx::ancestor`
-/// [`find_ancestor`]: `TraverseCtx::find_ancestor`
+/// [`ancestors`]: `TraverseCtx::ancestors`
 /// [`scopes`]: `TraverseCtx::scopes`
 /// [`symbols`]: `TraverseCtx::symbols`
 /// [`scopes_mut`]: `TraverseCtx::scopes_mut`
@@ -110,7 +110,7 @@ pub struct TraverseCtx<'a> {
     pub ast: AstBuilder<'a>,
 }
 
-/// Return value of closure when using [`TraverseCtx::find_ancestor`] or [`TraverseCtx::find_scope`].
+/// Return value of closure when using [`TraverseCtx::find_scope`].
 pub enum FinderRet<T> {
     Found(T),
     Stop,
@@ -157,41 +157,12 @@ impl<'a> TraverseCtx<'a> {
         self.ancestry.ancestor(level)
     }
 
-    /// Walk up trail of ancestors to find a node.
+    /// Get iterator over ancestors, starting with closest ancestor.
     ///
-    /// `finder` should return:
+    /// Shortcut for `ctx.ancestry.ancestors`.
-    /// * `FinderRet::Found(value)` to stop walking and return `Some(value)`.
+    #[inline]
-    /// * `FinderRet::Stop` to stop walking and return `None`.
+    pub fn ancestors<'b>(&'b self) -> impl Iterator<Item = &'b Ancestor<'a>> {
-    /// * `FinderRet::Continue` to continue walking up.
+        self.ancestry.ancestors()
-    ///
-    /// # Example
-    ///
-    /// ```
-    /// use oxc_ast::ast::ThisExpression;
-    /// use oxc_traverse::{Ancestor, FinderRet, Traverse, TraverseCtx};
-    ///
-    /// struct MyTraverse;
-    /// impl<'a> Traverse<'a> for MyTraverse {
-    ///     fn enter_this_expression(&mut self, this_expr: &mut ThisExpression, ctx: &mut TraverseCtx<'a>) {
-    ///         // Get name of function where `this` is bound.
-    ///         // NB: This example doesn't handle `this` in class fields or static blocks.
-    ///         let fn_id = ctx.find_ancestor(|ancestor| {
-    ///             match ancestor {
-    ///                 Ancestor::FunctionBody(func) => FinderRet::Found(func.id()),
-    ///                 Ancestor::FunctionParams(func) => FinderRet::Found(func.id()),
-    ///                 _ => FinderRet::Continue
-    ///             }
-    ///         });
-    ///     }
-    /// }
-    /// ```
-    ///
-    /// Shortcut for `self.ancestry.find_ancestor`.
-    pub fn find_ancestor<'c, F, O>(&'c self, finder: F) -> Option<O>
-    where
-        F: Fn(&'c Ancestor<'a>) -> FinderRet<O>,
-    {
-        self.ancestry.find_ancestor(finder)
     }
 
     /// Get depth in the AST.