[be] Move more logic into Imports module

compiler/forget/packages/babel-plugin-react-forget/src/Entrypoint/Imports.ts

@@ -66,3 +66,77 @@ export function isNonNamespacedImportOfReact(
       .every((specifier) => specifier.isImportSpecifier())
   );
 }
+
+export function findExistingImports(program: NodePath<t.Program>): {
+  didInsertUseMemoCache: boolean;
+  hasExistingReactImport: boolean;
+} {
+  let didInsertUseMemoCache = false;
+  let hasExistingReactImport = false;
+  program.traverse({
+    CallExpression(callExprPath) {
+      const callee = callExprPath.get("callee");
+      const args = callExprPath.get("arguments");
+      if (
+        callee.isIdentifier() &&
+        callee.node.name === "useMemoCache" &&
+        args.length === 1 &&
+        args[0].isNumericLiteral()
+      ) {
+        didInsertUseMemoCache = true;
+      }
+    },
+    ImportDeclaration(importDeclPath) {
+      if (isNonNamespacedImportOfReact(importDeclPath)) {
+        hasExistingReactImport = true;
+      }
+    },
+  });
+
+  return {
+    didInsertUseMemoCache,
+    hasExistingReactImport,
+  };
+}
+
+/**
+ * If an existing import of React exists (ie `import {useMemo} from 'React'`), inject useMemoCache
+ * into the list of destructured variables.
+ */
+export function updateExistingReactImportDeclaration(
+  program: NodePath<t.Program>
+): boolean {
+  let didInsertUseMemoCache = false;
+  program.traverse({
+    ImportDeclaration(importDeclPath) {
+      if (isNonNamespacedImportOfReact(importDeclPath)) {
+        importDeclPath.pushContainer(
+          "specifiers",
+          t.importSpecifier(
+            t.identifier("useMemoCache"),
+            t.identifier("unstable_useMemoCache")
+          )
+        );
+        didInsertUseMemoCache = true;
+      }
+    },
+  });
+  return didInsertUseMemoCache;
+}
+
+export function insertUseMemoCacheImportDeclaration(
+  program: NodePath<t.Program>
+): void {
+  program.unshiftContainer(
+    "body",
+    t.importDeclaration(
+      [
+        t.importSpecifier(
+          t.identifier("useMemoCache"),
+          t.identifier("unstable_useMemoCache")
+        ),
+      ],
+      t.stringLiteral("react")
+    )
+  );
+}

compiler/forget/packages/babel-plugin-react-forget/src/Entrypoint/Program.ts

@@ -16,7 +16,12 @@ import {
 import { GeneratedSource } from "../HIR";
 import { isComponentDeclaration } from "../Utils/ComponentDeclaration";
 import { insertGatedFunctionDeclaration } from "./Gating";
-import { addImportsToProgram, isNonNamespacedImportOfReact } from "./Imports";
+import {
+  addImportsToProgram,
+  findExistingImports,
+  insertUseMemoCacheImportDeclaration,
+  updateExistingReactImportDeclaration,
+} from "./Imports";
 import { addInstrumentForget } from "./Instrumentation";
 import { PluginOptions, parsePluginOptions } from "./Options";
 import { compileFn } from "./Pipeline";
@@ -250,65 +255,22 @@ export function compileProgram(
   // If there isn't already an import of * as React, insert it so useMemoCache doesn't
   // throw
   if (hasForgetCompiledCode) {
-    let didInsertUseMemoCache = false;
-    let hasExistingReactImport = false;
-    program.traverse({
-      CallExpression(callExprPath) {
-        const callee = callExprPath.get("callee");
-        const args = callExprPath.get("arguments");
-        if (
-          callee.isIdentifier() &&
-          callee.node.name === "useMemoCache" &&
-          args.length === 1 &&
-          args[0].isNumericLiteral()
-        ) {
-          didInsertUseMemoCache = true;
-        }
-      },
-      ImportDeclaration(importDeclPath) {
-        if (isNonNamespacedImportOfReact(importDeclPath)) {
-          hasExistingReactImport = true;
-        }
-      },
-    });
+    const { didInsertUseMemoCache, hasExistingReactImport } =
+      findExistingImports(program);
+
     // If Forget did successfully compile inject/update an import of
     // `import {unstable_useMemoCache as useMemoCache} from 'react'` and rename
     // `React.unstable_useMemoCache(n)` to `useMemoCache(n)`;
     if (didInsertUseMemoCache) {
       if (hasExistingReactImport) {
-        let didUpdateImport = false;
-        program.traverse({
-          ImportDeclaration(importDeclPath) {
-            if (isNonNamespacedImportOfReact(importDeclPath)) {
-              importDeclPath.pushContainer(
-                "specifiers",
-                t.importSpecifier(
-                  t.identifier("useMemoCache"),
-                  t.identifier("unstable_useMemoCache")
-                )
-              );
-              didUpdateImport = true;
-            }
-          },
-        });
+        const didUpdateImport = updateExistingReactImportDeclaration(program);
         if (didUpdateImport === false) {
           throw new Error(
             "Expected an ImportDeclaration of react in order to update ImportSpecifiers with useMemoCache"
           );
         }
       } else {
-        program.unshiftContainer(
-          "body",
-          t.importDeclaration(
-            [
-              t.importSpecifier(
-                t.identifier("useMemoCache"),
-                t.identifier("unstable_useMemoCache")
-              ),
-            ],
-            t.stringLiteral("react")
-          )
-        );
+        insertUseMemoCacheImportDeclaration(program);
       }
     }
     const externalFunctions = [];