[babel] Clean up import logic a bit
Addressing feedback from #2218
Lauren Tan committed
Oct 31, 2023 at 12:16 UTC
0a7bb7ac7a92238ba354638ea4304f96c3525caf
2 files changed
+52
-62
compiler/packages/babel-plugin-react-forget/src/Entrypoint/Imports.ts
+36
-16
@@ -3,7 +3,7 @@ import * as t from "@babel/types";
3
import { CompilerError } from "../CompilerError";
4
import { GeneratedSource } from "../HIR";
5
import { getOrInsertDefault } from "../Utils/utils";
6
-import { ExternalFunction } from "./Options";
6
+import { ExternalFunction, PluginOptions } from "./Options";
7
8
export function addImportsToProgram(
9
path: NodePath<t.Program>,
@@ -129,24 +129,44 @@ export function updateExistingReactImportDeclaration(
129
return didInsertUseMemoCache;
130
}
131
132
-export function insertUseMemoCacheImportDeclaration(
133
- program: NodePath<t.Program>
132
+export function updateUseMemoCacheImport(
133
+ program: NodePath<t.Program>,
134
+ options: PluginOptions
135
): void {
135
- program.unshiftContainer(
136
- "body",
137
- t.importDeclaration(
138
- [
139
- t.importSpecifier(
140
- t.identifier("useMemoCache"),
141
- t.identifier("unstable_useMemoCache")
142
- ),
143
- ],
144
- t.stringLiteral("react")
145
- )
146
- );
136
+ // If there isn't already an import of * as React, insert it so useMemoCache doesn't
137
+ // throw
138
+ const { didInsertUseMemoCache, hasExistingReactImport } =
139
+ findExistingImports(program);
140
+
141
+ // unstable_useMemoCache wasn't already imported, nothing to do
142
+ if (!didInsertUseMemoCache) {
143
+ return;
144
+ }
145
+
146
+ if (
147
+ options.useMemoCacheSource === "react" ||
148
+ options.useMemoCacheSource === null
149
+ ) {
150
+ // If Forget did successfully compile inject/update an import of
151
+ // `import {unstable_useMemoCache as useMemoCache} from 'react'` and rename
152
+ // `React.unstable_useMemoCache(n)` to `useMemoCache(n)`;
153
+ if (hasExistingReactImport) {
154
+ const didUpdateImport = updateExistingReactImportDeclaration(program);
155
+ if (didUpdateImport === false) {
156
+ throw new Error(
157
+ "Expected an ImportDeclaration of react in order to update ImportSpecifiers with useMemoCache"
158
+ );
159
+ }
160
+ } else {
161
+ addUseMemoCacheImportDeclaration(program, "react");
162
+ }
163
+ } else if (typeof options.useMemoCacheSource === "string") {
164
+ // import useMemoCache from userspace module
165
+ addUseMemoCacheImportDeclaration(program, options.useMemoCacheSource);
166
+ }
167
}
168
149
-export function insertUserspaceUseMemoCacheImportDeclaration(
169
+function addUseMemoCacheImportDeclaration(
170
program: NodePath<t.Program>,
171
moduleName: string
172
): void {
compiler/packages/babel-plugin-react-forget/src/Entrypoint/Program.ts
+16
-46
@@ -17,13 +17,7 @@ import { CodegenFunction } from "../ReactiveScopes";
17
import { isComponentDeclaration } from "../Utils/ComponentDeclaration";
18
import { assertExhaustive } from "../Utils/utils";
19
import { insertGatedFunctionDeclaration } from "./Gating";
20
-import {
21
- addImportsToProgram,
22
- findExistingImports,
23
- insertUseMemoCacheImportDeclaration,
24
- insertUserspaceUseMemoCacheImportDeclaration,
25
- updateExistingReactImportDeclaration,
26
-} from "./Imports";
20
+import { addImportsToProgram, updateUseMemoCacheImport } from "./Imports";
21
import { addInstrumentForget } from "./Instrumentation";
22
import { ExternalFunction, PluginOptions, parsePluginOptions } from "./Options";
23
import { compileFn } from "./Pipeline";
@@ -316,47 +310,23 @@ export function compileProgram(
310
}
311
);
312
319
- // If there isn't already an import of * as React, insert it so useMemoCache doesn't
320
- // throw
313
+ // Forget compiled the component, we need to update existing imports of unstable_useMemoCache
314
if (hasForgetMutatedOriginalSource) {
322
- const { didInsertUseMemoCache, hasExistingReactImport } =
323
- findExistingImports(program);
324
-
325
- // If Forget did successfully compile inject/update an import of
326
- // `import {unstable_useMemoCache as useMemoCache} from 'react'` and rename
327
- // `React.unstable_useMemoCache(n)` to `useMemoCache(n)`;
328
- if (didInsertUseMemoCache) {
329
- if (options.useMemoCacheSource === null) {
330
- if (hasExistingReactImport) {
331
- const didUpdateImport = updateExistingReactImportDeclaration(program);
332
- if (didUpdateImport === false) {
333
- throw new Error(
334
- "Expected an ImportDeclaration of react in order to update ImportSpecifiers with useMemoCache"
335
- );
336
- }
337
- } else {
338
- insertUseMemoCacheImportDeclaration(program);
339
- }
340
- } else if (typeof options.useMemoCacheSource === "string") {
341
- insertUserspaceUseMemoCacheImportDeclaration(
342
- program,
343
- options.useMemoCacheSource
344
- );
345
- }
346
- }
347
- const externalFunctions: ExternalFunction[] = [];
348
- // TODO: check for duplicate import specifiers
349
- if (options.gating != null) {
350
- externalFunctions.push(options.gating);
351
- }
352
- if (options.instrumentForget != null) {
353
- externalFunctions.push(options.instrumentForget);
354
- }
355
- if (options.environment?.enableEmitFreeze != null) {
356
- externalFunctions.push(options.environment.enableEmitFreeze);
357
- }
358
- addImportsToProgram(program, externalFunctions);
315
+ updateUseMemoCacheImport(program, options);
316
+ }
317
+
318
+ const externalFunctions: ExternalFunction[] = [];
319
+ // TODO: check for duplicate import specifiers
320
+ if (options.gating != null) {
321
+ externalFunctions.push(options.gating);
322
+ }
323
+ if (options.instrumentForget != null) {
324
+ externalFunctions.push(options.instrumentForget);
325
+ }
326
+ if (options.environment?.enableEmitFreeze != null) {
327
+ externalFunctions.push(options.environment.enableEmitFreeze);
328
}
329
+ addImportsToProgram(program, externalFunctions);
330
}
331
332
function shouldVisitNode(