From 93b79e4aa048677094a72d799ef2480761424303 Mon Sep 17 00:00:00 2001
From: Sathya Gunasekaran <gsathya.ceg@gmail.com>
Date: Mon, 20 Mar 2023 17:42:56 +0000
Subject: [PATCH] [Babel] Refactor visitor to separate function

This lets us extend to ArrowFunctionExpression
---
 compiler/forget/src/Babel/BabelPlugin.ts | 75 +++++++++++++-----------
 1 file changed, 41 insertions(+), 34 deletions(-)

diff --git a/compiler/forget/src/Babel/BabelPlugin.ts b/compiler/forget/src/Babel/BabelPlugin.ts
index 7a75ed24a2..249738cd5d 100644
--- a/compiler/forget/src/Babel/BabelPlugin.ts
+++ b/compiler/forget/src/Babel/BabelPlugin.ts
@@ -39,45 +39,52 @@ export default function ReactForgetBabelPlugin(
 ): BabelCore.PluginObj {
   let hasForgetCompiledCode: boolean = false;
 
+  function visitFn(
+    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);
+
+    if (pass.opts.gatingModule) {
+      // Rename existing function
+      invariant(fn.node.id, "FunctionDeclaration must have a name");
+      const original = fn.node.id;
+      fn.node.id = addSuffix(fn.node.id, "_uncompiled");
+
+      // Rename and append compiled function
+      invariant(compiled.id, "FunctionDeclaration must produce a name");
+      compiled.id = addSuffix(compiled.id, "_forget");
+      const compiledFn = fn.insertAfter(compiled)[0];
+      compiledFn.skip();
+
+      // Build and append gating test
+      compiledFn.insertAfter(buildGatingTest(fn, compiled.id, original));
+    } else {
+      fn.replaceWith(compiled);
+    }
+
+    // We are generating a new FunctionDeclaration node, so we must skip over it or this
+    // traversal will loop infinitely.
+    fn.skip();
+  }
+
   const visitor = {
     FunctionDeclaration(
       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);
-
-      if (pass.opts.gatingModule) {
-        // Rename existing function
-        invariant(fn.node.id, "FunctionDeclaration must have a name");
-        const original = fn.node.id;
-        fn.node.id = addSuffix(fn.node.id, "_uncompiled");
-
-        // Rename and append compiled function
-        invariant(compiled.id, "FunctionDeclaration must produce a name");
-        compiled.id = addSuffix(compiled.id, "_forget");
-        const compiledFn = fn.insertAfter(compiled)[0];
-        compiledFn.skip();
-
-        // Build and append gating test
-        compiledFn.insertAfter(buildGatingTest(fn, compiled.id, original));
-      } else {
-        fn.replaceWith(compiled);
-      }
-
-      // We are generating a new FunctionDeclaration node, so we must skip over it or this
-      // traversal will loop infinitely.
-      fn.skip();
+      visitFn(fn, pass);
     },
   };