[be] Move more logic into Imports module
Lauren Tan committed
Aug 11, 2023 at 16:39 UTC
25e3575c38d1d4755acd56b8dc5d94866bfa3d09
2 files changed
+85
-49
compiler/forget/packages/babel-plugin-react-forget/src/Entrypoint/Imports.ts
+74
@@ -66,3 +66,77 @@ export function isNonNamespacedImportOfReact(
66
.every((specifier) => specifier.isImportSpecifier())
67
);
68
}
69
+
70
+export function findExistingImports(program: NodePath<t.Program>): {
71
+ didInsertUseMemoCache: boolean;
72
+ hasExistingReactImport: boolean;
73
+} {
74
+ let didInsertUseMemoCache = false;
75
+ let hasExistingReactImport = false;
76
+ program.traverse({
77
+ CallExpression(callExprPath) {
78
+ const callee = callExprPath.get("callee");
79
+ const args = callExprPath.get("arguments");
80
+ if (
81
+ callee.isIdentifier() &&
82
+ callee.node.name === "useMemoCache" &&
83
+ args.length === 1 &&
84
+ args[0].isNumericLiteral()
85
+ ) {
86
+ didInsertUseMemoCache = true;
87
+ }
88
+ },
89
+ ImportDeclaration(importDeclPath) {
90
+ if (isNonNamespacedImportOfReact(importDeclPath)) {
91
+ hasExistingReactImport = true;
92
+ }
93
+ },
94
+ });
95
+
96
+ return {
97
+ didInsertUseMemoCache,
98
+ hasExistingReactImport,
99
+ };
100
+}
101
+
102
+/**
103
+ * If an existing import of React exists (ie `import {useMemo} from 'React'`), inject useMemoCache
104
+ * into the list of destructured variables.
105
+ */
106
+export function updateExistingReactImportDeclaration(
107
+ program: NodePath<t.Program>
108
+): boolean {
109
+ let didInsertUseMemoCache = false;
110
+ program.traverse({
111
+ ImportDeclaration(importDeclPath) {
112
+ if (isNonNamespacedImportOfReact(importDeclPath)) {
113
+ importDeclPath.pushContainer(
114
+ "specifiers",
115
+ t.importSpecifier(
116
+ t.identifier("useMemoCache"),
117
+ t.identifier("unstable_useMemoCache")
118
+ )
119
+ );
120
+ didInsertUseMemoCache = true;
121
+ }
122
+ },
123
+ });
124
+ return didInsertUseMemoCache;
125
+}
126
+
127
+export function insertUseMemoCacheImportDeclaration(
128
+ program: NodePath<t.Program>
129
+): void {
130
+ program.unshiftContainer(
131
+ "body",
132
+ t.importDeclaration(
133
+ [
134
+ t.importSpecifier(
135
+ t.identifier("useMemoCache"),
136
+ t.identifier("unstable_useMemoCache")
137
+ ),
138
+ ],
139
+ t.stringLiteral("react")
140
+ )
141
+ );
142
+}
compiler/forget/packages/babel-plugin-react-forget/src/Entrypoint/Program.ts
+11
-49
@@ -16,7 +16,12 @@ import {
16
import { GeneratedSource } from "../HIR";
17
import { isComponentDeclaration } from "../Utils/ComponentDeclaration";
18
import { insertGatedFunctionDeclaration } from "./Gating";
19
-import { addImportsToProgram, isNonNamespacedImportOfReact } from "./Imports";
19
+import {
20
+ addImportsToProgram,
21
+ findExistingImports,
22
+ insertUseMemoCacheImportDeclaration,
23
+ updateExistingReactImportDeclaration,
24
+} from "./Imports";
25
import { addInstrumentForget } from "./Instrumentation";
26
import { PluginOptions, parsePluginOptions } from "./Options";
27
import { compileFn } from "./Pipeline";
@@ -250,65 +255,22 @@ export function compileProgram(
255
// If there isn't already an import of * as React, insert it so useMemoCache doesn't
256
// throw
257
if (hasForgetCompiledCode) {
253
- let didInsertUseMemoCache = false;
254
- let hasExistingReactImport = false;
255
- program.traverse({
256
- CallExpression(callExprPath) {
257
- const callee = callExprPath.get("callee");
258
- const args = callExprPath.get("arguments");
259
- if (
260
- callee.isIdentifier() &&
261
- callee.node.name === "useMemoCache" &&
262
- args.length === 1 &&
263
- args[0].isNumericLiteral()
264
- ) {
265
- didInsertUseMemoCache = true;
266
- }
267
- },
268
- ImportDeclaration(importDeclPath) {
269
- if (isNonNamespacedImportOfReact(importDeclPath)) {
270
- hasExistingReactImport = true;
271
- }
272
- },
273
- });
258
+ const { didInsertUseMemoCache, hasExistingReactImport } =
259
+ findExistingImports(program);
260
+
261
// If Forget did successfully compile inject/update an import of
262
// `import {unstable_useMemoCache as useMemoCache} from 'react'` and rename
263
// `React.unstable_useMemoCache(n)` to `useMemoCache(n)`;
264
if (didInsertUseMemoCache) {
265
if (hasExistingReactImport) {
279
- let didUpdateImport = false;
280
- program.traverse({
281
- ImportDeclaration(importDeclPath) {
282
- if (isNonNamespacedImportOfReact(importDeclPath)) {
283
- importDeclPath.pushContainer(
284
- "specifiers",
285
- t.importSpecifier(
286
- t.identifier("useMemoCache"),
287
- t.identifier("unstable_useMemoCache")
288
- )
289
- );
290
- didUpdateImport = true;
291
- }
292
- },
293
- });
266
+ const didUpdateImport = updateExistingReactImportDeclaration(program);
267
if (didUpdateImport === false) {
268
throw new Error(
269
"Expected an ImportDeclaration of react in order to update ImportSpecifiers with useMemoCache"
270
);
271
}
272
} else {
300
- program.unshiftContainer(
301
- "body",
302
- t.importDeclaration(
303
- [
304
- t.importSpecifier(
305
- t.identifier("useMemoCache"),
306
- t.identifier("unstable_useMemoCache")
307
- ),
308
- ],
309
- t.stringLiteral("react")
310
- )
311
- );
273
+ insertUseMemoCacheImportDeclaration(program);
274
}
275
}
276
const externalFunctions = [];