From ff571d8d85adf6f26018686a9bd4a4c651f6430c Mon Sep 17 00:00:00 2001
From: Sathya Gunasekaran <gsathya.ceg@gmail.com>
Date: Mon, 20 Mar 2023 17:42:57 +0000
Subject: [PATCH] [Babel] Refactor compilation checks into a separate function

---
 compiler/forget/src/Babel/BabelPlugin.ts | 35 +++++++++++++++++-------
 1 file changed, 25 insertions(+), 10 deletions(-)

diff --git a/compiler/forget/src/Babel/BabelPlugin.ts b/compiler/forget/src/Babel/BabelPlugin.ts
index 249738cd5d..a24eb2ee0c 100644
--- a/compiler/forget/src/Babel/BabelPlugin.ts
+++ b/compiler/forget/src/Babel/BabelPlugin.ts
@@ -43,16 +43,6 @@ export default function ReactForgetBabelPlugin(
     fn: BabelCore.NodePath<t.FunctionDeclaration>,
     pass: BabelPluginPass
   ): void {
-    if (pass.opts.enableOnlyOnUseForgetDirective) {
-      if (!hasUseForgetDirective(fn.node.body.directives)) {
-        return;
-      }
-    }
-
-    if (fn.scope.getProgramParent() !== fn.scope.parent) {
-      return;
-    }
-
     hasForgetCompiledCode = true;
     const compiled = compile(fn, pass.opts.environment);
 
@@ -84,6 +74,10 @@ export default function ReactForgetBabelPlugin(
       fn: BabelCore.NodePath<t.FunctionDeclaration>,
       pass: BabelPluginPass
     ): void {
+      if (!shouldCompile(fn, pass)) {
+        return;
+      }
+
       visitFn(fn, pass);
     },
   };
@@ -120,6 +114,27 @@ export default function ReactForgetBabelPlugin(
   };
 }
 
+function shouldCompile(
+  fn: BabelCore.NodePath<t.FunctionDeclaration | t.ArrowFunctionExpression>,
+  pass: BabelPluginPass
+): boolean {
+  if (pass.opts.enableOnlyOnUseForgetDirective) {
+    const body = fn.get("body");
+    if (!body.isBlockStatement()) {
+      return false;
+    }
+    if (!hasUseForgetDirective(body.node.directives)) {
+      return false;
+    }
+  }
+
+  if (fn.scope.getProgramParent() !== fn.scope.parent) {
+    return false;
+  }
+
+  return true;
+}
+
 function buildGatingTest(
   uncompiled: BabelCore.NodePath<t.FunctionDeclaration>,
   compiled: t.Identifier,