From 298bf55f25ce698c16bcd9fc9cdb60e810b8a6fb Mon Sep 17 00:00:00 2001
From: Lauren Tan <poteto@users.noreply.github.com>
Date: Mon, 14 Aug 2023 12:05:52 -0400
Subject: [PATCH] Address feedback

---
 .../src/Entrypoint/Program.ts                 | 67 +++++++++++--------
 1 file changed, 40 insertions(+), 27 deletions(-)

diff --git a/compiler/forget/packages/babel-plugin-react-forget/src/Entrypoint/Program.ts b/compiler/forget/packages/babel-plugin-react-forget/src/Entrypoint/Program.ts
index 486e53f459..609a1fca22 100644
--- a/compiler/forget/packages/babel-plugin-react-forget/src/Entrypoint/Program.ts
+++ b/compiler/forget/packages/babel-plugin-react-forget/src/Entrypoint/Program.ts
@@ -96,31 +96,46 @@ function compileAndInsertNewFunctionDeclaration(
     });
     const originalIdent = fnPath.node.id;
 
-    let gatedFn = null;
-    if (pass.opts.gating != null) {
-      gatedFn = insertGatedFunctionDeclaration(
-        fnPath,
-        compiledFn,
-        originalIdent,
-        pass.opts.gating
-      );
-    } else {
-      fnPath.replaceWith(compiledFn);
-    }
-
-    if (pass.opts.instrumentForget != null) {
-      const instrumentFnName = pass.opts.instrumentForget.importSpecifierName;
-      addInstrumentForget(fnPath, originalIdent.name, instrumentFnName);
-      if (pass.opts.gating != null && gatedFn != null) {
-        addInstrumentForget(gatedFn, originalIdent.name, instrumentFnName);
-      }
-    }
+    insertNewFunctionDeclaration(fnPath, originalIdent, compiledFn, pass);
     hasForgetMutatedOriginalSource = true;
   }
 
   return hasForgetMutatedOriginalSource;
 }
 
+function insertNewFunctionDeclaration(
+  fnPath: NodePath<t.FunctionDeclaration>,
+  originalIdent: t.Identifier,
+  compiledFn: t.FunctionDeclaration,
+  pass: CompilerPass
+): void {
+  let gatedFn = null;
+  if (pass.opts.gating != null) {
+    gatedFn = insertGatedFunctionDeclaration(
+      fnPath,
+      compiledFn,
+      originalIdent,
+      pass.opts.gating
+    );
+  } else {
+    fnPath.replaceWith(compiledFn);
+  }
+
+  if (pass.opts.instrumentForget != null) {
+    const instrumentFnName = pass.opts.instrumentForget.importSpecifierName;
+    addInstrumentForget(fnPath, originalIdent.name, instrumentFnName);
+    if (pass.opts.gating != null) {
+      CompilerError.invariant(gatedFn != null, {
+        reason: "Should have inserted a gated function declaration",
+        description: null,
+        loc: null,
+        suggestions: null,
+      });
+      addInstrumentForget(gatedFn, originalIdent.name, instrumentFnName);
+    }
+  }
+}
+
 export function compileProgram(
   program: NodePath<t.Program>,
   pass: CompilerPass
@@ -205,10 +220,9 @@ export function compileProgram(
           return;
         }
 
-        hasForgetMutatedOriginalSource = compileAndInsertNewFunctionDeclaration(
-          fn,
-          pass
-        );
+        if (compileAndInsertNewFunctionDeclaration(fn, pass) === true) {
+          hasForgetMutatedOriginalSource = true;
+        }
       },
 
       ArrowFunctionExpression(
@@ -239,10 +253,9 @@ export function compileProgram(
           return;
         }
 
-        hasForgetMutatedOriginalSource = compileAndInsertNewFunctionDeclaration(
-          loweredFn,
-          pass
-        );
+        if (compileAndInsertNewFunctionDeclaration(loweredFn, pass) === true) {
+          hasForgetMutatedOriginalSource = true;
+        }
       },
     },
     {