From 611d1a4ba0008d67b5dcda485177f0813b2b573e Mon Sep 17 00:00:00 2001
From: liuxingbaoyu <30521560+liuxingbaoyu@users.noreply.github.com>
Date: Tue, 10 Dec 2024 18:17:35 +0800
Subject: [PATCH] Improve `babel-plugin-jest-hoist` logic (#15415)

---
 CHANGELOG.md                                  |  1 +
 packages/babel-plugin-jest-hoist/src/index.ts | 94 ++++++++-----------
 scripts/lintTs.mjs                            |  6 ++
 3 files changed, 48 insertions(+), 53 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 363abee8a97e..cd34d1360be7 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -50,6 +50,7 @@
 ### Fixes
 
 - `[babel-plugin-jest-hoist]` Use `denylist` instead of the deprecated `blacklist` for Babel 8 support ([#14109](https://github.com/jestjs/jest/pull/14109))
+- `[babel-plugin-jest-hoist]` Do not rely on buggy Babel behaviour ([#15415](https://github.com/jestjs/jest/pull/15415))
 - `[expect]` Check error instance type for `toThrow/toThrowError` ([#14576](https://github.com/jestjs/jest/pull/14576))
 - `[expect]` Improve diff for failing `expect.objectContaining` ([#15038](https://github.com/jestjs/jest/pull/15038))
 - `[expect]` Use `Array.isArray` to check if an array is an `Array` ([#15101](https://github.com/jestjs/jest/pull/15101))
diff --git a/packages/babel-plugin-jest-hoist/src/index.ts b/packages/babel-plugin-jest-hoist/src/index.ts
index 381fdad2560e..2885e386ac57 100644
--- a/packages/babel-plugin-jest-hoist/src/index.ts
+++ b/packages/babel-plugin-jest-hoist/src/index.ts
@@ -10,20 +10,17 @@ import type {PluginObj} from '@babel/core';
 import {statement} from '@babel/template';
 import type {NodePath} from '@babel/traverse';
 import {
-  type BlockStatement,
   type CallExpression,
-  type EmptyStatement,
   type Expression,
   type Identifier,
   type ImportDeclaration,
   type MemberExpression,
   type Node,
-  type Program,
+  type Statement,
   type Super,
   type VariableDeclaration,
   type VariableDeclarator,
   callExpression,
-  emptyStatement,
   isIdentifier,
   variableDeclaration,
 } from '@babel/types';
@@ -312,53 +309,6 @@ const extractJestObjExprIfHoistable = (expr: NodePath): JestObjInfo | null => {
   return null;
 };
 
-function visitBlock(block: NodePath<BlockStatement> | NodePath<Program>) {
-  // use a temporary empty statement instead of the real first statement, which may itself be hoisted
-  const [varsHoistPoint, callsHoistPoint] = block.unshiftContainer<
-    BlockStatement | Program,
-    'body',
-    [EmptyStatement, EmptyStatement]
-  >('body', [emptyStatement(), emptyStatement()]);
-  block.traverse({
-    CallExpression: visitCallExpr,
-    VariableDeclarator: visitVariableDeclarator,
-    // do not traverse into nested blocks, or we'll hoist calls in there out to this block
-    denylist: ['BlockStatement'],
-  });
-  callsHoistPoint.remove();
-  varsHoistPoint.remove();
-
-  function visitCallExpr(callExpr: NodePath<CallExpression>) {
-    if (hoistedJestGetters.has(callExpr.node)) {
-      const mockStmt = callExpr.getStatementParent();
-
-      if (mockStmt) {
-        const mockStmtParent = mockStmt.parentPath;
-        if (mockStmtParent.isBlock()) {
-          const mockStmtNode = mockStmt.node;
-          mockStmt.remove();
-          callsHoistPoint.insertBefore(mockStmtNode);
-        }
-      }
-    }
-  }
-
-  function visitVariableDeclarator(varDecl: NodePath<VariableDeclarator>) {
-    if (hoistedVariables.has(varDecl.node)) {
-      // should be assert function, but it's not. So let's cast below
-      varDecl.parentPath.assertVariableDeclaration();
-
-      const {kind, declarations} = varDecl.parent as VariableDeclaration;
-      if (declarations.length === 1) {
-        varDecl.parentPath.remove();
-      } else {
-        varDecl.remove();
-      }
-      varsHoistPoint.insertBefore(variableDeclaration(kind, [varDecl.node]));
-    }
-  }
-}
-
 /* eslint-disable sort-keys */
 export default function jestHoist(): PluginObj<{
   declareJestObjGetterIdentifier: () => Identifier;
@@ -404,8 +354,46 @@ export default function jestHoist(): PluginObj<{
     },
     // in `post` to make sure we come after an import transform and can unshift above the `require`s
     post({path: program}) {
-      visitBlock(program);
-      program.traverse({BlockStatement: visitBlock});
+      type Item = {calls: Array<Statement>; vars: Array<Statement>};
+
+      const stack: Array<Item> = [{calls: [], vars: []}];
+      program.traverse({
+        BlockStatement: {
+          enter() {
+            stack.push({calls: [], vars: []});
+          },
+          exit(path) {
+            const item = stack.pop()!;
+            path.node.body.unshift(...item.vars, ...item.calls);
+          },
+        },
+        CallExpression(callExpr: NodePath<CallExpression>) {
+          if (hoistedJestGetters.has(callExpr.node)) {
+            const mockStmt = callExpr.getStatementParent();
+
+            if (mockStmt && mockStmt.parentPath.isBlock()) {
+              stack.at(-1)!.calls.push(mockStmt.node);
+              mockStmt.remove();
+            }
+          }
+        },
+        VariableDeclarator(varDecl: NodePath<VariableDeclarator>) {
+          if (hoistedVariables.has(varDecl.node)) {
+            // should be assert function, but it's not. So let's cast below
+            varDecl.parentPath.assertVariableDeclaration();
+
+            const {kind, declarations} = varDecl.parent as VariableDeclaration;
+            if (declarations.length === 1) {
+              varDecl.parentPath.remove();
+            } else {
+              varDecl.remove();
+            }
+            stack.at(-1)!.vars.push(variableDeclaration(kind, [varDecl.node]));
+          }
+        },
+      });
+      const item = stack.pop()!;
+      program.node.body.unshift(...item.vars, ...item.calls);
     },
   };
 }
diff --git a/scripts/lintTs.mjs b/scripts/lintTs.mjs
index 5f92e8380c9e..01c4d353a74d 100644
--- a/scripts/lintTs.mjs
+++ b/scripts/lintTs.mjs
@@ -111,6 +111,12 @@ try {
                   '@typescript-eslint/no-redundant-type-constituents': 'off',
                 },
               },
+              // {
+              //   files: ['packages/babel-plugin-jest-hoist/src/index.ts'],
+              //   rules: {
+              //     '@typescript-eslint/strict-boolean-expressions': 'off',
+              //   },
+              // },
             ],
             parser: '@typescript-eslint/parser',
             parserOptions: {