impl undo_blockify_for_stmt_body

swc-project · Dec 8, 2022 · b93352f · b93352f
1 parent 3924889
commit b93352f
Showing 1 changed file with 21 additions and 4 deletions.
diff --git a/crates/swc_ecma_transforms_compat/src/es2015/block_scoping/mod.rs b/crates/swc_ecma_transforms_compat/src/es2015/block_scoping/mod.rs
@@ -377,12 +377,26 @@ impl BlockScoping {
     /// }
     /// ```
     /// which fixes https://github.com/swc-project/swc/issues/6573
-    fn blockify_for_stmt_body(&self, body: &mut Box<Stmt>) {
+    fn blockify_for_stmt_body(&self, body: &mut Box<Stmt>) -> bool {
         if !body.is_block() {
             *body = Box::new(Stmt::Block(BlockStmt {
                 span: Default::default(),
                 stmts: vec![*body.take()],
             }));
+            true
+        } else {
+            false
+        }
+    }
+
+    fn undo_blockify_for_stmt_body(&self, body: &mut Box<Stmt>, blockifyed: bool) {
+        if blockifyed {
+            let stmt = body
+                .as_mut_block()
+                .and_then(|block| (block.stmts.len() == 1).then(|| block.stmts[0].take()));
+            if let Some(stmt) = stmt {
+                *body = Box::new(stmt)
+            }
         }
     }
 }
@@ -424,7 +438,7 @@ impl VisitMut for BlockScoping {
     }
 
     fn visit_mut_for_in_stmt(&mut self, node: &mut ForInStmt) {
-        self.blockify_for_stmt_body(&mut node.body);
+        let blockifyed = self.blockify_for_stmt_body(&mut node.body);
         let lexical_var = if let VarDeclOrPat::VarDecl(decl) = &node.left {
             find_lexical_vars(decl)
         } else {
@@ -445,10 +459,11 @@ impl VisitMut for BlockScoping {
 
         self.visit_mut_with_scope(kind, &mut node.body);
         self.handle_capture_of_vars(&mut node.body);
+        self.undo_blockify_for_stmt_body(&mut node.body, blockifyed);
     }
 
     fn visit_mut_for_of_stmt(&mut self, node: &mut ForOfStmt) {
-        self.blockify_for_stmt_body(&mut node.body);
+        let blockifyed = self.blockify_for_stmt_body(&mut node.body);
         let vars = if let VarDeclOrPat::VarDecl(decl) = &node.left {
             find_lexical_vars(decl)
         } else {
@@ -470,10 +485,11 @@ impl VisitMut for BlockScoping {
 
         self.visit_mut_with_scope(kind, &mut node.body);
         self.handle_capture_of_vars(&mut node.body);
+        self.undo_blockify_for_stmt_body(&mut node.body, blockifyed);
     }
 
     fn visit_mut_for_stmt(&mut self, node: &mut ForStmt) {
-        self.blockify_for_stmt_body(&mut node.body);
+        let blockifyed = self.blockify_for_stmt_body(&mut node.body);
         let lexical_var = if let Some(VarDeclOrExpr::VarDecl(decl)) = &node.init {
             find_lexical_vars(decl)
         } else {
@@ -494,6 +510,7 @@ impl VisitMut for BlockScoping {
         };
         self.visit_mut_with_scope(kind, &mut node.body);
         self.handle_capture_of_vars(&mut node.body);
+        self.undo_blockify_for_stmt_body(&mut node.body, blockifyed);
     }
 
     fn visit_mut_function(&mut self, f: &mut Function) {