feat(visit): Add experimental traverse APIs (#9464)

swc-project · Aug 20, 2024 · 3ee8980 · 3ee8980
1 parent a212ab0
commit 3ee8980
Show file tree

Hide file tree

Showing 7 changed files with 144 additions and 7 deletions.
diff --git a/.changeset/shaggy-cats-promise.md b/.changeset/shaggy-cats-promise.md
@@ -0,0 +1,9 @@
+---
+swc_css_visit: patch
+swc_html_visit: patch
+swc_xml_visit: patch
+swc_ecma_visit: patch
+swc_core: patch
+---
+
+feat(visit): Add experimental traverse APIs
diff --git a/crates/swc_css_visit/src/generated.rs b/crates/swc_css_visit/src/generated.rs
@@ -113036,8 +113036,9 @@ pub enum NodeRef<'ast> {
     WqName(&'ast WqName),
 }
 impl<'ast> NodeRef<'ast> {
+    #[doc = r" This is not a part of semver-stable API. It is experimental and subject to change."]
     #[allow(unreachable_patterns)]
-    pub fn raw_children(&'ast self) -> Box<dyn 'ast + Iterator<Item = NodeRef<'ast>>> {
+    pub fn experimental_raw_children<'a>(&'a self) -> Box<dyn 'a + Iterator<Item = NodeRef<'ast>>> {
         match self {
             NodeRef::AbsoluteColorBase(node) => match node {
                 AbsoluteColorBase::HexColor(v0) => {
@@ -114614,5 +114615,23 @@ impl<'ast> NodeRef<'ast> {
         }
     }
 }
+impl<'ast> NodeRef<'ast> {
+    #[doc = r" Visit all nodes in self in preorder."]
+    #[doc = r""]
+    #[doc = r" This is not a part of semver-stable API. It is"]
+    #[doc = r" experimental and subject to change."]
+    pub fn experimental_traverse(&'ast self) -> Box<dyn 'ast + Iterator<Item = NodeRef<'ast>>> {
+        let mut queue = std::collections::VecDeque::<NodeRef<'ast>>::new();
+        queue.push_back(*self);
+        Box::new(std::iter::from_fn(move || {
+            let node: NodeRef<'ast> = queue.pop_front()?;
+            {
+                let children = node.experimental_raw_children();
+                queue.extend(children);
+            }
+            Some(node)
+        }))
+    }
+}
 #[cfg(any(docsrs, feature = "path"))]
 pub use self::fields::{AstParentKind, AstParentNodeRef};
diff --git a/crates/swc_ecma_visit/src/generated.rs b/crates/swc_ecma_visit/src/generated.rs
@@ -143163,8 +143163,9 @@ pub enum NodeRef<'ast> {
     YieldExpr(&'ast YieldExpr),
 }
 impl<'ast> NodeRef<'ast> {
+    #[doc = r" This is not a part of semver-stable API. It is experimental and subject to change."]
     #[allow(unreachable_patterns)]
-    pub fn raw_children(&'ast self) -> Box<dyn 'ast + Iterator<Item = NodeRef<'ast>>> {
+    pub fn experimental_raw_children<'a>(&'a self) -> Box<dyn 'a + Iterator<Item = NodeRef<'ast>>> {
         match self {
             NodeRef::Accessibility(node) => match node {
                 _ => Box::new(::std::iter::empty::<NodeRef<'ast>>()),
@@ -145509,5 +145510,23 @@ impl<'ast> NodeRef<'ast> {
         }
     }
 }
+impl<'ast> NodeRef<'ast> {
+    #[doc = r" Visit all nodes in self in preorder."]
+    #[doc = r""]
+    #[doc = r" This is not a part of semver-stable API. It is"]
+    #[doc = r" experimental and subject to change."]
+    pub fn experimental_traverse(&'ast self) -> Box<dyn 'ast + Iterator<Item = NodeRef<'ast>>> {
+        let mut queue = std::collections::VecDeque::<NodeRef<'ast>>::new();
+        queue.push_back(*self);
+        Box::new(std::iter::from_fn(move || {
+            let node: NodeRef<'ast> = queue.pop_front()?;
+            {
+                let children = node.experimental_raw_children();
+                queue.extend(children);
+            }
+            Some(node)
+        }))
+    }
+}
 #[cfg(any(docsrs, feature = "path"))]
 pub use self::fields::{AstParentKind, AstParentNodeRef};
diff --git a/crates/swc_ecma_visit/tests/main.rs b/crates/swc_ecma_visit/tests/main.rs
@@ -1,6 +1,6 @@
 use swc_common::{chain, DUMMY_SP};
-use swc_ecma_ast::{Module, Program};
-use swc_ecma_visit::{Visit, VisitWith};
+use swc_ecma_ast::*;
+use swc_ecma_visit::{NodeRef, Visit, VisitWith};
 
 #[test]
 fn should_visit_program() {
@@ -29,3 +29,30 @@ fn should_visit_program() {
 
     assert_eq!(counter, 1);
 }
+
+#[test]
+fn traverse_lookup() {
+    let node = Expr::Call(CallExpr {
+        span: DUMMY_SP,
+        callee: Callee::Expr(
+            AwaitExpr {
+                span: DUMMY_SP,
+                arg: Ident::new_no_ctxt("foo".into(), DUMMY_SP).into(),
+            }
+            .into(),
+        ),
+        args: Vec::new(),
+        ..Default::default()
+    });
+
+    let node_ref = NodeRef::from(&node);
+    let iter = node_ref.experimental_traverse();
+
+    let mut has_await = false;
+
+    for node in iter {
+        has_await |= matches!(node, NodeRef::AwaitExpr(..));
+    }
+
+    assert!(has_await);
+}
diff --git a/crates/swc_html_visit/src/generated.rs b/crates/swc_html_visit/src/generated.rs
@@ -11289,8 +11289,9 @@ pub enum NodeRef<'ast> {
     TokenAndSpan(&'ast TokenAndSpan),
 }
 impl<'ast> NodeRef<'ast> {
+    #[doc = r" This is not a part of semver-stable API. It is experimental and subject to change."]
     #[allow(unreachable_patterns)]
-    pub fn raw_children(&'ast self) -> Box<dyn 'ast + Iterator<Item = NodeRef<'ast>>> {
+    pub fn experimental_raw_children<'a>(&'a self) -> Box<dyn 'a + Iterator<Item = NodeRef<'ast>>> {
         match self {
             NodeRef::Attribute(node) => {
                 let iterator = ::std::iter::empty::<NodeRef<'ast>>().chain(
@@ -11381,5 +11382,23 @@ impl<'ast> NodeRef<'ast> {
         }
     }
 }
+impl<'ast> NodeRef<'ast> {
+    #[doc = r" Visit all nodes in self in preorder."]
+    #[doc = r""]
+    #[doc = r" This is not a part of semver-stable API. It is"]
+    #[doc = r" experimental and subject to change."]
+    pub fn experimental_traverse(&'ast self) -> Box<dyn 'ast + Iterator<Item = NodeRef<'ast>>> {
+        let mut queue = std::collections::VecDeque::<NodeRef<'ast>>::new();
+        queue.push_back(*self);
+        Box::new(std::iter::from_fn(move || {
+            let node: NodeRef<'ast> = queue.pop_front()?;
+            {
+                let children = node.experimental_raw_children();
+                queue.extend(children);
+            }
+            Some(node)
+        }))
+    }
+}
 #[cfg(any(docsrs, feature = "path"))]
 pub use self::fields::{AstParentKind, AstParentNodeRef};
diff --git a/crates/swc_xml_visit/src/generated.rs b/crates/swc_xml_visit/src/generated.rs
@@ -10857,8 +10857,9 @@ pub enum NodeRef<'ast> {
     TokenAndSpan(&'ast TokenAndSpan),
 }
 impl<'ast> NodeRef<'ast> {
+    #[doc = r" This is not a part of semver-stable API. It is experimental and subject to change."]
     #[allow(unreachable_patterns)]
-    pub fn raw_children(&'ast self) -> Box<dyn 'ast + Iterator<Item = NodeRef<'ast>>> {
+    pub fn experimental_raw_children<'a>(&'a self) -> Box<dyn 'a + Iterator<Item = NodeRef<'ast>>> {
         match self {
             NodeRef::Attribute(node) => {
                 let iterator = ::std::iter::empty::<NodeRef<'ast>>().chain(
@@ -10942,5 +10943,23 @@ impl<'ast> NodeRef<'ast> {
         }
     }
 }
+impl<'ast> NodeRef<'ast> {
+    #[doc = r" Visit all nodes in self in preorder."]
+    #[doc = r""]
+    #[doc = r" This is not a part of semver-stable API. It is"]
+    #[doc = r" experimental and subject to change."]
+    pub fn experimental_traverse(&'ast self) -> Box<dyn 'ast + Iterator<Item = NodeRef<'ast>>> {
+        let mut queue = std::collections::VecDeque::<NodeRef<'ast>>::new();
+        queue.push_back(*self);
+        Box::new(std::iter::from_fn(move || {
+            let node: NodeRef<'ast> = queue.pop_front()?;
+            {
+                let children = node.experimental_raw_children();
+                queue.extend(children);
+            }
+            Some(node)
+        }))
+    }
+}
 #[cfg(any(docsrs, feature = "path"))]
 pub use self::fields::{AstParentKind, AstParentNodeRef};
diff --git a/tools/generate-code/src/generators/visitor.rs b/tools/generate-code/src/generators/visitor.rs
@@ -1599,14 +1599,39 @@ fn define_fields(crate_name: &Ident, node_types: &[&Item]) -> Vec<Item> {
             ));
             items.push(parse_quote!(
                 impl<'ast> NodeRef<'ast> {
+                    /// This is not a part of semver-stable API. It is experimental and subject to change.
                     #[allow(unreachable_patterns)]
-                    pub fn raw_children(&'ast self) -> Box<dyn 'ast + Iterator<Item = NodeRef<'ast>>> {
+                    pub fn experimental_raw_children<'a>(&'a self) -> Box<dyn 'a + Iterator<Item = NodeRef<'ast>>> {
                         match self {
                             #(#node_ref_iter_next_arms)*
                         }
                     }
                 }
             ));
+
+            items.push(parse_quote!(
+                impl<'ast> NodeRef<'ast> {
+                    /// Visit all nodes in self in preorder.
+                    ///
+                    /// This is not a part of semver-stable API. It is
+                    /// experimental and subject to change.
+                    pub fn experimental_traverse(
+                        &'ast self,
+                    ) -> Box<dyn 'ast + Iterator<Item = NodeRef<'ast>>> {
+                        let mut queue = std::collections::VecDeque::<NodeRef<'ast>>::new();
+                        queue.push_back(*self);
+
+                        Box::new(std::iter::from_fn(move || {
+                            let node: NodeRef<'ast> = queue.pop_front()?;
+                            {
+                                let children = node.experimental_raw_children();
+                                queue.extend(children);
+                            }
+                            Some(node)
+                        }))
+                    }
+                }
+            ));
         }
 
         items.insert(