Allow prefixed hooks for compiling bundled code
We're doing some internal benchmarking using a lightweight bundler that @pieterv wrote for experimentation purposes. It's designed to fully preserve Flow type annotations so we can experiment with type-driven compilation and test out what benefits we might get from "cross-module" compilation more easily (ie by just bundling together a few modules so we can see them all as one). However, the bundler renames local variables and imports, so that a reference to `useMemo()` might end up as `React$useMemo()` or similar. This PR adds a flag to tell the compiler that builtin hooks might be prefixed and resolve them appropriately.
Joe Savona committed
Jan 30, 2024 at 22:11 UTC
d55420c430cee6338fcc63ad663cb4b778af55de
7 files changed
+181
-24
compiler/packages/babel-plugin-react-forget/src/Entrypoint/Program.ts
+24
-12
@@ -376,11 +376,12 @@ function shouldVisitNode(fn: BabelFn, pass: CompilerPass): boolean {
376
return false;
377
}
378
case "infer": {
379
+ const hookPattern = pass.opts.environment?.hookPattern ?? null;
380
return (
381
// Component declarations are known components
382
(fn.isFunctionDeclaration() && isComponentDeclaration(fn.node)) ||
383
// Otherwise check if this is a component or hook-like function
383
- isComponentOrHookLike(fn)
384
+ isComponentOrHookLike(fn, hookPattern)
385
);
386
}
387
case "all": {
@@ -414,7 +415,10 @@ function hasUseMemoCacheCall(
415
return hasUseMemoCache;
416
}
417
417
-function isHookName(s: string): boolean {
418
+function isHookName(s: string, hookPattern: string | null): boolean {
419
+ if (hookPattern !== null) {
420
+ return new RegExp(hookPattern).test(s);
421
+ }
422
return /^use[A-Z0-9]/.test(s);
423
}
424
@@ -423,13 +427,16 @@ function isHookName(s: string): boolean {
427
* containing a hook name.
428
*/
429
426
-function isHook(path: NodePath<t.Expression | t.PrivateName>): boolean {
430
+function isHook(
431
+ path: NodePath<t.Expression | t.PrivateName>,
432
+ hookPattern: string | null
433
+): boolean {
434
if (path.isIdentifier()) {
428
- return isHookName(path.node.name);
435
+ return isHookName(path.node.name, hookPattern);
436
} else if (
437
path.isMemberExpression() &&
438
!path.node.computed &&
432
- isHook(path.get("property"))
439
+ isHook(path.get("property"), hookPattern)
440
) {
441
const obj = path.get("object").node;
442
const isPascalCaseNameSpace = /^[A-Z].*/;
@@ -518,18 +525,20 @@ function isValidComponentParams(
525
function isComponentOrHookLike(
526
node: NodePath<
527
t.FunctionDeclaration | t.ArrowFunctionExpression | t.FunctionExpression
521
- >
528
+ >,
529
+ hookPattern: string | null
530
): boolean {
531
const functionName = getFunctionName(node);
532
// Check if the name is component or hook like:
533
if (functionName !== null && isComponentName(functionName)) {
534
return (
535
// As an added check we also look for hook invocations or JSX
528
- callsHooksOrCreatesJsx(node) && isValidComponentParams(node.get("params"))
536
+ callsHooksOrCreatesJsx(node, hookPattern) &&
537
+ isValidComponentParams(node.get("params"))
538
);
530
- } else if (functionName !== null && isHook(functionName)) {
539
+ } else if (functionName !== null && isHook(functionName, hookPattern)) {
540
// Hooks have hook invocations or JSX, but can take any # of arguments
532
- return callsHooksOrCreatesJsx(node);
541
+ return callsHooksOrCreatesJsx(node, hookPattern);
542
}
543
544
/*
@@ -539,7 +548,7 @@ function isComponentOrHookLike(
548
if (node.isFunctionExpression() || node.isArrowFunctionExpression()) {
549
if (isForwardRefCallback(node) || isMemoCallback(node)) {
550
// As an added check we also look for hook invocations or JSX
542
- return callsHooksOrCreatesJsx(node);
551
+ return callsHooksOrCreatesJsx(node, hookPattern);
552
} else {
553
return false;
554
}
@@ -547,7 +556,10 @@ function isComponentOrHookLike(
556
return false;
557
}
558
550
-function callsHooksOrCreatesJsx(node: NodePath<t.Node>): boolean {
559
+function callsHooksOrCreatesJsx(
560
+ node: NodePath<t.Node>,
561
+ hookPattern: string | null
562
+): boolean {
563
let invokesHooks = false;
564
let createsJsx = false;
565
node.traverse({
@@ -556,7 +568,7 @@ function callsHooksOrCreatesJsx(node: NodePath<t.Node>): boolean {
568
},
569
CallExpression(call) {
570
const callee = call.get("callee");
559
- if (callee.isExpression() && isHook(callee)) {
571
+ if (callee.isExpression() && isHook(callee, hookPattern)) {
572
invokesHooks = true;
573
}
574
},
compiler/packages/babel-plugin-react-forget/src/HIR/Environment.ts
+27
-1
@@ -344,6 +344,19 @@ const EnvironmentConfigSchema = z.object({
344
* non-ideal.
345
*/
346
enableTreatFunctionDepsAsConditional: z.boolean().default(false),
347
+
348
+ /**
349
+ * If specified, this value is used as a pattern for determing which global values should be
350
+ * treated as hooks. The pattern should have a single capture group, which will be used as
351
+ * the hook name for the purposes of resolving hook definitions (for builtin hooks)_.
352
+ *
353
+ * For example, by default `React$useState` would not be treated as a hook. By specifying
354
+ * `hookPattern: 'React$(\w+)'`, the compiler will treat this value equivalently to `useState()`.
355
+ *
356
+ * This setting is intended for cases where Forget is compiling code that has been prebundled
357
+ * and identifiers have been changed.
358
+ */
359
+ hookPattern: z.string().nullable().default(null),
360
});
361
362
export type EnvironmentConfig = z.infer<typeof EnvironmentConfigSchema>;
@@ -447,7 +460,20 @@ export class Environment {
460
}
461
462
getGlobalDeclaration(name: string): Global | null {
450
- let resolvedGlobal: Global | null = this.#globals.get(name) ?? null;
463
+ let resolvedName = name;
464
+
465
+ if (this.config.hookPattern != null) {
466
+ const match = new RegExp(this.config.hookPattern).exec(name);
467
+ if (
468
+ match != null &&
469
+ typeof match[1] === "string" &&
470
+ isHookName(match[1])
471
+ ) {
472
+ resolvedName = match[1];
473
+ }
474
+ }
475
+
476
+ let resolvedGlobal: Global | null = this.#globals.get(resolvedName) ?? null;
477
if (resolvedGlobal === null) {
478
// Hack, since we don't track module level declarations and imports
479
if (isHookName(name)) {
compiler/packages/babel-plugin-react-forget/src/HIR/HIR.ts
+9
-3
@@ -1233,9 +1233,15 @@ export function isUseInsertionEffectHookType(id: Identifier): boolean {
1233
}
1234
1235
export function getHookKind(env: Environment, id: Identifier): HookKind | null {
1236
- const idType = id.type;
1237
- if (idType.kind === "Function") {
1238
- const signature = env.getFunctionSignature(idType);
1236
+ return getHookKindForType(env, id.type);
1237
+}
1238
+
1239
+export function getHookKindForType(
1240
+ env: Environment,
1241
+ type: Type
1242
+): HookKind | null {
1243
+ if (type.kind === "Function") {
1244
+ const signature = env.getFunctionSignature(type);
1245
return signature?.hookKind ?? null;
1246
}
1247
return null;
compiler/packages/babel-plugin-react-forget/src/Inference/DropManualMemoization.ts
+6
-5
@@ -14,6 +14,7 @@ import {
14
Instruction,
15
Place,
16
SpreadPattern,
17
+ getHookKindForType,
18
makeInstructionId,
19
} from "../HIR";
20
import { createTemporaryPlace, markInstructionIds } from "../HIR/HIRBuilder";
@@ -43,11 +44,11 @@ export function dropManualMemoization(func: HIRFunction): void {
44
break;
45
}
46
case "LoadGlobal": {
46
- if (
47
- instr.value.name === "useMemo" ||
48
- instr.value.name === "useCallback"
49
- ) {
50
- hooks.set(instr.lvalue.identifier.id, instr.value.name);
47
+ const global = func.env.getGlobalDeclaration(instr.value.name);
48
+ const hookKind =
49
+ global !== null ? getHookKindForType(func.env, global) : null;
50
+ if (hookKind === "useMemo" || hookKind === "useCallback") {
51
+ hooks.set(instr.lvalue.identifier.id, hookKind);
52
} else if (instr.value.name === "React") {
53
react.add(instr.lvalue.identifier.id);
54
}
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/hooks-with-prefix.expect.md
new
+75
@@ -0,0 +1,75 @@
1
+
2
+## Input
3
+
4
+```javascript
5
+// @hookPattern:"React\$(\w+)"
6
+
7
+import * as React from "react";
8
+import { makeArray } from "shared-runtime";
9
+
10
+const React$useState = React.useState;
11
+const React$useMemo = React.useMemo;
12
+
13
+function Component() {
14
+ const [state, setState] = React$useState(0);
15
+ const doubledArray = React$useMemo(() => {
16
+ return makeArray(state);
17
+ }, [state]);
18
+ return <div>{doubledArray.join("")}</div>;
19
+}
20
+
21
+export const FIXTURE_ENTRYPOINT = {
22
+ fn: Component,
23
+ params: [{}],
24
+};
25
+
26
+```
27
+
28
+## Code
29
+
30
+```javascript
31
+import { unstable_useMemoCache as useMemoCache } from "react"; // @hookPattern:"React\$(\w+)"
32
+
33
+import * as React from "react";
34
+import { makeArray } from "shared-runtime";
35
+
36
+const React$useState = React.useState;
37
+const React$useMemo = React.useMemo;
38
+
39
+function Component() {
40
+ const $ = useMemoCache(5);
41
+ const [state] = React$useState(0);
42
+ let t15;
43
+ let t0;
44
+ if ($[0] !== state) {
45
+ t15 = makeArray(state);
46
+ const doubledArray = t15;
47
+
48
+ t0 = doubledArray.join("");
49
+ $[0] = state;
50
+ $[1] = t0;
51
+ $[2] = t15;
52
+ } else {
53
+ t0 = $[1];
54
+ t15 = $[2];
55
+ }
56
+ let t1;
57
+ if ($[3] !== t0) {
58
+ t1 = <div>{t0}</div>;
59
+ $[3] = t0;
60
+ $[4] = t1;
61
+ } else {
62
+ t1 = $[4];
63
+ }
64
+ return t1;
65
+}
66
+
67
+export const FIXTURE_ENTRYPOINT = {
68
+ fn: Component,
69
+ params: [{}],
70
+};
71
+
72
+```
73
+
74
+### Eval output
75
+(kind: ok) <div>0</div>
\ No newline at end of file
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/hooks-with-prefix.js
new
+20
@@ -0,0 +1,20 @@
1
+// @hookPattern:"React\$(\w+)"
2
+
3
+import * as React from "react";
4
+import { makeArray } from "shared-runtime";
5
+
6
+const React$useState = React.useState;
7
+const React$useMemo = React.useMemo;
8
+
9
+function Component() {
10
+ const [state, setState] = React$useState(0);
11
+ const doubledArray = React$useMemo(() => {
12
+ return makeArray(state);
13
+ }, [state]);
14
+ return <div>{doubledArray.join("")}</div>;
15
+}
16
+
17
+export const FIXTURE_ENTRYPOINT = {
18
+ fn: Component,
19
+ params: [{}],
20
+};
compiler/packages/fixture-test-utils/src/compiler-utils.ts
+20
-3
@@ -37,6 +37,7 @@ export function transformFixtureInput(
37
let compilationMode: CompilationMode = "all";
38
let enableUseMemoCachePolyfill = false;
39
let panicThreshold: PanicThresholdOptions = "ALL_ERRORS";
40
+ let hookPattern: string | null = null;
41
42
if (firstLine.indexOf("@compilationMode(annotation)") !== -1) {
43
assert(
@@ -85,9 +86,24 @@ export function transformFixtureInput(
86
}
87
88
let eslintSuppressionRules: Array<string> | null = null;
88
- const match = /@eslintSuppressionRules\(([^)]+)\)/.exec(firstLine);
89
- if (match != null) {
90
- eslintSuppressionRules = match[1].split("|");
89
+ const eslintSuppressionMatch = /@eslintSuppressionRules\(([^)]+)\)/.exec(
90
+ firstLine
91
+ );
92
+ if (eslintSuppressionMatch != null) {
93
+ eslintSuppressionRules = eslintSuppressionMatch[1].split("|");
94
+ }
95
+
96
+ const hookPatternMatch = /@hookPattern:"([^"]+)"/.exec(firstLine);
97
+ if (
98
+ hookPatternMatch &&
99
+ hookPatternMatch.length > 1 &&
100
+ hookPatternMatch[1].trim().length > 0
101
+ ) {
102
+ hookPattern = hookPatternMatch[1].trim();
103
+ } else if (firstLine.includes("@hookPattern")) {
104
+ throw new Error(
105
+ 'Invalid @hookPattern:"..." pragma, must contain the prefix between balanced double quotes eg @hookPattern:"pattern"'
106
+ );
107
}
108
109
const config = parseConfigPragmaFn(firstLine);
@@ -131,6 +147,7 @@ export function transformFixtureInput(
147
enableEmitInstrumentForget,
148
enableEmitHookGuards,
149
assertValidMutableRanges: true,
150
+ hookPattern,
151
},
152
compilationMode,
153
logger: null,