diff --git a/compiler/forget/packages/babel-plugin-react-forget/src/Entrypoint/Imports.ts b/compiler/forget/packages/babel-plugin-react-forget/src/Entrypoint/Imports.ts index 79dc41653d..0fca72f1f9 100644 --- a/compiler/forget/packages/babel-plugin-react-forget/src/Entrypoint/Imports.ts +++ b/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): { + 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 +): 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 +): void { + program.unshiftContainer( + "body", + t.importDeclaration( + [ + t.importSpecifier( + t.identifier("useMemoCache"), + t.identifier("unstable_useMemoCache") + ), + ], + t.stringLiteral("react") + ) + ); +} 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 b142d99fcf..6480098a26 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 @@ -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 = [];