From 8bb537209cd046691516d72a4276820428c48051 Mon Sep 17 00:00:00 2001
From: Tim Seckinger <seckinger.tim@gmail.com>
Date: Wed, 28 Oct 2020 21:09:31 +0100
Subject: [PATCH] refactor: pure traversal solution

---
 packages/babel-plugin-jest-hoist/src/index.ts | 62 +++++++++++++------
 1 file changed, 43 insertions(+), 19 deletions(-)

diff --git a/packages/babel-plugin-jest-hoist/src/index.ts b/packages/babel-plugin-jest-hoist/src/index.ts
index 5da3dbd0d825..15e2c5bf5ae7 100644
--- a/packages/babel-plugin-jest-hoist/src/index.ts
+++ b/packages/babel-plugin-jest-hoist/src/index.ts
@@ -8,12 +8,15 @@
 
 import type {NodePath} from '@babel/traverse';
 import {
+  BlockStatement,
+  CallExpression,
   Expression,
   Identifier,
   Node,
   Program,
   callExpression,
-  clone,
+  emptyStatement,
+  isIdentifier,
 } from '@babel/types';
 import {statement} from '@babel/template';
 import type {PluginObj} from '@babel/core';
@@ -113,7 +116,7 @@ FUNCTIONS.mock = args => {
 
     const ids: Set<NodePath<Identifier>> = new Set();
     const parentScope = moduleFactory.parentPath.scope;
-    // @ts-expect-error: ReferencedIdentifier is not known on visitors
+    // @ts-expect-error: ReferencedIdentifier and blacklist are not known on visitors
     moduleFactory.traverse(IDVisitor, {ids});
     for (const id of ids) {
       const {name} = id.node;
@@ -276,31 +279,52 @@ export default (): PluginObj<{
         jestObjExpr.replaceWith(
           callExpression(this.declareJestObjGetterIdentifier(), []),
         );
-        exprStmt.setData('_jestShouldHoist', true);
       }
     },
   },
   // in `post` to make sure we come after an import transform and can unshift above the `require`s
   post({path: program}: {path: NodePath<Program>}) {
+    const self = this;
+    visitBlock(program);
     program.traverse({
-      ExpressionStatement: exprStmt => {
-        if (exprStmt.getData('_jestShouldHoist')) {
-          const prevSiblings = exprStmt.getAllPrevSiblings();
-          const unhoistedSiblings = prevSiblings
-            .reverse()
-            .filter(
-              s =>
-                !s.getData('_jestShouldHoist') && !s.getData('_jestWasHoisted'),
-            );
-          const exprNode = exprStmt.node;
-          if (unhoistedSiblings.length) {
-            const hoisted = unhoistedSiblings[0].insertBefore(clone(exprNode));
-            hoisted[0].setData('_jestWasHoisted', true);
-            exprStmt.remove();
+      BlockStatement: visitBlock,
+    });
+
+    function visitBlock(block: NodePath<BlockStatement> | NodePath<Program>) {
+      // use a temporary empty statement instead of the real first statement, which may itself be hoisted
+      const [firstNonHoistedStatementOfBlock] = block.unshiftContainer(
+        'body',
+        emptyStatement(),
+      );
+      block.traverse({
+        CallExpression: visitCallExpr,
+        // do not traverse into nested blocks, or we'll hoist calls in there out to this block
+        // @ts-expect-error blacklist is not known
+        blacklist: ['BlockStatement'],
+      });
+      firstNonHoistedStatementOfBlock.remove();
+
+      function visitCallExpr(callExpr: NodePath<CallExpression>) {
+        const {
+          node: {callee},
+        } = callExpr;
+        if (
+          isIdentifier(callee) &&
+          callee.name === self.jestObjGetterIdentifier?.name
+        ) {
+          const mockStmt = callExpr.getStatementParent();
+
+          if (mockStmt) {
+            const mockStmtNode = mockStmt.node;
+            const mockStmtParent = mockStmt.parentPath;
+            if (mockStmtParent.isBlock()) {
+              mockStmt.remove();
+              firstNonHoistedStatementOfBlock.insertBefore(mockStmtNode);
+            }
           }
         }
-      },
-    });
+      }
+    }
   },
 });
 /* eslint-enable */