From d50689eab0e82a68c84fcacdd8ce5c9009a65100 Mon Sep 17 00:00:00 2001
From: Fy <1114550440@qq.com>
Date: Wed, 22 Mar 2023 19:48:13 +0800
Subject: [PATCH] fix(es/visit): Fix handling of `Program` of `AndThen` (#7120)

---
 crates/swc_ecma_visit/src/lib.rs    | 40 +++++++++++++++++++++++++++++
 crates/swc_ecma_visit/tests/main.rs | 31 ++++++++++++++++++++++
 2 files changed, 71 insertions(+)
 create mode 100644 crates/swc_ecma_visit/tests/main.rs

diff --git a/crates/swc_ecma_visit/src/lib.rs b/crates/swc_ecma_visit/src/lib.rs
index 0a1afb8eccb5..ab4ff87a48b1 100644
--- a/crates/swc_ecma_visit/src/lib.rs
+++ b/crates/swc_ecma_visit/src/lib.rs
@@ -20,6 +20,12 @@ where
 
     B: Fold,
 {
+    #[inline(always)]
+    fn fold_program(&mut self, n: Program) -> Program {
+        let n = self.first.fold_program(n);
+        self.second.fold_program(n)
+    }
+
     #[inline(always)]
     fn fold_module(&mut self, n: Module) -> Module {
         let n = self.first.fold_module(n);
@@ -38,6 +44,11 @@ where
     A: VisitMut,
     B: VisitMut,
 {
+    fn visit_mut_program(&mut self, n: &mut Program) {
+        self.first.visit_mut_program(n);
+        self.second.visit_mut_program(n);
+    }
+
     fn visit_mut_module(&mut self, n: &mut Module) {
         self.first.visit_mut_module(n);
         self.second.visit_mut_module(n)
@@ -54,6 +65,11 @@ where
     A: Visit,
     B: Visit,
 {
+    fn visit_program(&mut self, n: &Program) {
+        self.first.visit_program(n);
+        self.second.visit_program(n);
+    }
+
     fn visit_module(&mut self, n: &Module) {
         self.first.visit_module(n);
         self.second.visit_module(n);
@@ -69,6 +85,19 @@ impl<V> Fold for Repeat<V>
 where
     V: Fold + Repeated,
 {
+    fn fold_program(&mut self, mut node: Program) -> Program {
+        loop {
+            self.pass.reset();
+            node = node.fold_with(&mut self.pass);
+
+            if !self.pass.changed() {
+                break;
+            }
+        }
+
+        node
+    }
+
     fn fold_module(&mut self, mut node: Module) -> Module {
         loop {
             self.pass.reset();
@@ -100,6 +129,17 @@ impl<V> VisitMut for Repeat<V>
 where
     V: VisitMut + Repeated,
 {
+    fn visit_mut_program(&mut self, node: &mut Program) {
+        loop {
+            self.pass.reset();
+            node.visit_mut_with(&mut self.pass);
+
+            if !self.pass.changed() {
+                break;
+            }
+        }
+    }
+
     fn visit_mut_module(&mut self, node: &mut Module) {
         loop {
             self.pass.reset();
diff --git a/crates/swc_ecma_visit/tests/main.rs b/crates/swc_ecma_visit/tests/main.rs
new file mode 100644
index 000000000000..bf128f770876
--- /dev/null
+++ b/crates/swc_ecma_visit/tests/main.rs
@@ -0,0 +1,31 @@
+use swc_common::{chain, DUMMY_SP};
+use swc_ecma_ast::{Module, Program};
+use swc_ecma_visit::{Visit, VisitWith};
+
+#[test]
+fn should_visit_program() {
+    struct Pass1<'a>(&'a mut usize);
+    struct Pass2;
+
+    impl<'a> Visit for Pass1<'a> {
+        fn visit_program(&mut self, _program: &Program) {
+            *self.0 += 1;
+        }
+    }
+
+    impl Visit for Pass2 {}
+
+    let n = Program::Module(Module {
+        span: DUMMY_SP,
+        body: vec![],
+        shebang: None,
+    });
+
+    let mut counter = 0;
+    let p1 = Pass1(&mut counter);
+    let p2 = Pass2;
+    let mut pass = chain!(p1, p2);
+    n.visit_with(&mut pass);
+
+    assert_eq!(counter, 1);
+}