Infer type of React function
Infer if a function is a component or hook when we're deciding to compile a function and store that in the environment. This is used in passes like InferReferenceEffects rather than having to re-parse the name in each pass.
Sathya Gunasekaran committed
Feb 29, 2024 at 14:47 UTC
ad8f19675cbde7af76f824a85c051bc57d03a173
6 files changed
+93
-47
compiler/apps/playground/components/Editor/EditorImpl.tsx
+39
-10
@@ -20,6 +20,7 @@ import {
20
run,
21
ValueKind,
22
} from "babel-plugin-react-forget";
23
+import { ReactFunctionType } from "babel-plugin-react-forget/dist/HIR/Environment";
24
import clsx from "clsx";
25
import invariant from "invariant";
26
import { useSnackbar } from "notistack";
@@ -142,6 +143,26 @@ const COMMON_HOOKS: Array<[string, Hook]> = [
143
],
144
];
145
146
+function isHookName(s: string): boolean {
147
+ return /^use[A-Z0-9]/.test(s);
148
+}
149
+
150
+function getReactFunctionType(
151
+ id: NodePath<t.Identifier | null | undefined>,
152
+): ReactFunctionType {
153
+ if (id && id.node && id.isIdentifier()) {
154
+ if (isHookName(id.node.name)) {
155
+ return "Hook";
156
+ }
157
+
158
+ const isPascalCaseNameSpace = /^[A-Z].*/;
159
+ if (isPascalCaseNameSpace.test(id.node.name)) {
160
+ return "Component";
161
+ }
162
+ }
163
+ return "Other";
164
+}
165
+
166
function compile(source: string): CompilerOutput {
167
const results = new Map<string, PrintedCompilerPipelineValue[]>();
168
const error = new CompilerError();
@@ -173,10 +194,16 @@ function compile(source: string): CompilerOutput {
194
continue;
195
}
196
176
- for (const result of run(fn, {
177
- ...config,
178
- customHooks: new Map([...COMMON_HOOKS]),
179
- }, null)) {
197
+ const id = fn.get("id");
198
+ for (const result of run(
199
+ fn,
200
+ {
201
+ ...config,
202
+ customHooks: new Map([...COMMON_HOOKS]),
203
+ },
204
+ getReactFunctionType(id),
205
+ null,
206
+ )) {
207
const fnName = fn.node.id?.name ?? null;
208
switch (result.kind) {
209
case "ast": {
@@ -238,12 +265,14 @@ function compile(source: string): CompilerOutput {
265
// Handle unexpected failures by logging (to get a stack trace)
266
// and reporting
267
console.error(err);
241
- error.details.push(new CompilerErrorDetail({
242
- severity: ErrorSeverity.Invariant,
243
- reason: `Unexpected failure when transforming input! ${err}`,
244
- loc: null,
245
- suggestions: null
246
- }));
268
+ error.details.push(
269
+ new CompilerErrorDetail({
270
+ severity: ErrorSeverity.Invariant,
271
+ reason: `Unexpected failure when transforming input! ${err}`,
272
+ loc: null,
273
+ suggestions: null,
274
+ }),
275
+ );
276
}
277
}
278
if (error.hasErrors()) {
compiler/packages/babel-plugin-react-forget/scripts/jest/makeTransform.ts
+1
-1
@@ -168,7 +168,7 @@ function ReactForgetFunctionTransform() {
168
}
169
}
170
171
- const compiled = compile(fn, forgetOptions, null);
171
+ const compiled = compile(fn, forgetOptions, "Other", null);
172
compiledFns.add(compiled);
173
174
const fun = t.functionDeclaration(
compiler/packages/babel-plugin-react-forget/src/Entrypoint/Pipeline.ts
+9
-3
@@ -17,7 +17,11 @@ import {
17
lower,
18
mergeConsecutiveBlocks,
19
} from "../HIR";
20
-import { Environment, EnvironmentConfig } from "../HIR/Environment";
20
+import {
21
+ Environment,
22
+ EnvironmentConfig,
23
+ ReactFunctionType,
24
+} from "../HIR/Environment";
25
import { findContextIdentifiers } from "../HIR/FindContextIdentifiers";
26
import {
27
analyseFunctions,
@@ -90,10 +94,11 @@ export function* run(
94
t.FunctionDeclaration | t.ArrowFunctionExpression | t.FunctionExpression
95
>,
96
config: EnvironmentConfig,
97
+ fnType: ReactFunctionType,
98
filename: string | null
99
): Generator<CompilerPipelineValue, CodegenFunction> {
100
const contextIdentifiers = findContextIdentifiers(func);
96
- const env = new Environment(config, contextIdentifiers);
101
+ const env = new Environment(fnType, config, contextIdentifiers);
102
yield {
103
kind: "debug",
104
name: "EnvironmentConfig",
@@ -380,9 +385,10 @@ export function compileFn(
385
t.FunctionDeclaration | t.ArrowFunctionExpression | t.FunctionExpression
386
>,
387
config: EnvironmentConfig,
388
+ fnType: ReactFunctionType,
389
filename: string | null
390
): CodegenFunction {
385
- let generator = run(func, config, filename);
391
+ let generator = run(func, config, fnType, filename);
392
while (true) {
393
const next = generator.next();
394
if (next.done) {
compiler/packages/babel-plugin-react-forget/src/Entrypoint/Program.ts
+38
-27
@@ -14,6 +14,7 @@ import {
14
} from "../CompilerError";
15
import {
16
ExternalFunction,
17
+ ReactFunctionType,
18
parseEnvironmentConfig,
19
tryParseExternalFunction,
20
} from "../HIR/Environment";
@@ -219,7 +220,8 @@ export function compileProgram(
220
const compiledFns: CompileResult[] = [];
221
222
const traverseFunction = (fn: BabelFn, pass: CompilerPass): void => {
222
- if (!shouldVisitNode(fn, pass) || ALREADY_COMPILED.has(fn.node)) {
223
+ const fnType = getReactFunctionType(fn, pass);
224
+ if (fnType === null || ALREADY_COMPILED.has(fn.node)) {
225
return;
226
}
227
@@ -262,7 +264,7 @@ export function compileProgram(
264
}
265
const config = environment.unwrap();
266
265
- compiledFn = compileFn(fn, config, pass.filename);
267
+ compiledFn = compileFn(fn, config, fnType, pass.filename);
268
pass.opts.logger?.logEvent(pass.filename, {
269
kind: "CompileSuccess",
270
fnLoc: fn.node.loc ?? null,
@@ -387,10 +389,14 @@ export function compileProgram(
389
}
390
}
391
390
-function shouldVisitNode(fn: BabelFn, pass: CompilerPass): boolean {
392
+function getReactFunctionType(
393
+ fn: BabelFn,
394
+ pass: CompilerPass
395
+): ReactFunctionType | null {
396
if (hasUseMemoCacheCall(fn)) {
392
- return false;
397
+ return null;
398
}
399
+ const hookPattern = pass.opts.environment?.hookPattern ?? null;
400
if (fn.node.body.type === "BlockStatement") {
401
// Opt-outs disable compilation regardless of mode
402
const useNoForget = findDirectiveDisablingMemoization(
@@ -407,30 +413,38 @@ function shouldVisitNode(fn: BabelFn, pass: CompilerPass): boolean {
413
suggestions: null,
414
},
415
});
410
- return false;
416
+ return null;
417
}
418
// Otherwise opt-ins enable compilation regardless of mode
419
if (findDirectiveEnablingMemoization(fn.node.body.directives) != null) {
414
- return true;
420
+ return getComponentOrHookLike(fn, hookPattern) ?? "Other";
421
}
422
}
423
switch (pass.opts.compilationMode) {
424
case "annotation": {
425
// opt-ins are checked above
420
- return false;
426
+ return null;
427
}
428
case "infer": {
423
- const hookPattern = pass.opts.environment?.hookPattern ?? null;
424
- return (
425
- // Component and hook declarations are known components/hooks
426
- (fn.isFunctionDeclaration() &&
427
- (isComponentDeclaration(fn.node) || isHookDeclaration(fn.node))) ||
428
- // Otherwise check if this is a component or hook-like function
429
- isComponentOrHookLike(fn, hookPattern)
430
- );
429
+ // Component and hook declarations are known components/hooks
430
+ if (fn.isFunctionDeclaration()) {
431
+ if (isComponentDeclaration(fn.node)) {
432
+ return "Component";
433
+ } else if (isHookDeclaration(fn.node)) {
434
+ return "Hook";
435
+ }
436
+ }
437
+
438
+ // Otherwise check if this is a component or hook-like function
439
+ return getComponentOrHookLike(fn, hookPattern);
440
}
441
case "all": {
433
- return fn.scope.getProgramParent() === fn.scope.parent;
442
+ // Compile only top level functions
443
+ if (fn.scope.getProgramParent() !== fn.scope.parent) {
444
+ return null;
445
+ }
446
+
447
+ return getComponentOrHookLike(fn, hookPattern) ?? "Other";
448
}
449
default: {
450
assertExhaustive(
@@ -567,23 +581,22 @@ function isValidComponentParams(
581
* Adapted from the ESLint rule at
582
* https://github.com/facebook/react/blob/main/packages/eslint-plugin-react-hooks/src/RulesOfHooks.js#L90-L103
583
*/
570
-function isComponentOrHookLike(
584
+function getComponentOrHookLike(
585
node: NodePath<
586
t.FunctionDeclaration | t.ArrowFunctionExpression | t.FunctionExpression
587
>,
588
hookPattern: string | null
575
-): boolean {
589
+): ReactFunctionType | null {
590
const functionName = getFunctionName(node);
591
// Check if the name is component or hook like:
592
if (functionName !== null && isComponentName(functionName)) {
579
- return (
580
- // As an added check we also look for hook invocations or JSX
593
+ let isComponent =
594
callsHooksOrCreatesJsx(node, hookPattern) &&
582
- isValidComponentParams(node.get("params"))
583
- );
595
+ isValidComponentParams(node.get("params"));
596
+ return isComponent ? "Component" : null;
597
} else if (functionName !== null && isHook(functionName, hookPattern)) {
598
// Hooks have hook invocations or JSX, but can take any # of arguments
586
- return callsHooksOrCreatesJsx(node, hookPattern);
599
+ return callsHooksOrCreatesJsx(node, hookPattern) ? "Hook" : null;
600
}
601
602
/*
@@ -593,12 +606,10 @@ function isComponentOrHookLike(
606
if (node.isFunctionExpression() || node.isArrowFunctionExpression()) {
607
if (isForwardRefCallback(node) || isMemoCallback(node)) {
608
// As an added check we also look for hook invocations or JSX
596
- return callsHooksOrCreatesJsx(node, hookPattern);
597
- } else {
598
- return false;
609
+ return callsHooksOrCreatesJsx(node, hookPattern) ? "Component" : null;
610
}
611
}
601
- return false;
612
+ return null;
613
}
614
615
function callsHooksOrCreatesJsx(
compiler/packages/babel-plugin-react-forget/src/HIR/Environment.ts
+5
@@ -415,6 +415,8 @@ export function parseConfigPragma(pragma: string): EnvironmentConfig {
415
416
export type PartialEnvironmentConfig = Partial<EnvironmentConfig>;
417
418
+export type ReactFunctionType = "Component" | "Hook" | "Other";
419
+
420
export class Environment {
421
#globals: GlobalRegistry;
422
#shapes: ShapeRegistry;
@@ -422,14 +424,17 @@ export class Environment {
424
#nextBlock: number = 0;
425
#nextScope: number = 0;
426
config: EnvironmentConfig;
427
+ fnType: ReactFunctionType;
428
429
#contextIdentifiers: Set<t.Identifier>;
430
#hoistedIdentifiers: Set<t.Identifier>;
431
432
constructor(
433
+ fnType: ReactFunctionType,
434
config: EnvironmentConfig,
435
contextIdentifiers: Set<t.Identifier>
436
) {
437
+ this.fnType = fnType;
438
this.config = config;
439
this.#shapes = new Map(DEFAULT_SHAPES);
440
this.#globals = new Map(DEFAULT_GLOBALS);
compiler/packages/babel-plugin-react-forget/src/Inference/InferReferenceEffects.ts
+1
-6
@@ -135,8 +135,7 @@ export default function inferReferenceEffects(
135
reason: new Set([ValueReason.ReactiveFunctionArgument]),
136
};
137
138
- const isComponent = isComponentName(fn.id);
139
- if (isComponent) {
138
+ if (fn.env.fnType === "Component") {
139
CompilerError.invariant(fn.params.length <= 2, {
140
reason:
141
"Expected React component to have not more than two parameters: one for props and for ref",
@@ -619,10 +618,6 @@ class InferenceState {
618
}
619
}
620
622
-function isComponentName(name: string | null): boolean {
623
- return name !== null && /^[A-Z]/.test(name);
624
-}
625
-
621
function inferParam(
622
param: Place | SpreadPattern,
623
initialState: InferenceState,