@samitouri / QOS-React / commits / a07f5a3db5

[compiler] Provide support for custom fbt-like macro functions

ghstack-source-id: e3c6455ac2240914c3f25f3266a0cbb4a63971b5 Pull Request resolved: https://github.com/facebook/react/pull/29893

Joe Savona committed Jun 13, 2024 at 17:18 UTC a07f5a3db5deb5a429bf2617525b6e66dc777e8c
7 files changed +198 -12
compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts
+12
@@ -119,6 +119,18 @@ export type Hook = z.infer<typeof HookSchema>;
119 const EnvironmentConfigSchema = z.object({
120 customHooks: z.map(z.string(), HookSchema).optional().default(new Map()),
121
122 + /**
123 + * A list of functions which the application compiles as macros, where
124 + * the compiler must ensure they are not compiled to rename the macro or separate the
125 + * "function" from its argument.
126 + *
127 + * For example, Meta has some APIs such as `featureflag("name-of-feature-flag")` which
128 + * are rewritten by a plugin. Assigning `featureflag` to a temporary would break the
129 + * plugin since it looks specifically for the name of the function being invoked, not
130 + * following aliases.
131 + */
132 + customMacros: z.nullable(z.array(z.string())).default(null),
133 +
134 /**
135 * Enable a check that resets the memoization cache when the source code of the file changes.
136 * This is intended to support hot module reloading (HMR), where the same runtime component
compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/CodegenReactiveFunction.ts
+1 -1
@@ -43,7 +43,7 @@ import { Err, Ok, Result } from "../Utils/Result";
43 import { GuardKind } from "../Utils/RuntimeDiagnosticConstants";
44 import { assertExhaustive } from "../Utils/utils";
45 import { buildReactiveFunction } from "./BuildReactiveFunction";
46 -import { SINGLE_CHILD_FBT_TAGS } from "./MemoizeFbtOperandsInSameScope";
46 +import { SINGLE_CHILD_FBT_TAGS } from "./MemoizeFbtAndMacroOperandsInSameScope";
47 import { ReactiveFunctionVisitor, visitReactiveFunction } from "./visitors";
48
49 export const MEMO_CACHE_SENTINEL = "react.memo_cache_sentinel";
compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/MemoizeFbtAndMacroOperandsInSameScope.ts renamed
+31 -10
@@ -14,8 +14,15 @@ import {
14 } from "../HIR";
15 import { eachReactiveValueOperand } from "./visitors";
16
17 -/*
18 - * This pass supports the `fbt` translation system (https://facebook.github.io/fbt/).
17 +/**
18 + * This pass supports the
19 + * This pass supports the `fbt` translation system (https://facebook.github.io/fbt/)
20 + * as well as similar user-configurable macro-like APIs where it's important that
21 + * the name of the function not be changed, and it's literal arguments not be turned
22 + * into temporaries.
23 + *
24 + * ## FBT
25 + *
26 * FBT provides the `<fbt>` JSX element and `fbt()` calls (which take params in the
27 * form of `<fbt:param>` children or `fbt.param()` arguments, respectively). These
28 * tags/functions have restrictions on what types of syntax may appear as props/children/
@@ -26,13 +33,22 @@ import { eachReactiveValueOperand } from "./visitors";
33 * operands to fbt tags/calls have the same scope as the tag/call itself.
34 *
35 * Note that this still allows the props/arguments of `<fbt:param>`/`fbt.param()`
29 - * to be independently memoized
36 + * to be independently memoized.
37 + *
38 + * ## User-defined macro-like function
39 + *
40 + * Users can also specify their own functions to be treated similarly to fbt via the
41 + * `customMacros` environment configuration.
42 */
31 -export function memoizeFbtOperandsInSameScope(fn: HIRFunction): void {
43 +export function memoizeFbtAndMacroOperandsInSameScope(fn: HIRFunction): void {
44 + const fbtMacroTags = new Set([
45 + ...FBT_TAGS,
46 + ...(fn.env.config.customMacros ?? []),
47 + ]);
48 const fbtValues: Set<IdentifierId> = new Set();
49 while (true) {
50 let size = fbtValues.size;
35 - visit(fn, fbtValues);
51 + visit(fn, fbtMacroTags, fbtValues);
52 if (size === fbtValues.size) {
53 break;
54 }
@@ -50,7 +66,11 @@ export const SINGLE_CHILD_FBT_TAGS: Set<string> = new Set([
66 "fbs:param",
67 ]);
68
53 -function visit(fn: HIRFunction, fbtValues: Set<IdentifierId>): void {
69 +function visit(
70 + fn: HIRFunction,
71 + fbtMacroTags: Set<string>,
72 + fbtValues: Set<IdentifierId>
73 +): void {
74 for (const [, block] of fn.body.blocks) {
75 for (const instruction of block.instructions) {
76 const { lvalue, value } = instruction;
@@ -60,7 +80,7 @@ function visit(fn: HIRFunction, fbtValues: Set<IdentifierId>): void {
80 if (
81 value.kind === "Primitive" &&
82 typeof value.value === "string" &&
63 - FBT_TAGS.has(value.value)
83 + fbtMacroTags.has(value.value)
84 ) {
85 /*
86 * We don't distinguish between tag names and strings, so record
@@ -69,7 +89,7 @@ function visit(fn: HIRFunction, fbtValues: Set<IdentifierId>): void {
89 fbtValues.add(lvalue.identifier.id);
90 } else if (
91 value.kind === "LoadGlobal" &&
72 - FBT_TAGS.has(value.binding.name)
92 + fbtMacroTags.has(value.binding.name)
93 ) {
94 // Record references to `fbt` as a global
95 fbtValues.add(lvalue.identifier.id);
@@ -96,7 +116,7 @@ function visit(fn: HIRFunction, fbtValues: Set<IdentifierId>): void {
116 );
117 }
118 } else if (
99 - isFbtJsxExpression(fbtValues, value) ||
119 + isFbtJsxExpression(fbtMacroTags, fbtValues, value) ||
120 isFbtJsxChild(fbtValues, lvalue, value)
121 ) {
122 const fbtScope = lvalue.identifier.scope;
@@ -141,6 +161,7 @@ function isFbtCallExpression(
161 }
162
163 function isFbtJsxExpression(
164 + fbtMacroTags: Set<string>,
165 fbtValues: Set<IdentifierId>,
166 value: ReactiveValue
167 ): boolean {
@@ -148,7 +169,7 @@ function isFbtJsxExpression(
169 value.kind === "JsxExpression" &&
170 ((value.tag.kind === "Identifier" &&
171 fbtValues.has(value.tag.identifier.id)) ||
151 - (value.tag.kind === "BuiltinTag" && FBT_TAGS.has(value.tag.name)))
172 + (value.tag.kind === "BuiltinTag" && fbtMacroTags.has(value.tag.name)))
173 );
174 }
175
compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/index.ts
+1 -1
@@ -19,7 +19,7 @@ export { extractScopeDeclarationsFromDestructuring } from "./ExtractScopeDeclara
19 export { flattenReactiveLoops } from "./FlattenReactiveLoops";
20 export { flattenScopesWithHooksOrUse } from "./FlattenScopesWithHooksOrUse";
21 export { inferReactiveScopeVariables } from "./InferReactiveScopeVariables";
22 -export { memoizeFbtOperandsInSameScope } from "./MemoizeFbtOperandsInSameScope";
22 +export { memoizeFbtAndMacroOperandsInSameScope as memoizeFbtOperandsInSameScope } from "./MemoizeFbtAndMacroOperandsInSameScope";
23 export { mergeOverlappingReactiveScopes } from "./MergeOverlappingReactiveScopes";
24 export { mergeReactiveScopesThatInvalidateTogether } from "./MergeReactiveScopesThatInvalidateTogether";
25 export { printReactiveFunction } from "./PrintReactiveFunction";
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/meta-isms/repro-cx-assigned-to-temporary.expect.md new
+100
@@ -0,0 +1,100 @@
1 +
2 +## Input
3 +
4 +```javascript
5 +// @compilationMode(infer) @enableAssumeHooksFollowRulesOfReact:false @customMacros(cx)
6 +import { identity } from "shared-runtime";
7 +
8 +const DARK = "dark";
9 +
10 +function Component() {
11 + const theme = useTheme();
12 + return (
13 + <div
14 + className={cx({
15 + "styles/light": true,
16 + "styles/dark": theme.getTheme() === DARK,
17 + })}
18 + />
19 + );
20 +}
21 +
22 +function cx(obj) {
23 + const classes = [];
24 + for (const [key, value] of Object.entries(obj)) {
25 + if (value) {
26 + classes.push(key);
27 + }
28 + }
29 + return classes.join(" ");
30 +}
31 +
32 +function useTheme() {
33 + return {
34 + getTheme() {
35 + return DARK;
36 + },
37 + };
38 +}
39 +
40 +export const FIXTURE_ENTRYPOINT = {
41 + fn: Component,
42 + params: [{}],
43 +};
44 +
45 +```
46 +
47 +## Code
48 +
49 +```javascript
50 +import { c as _c } from "react/compiler-runtime"; // @compilationMode(infer) @enableAssumeHooksFollowRulesOfReact:false @customMacros(cx)
51 +import { identity } from "shared-runtime";
52 +
53 +const DARK = "dark";
54 +
55 +function Component() {
56 + const $ = _c(2);
57 + const theme = useTheme();
58 +
59 + const t0 = cx({
60 + "styles/light": true,
61 + "styles/dark": theme.getTheme() === DARK,
62 + });
63 + let t1;
64 + if ($[0] !== t0) {
65 + t1 = <div className={t0} />;
66 + $[0] = t0;
67 + $[1] = t1;
68 + } else {
69 + t1 = $[1];
70 + }
71 + return t1;
72 +}
73 +
74 +function cx(obj) {
75 + const classes = [];
76 + for (const [key, value] of Object.entries(obj)) {
77 + if (value) {
78 + classes.push(key);
79 + }
80 + }
81 + return classes.join(" ");
82 +}
83 +
84 +function useTheme() {
85 + return {
86 + getTheme() {
87 + return DARK;
88 + },
89 + };
90 +}
91 +
92 +export const FIXTURE_ENTRYPOINT = {
93 + fn: Component,
94 + params: [{}],
95 +};
96 +
97 +```
98 +
99 +### Eval output
100 +(kind: ok) <div class="styles/light styles/dark"></div>
\ No newline at end of file
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/meta-isms/repro-cx-assigned-to-temporary.js new
+39
@@ -0,0 +1,39 @@
1 +// @compilationMode(infer) @enableAssumeHooksFollowRulesOfReact:false @customMacros(cx)
2 +import { identity } from "shared-runtime";
3 +
4 +const DARK = "dark";
5 +
6 +function Component() {
7 + const theme = useTheme();
8 + return (
9 + <div
10 + className={cx({
11 + "styles/light": true,
12 + "styles/dark": theme.getTheme() === DARK,
13 + })}
14 + />
15 + );
16 +}
17 +
18 +function cx(obj) {
19 + const classes = [];
20 + for (const [key, value] of Object.entries(obj)) {
21 + if (value) {
22 + classes.push(key);
23 + }
24 + }
25 + return classes.join(" ");
26 +}
27 +
28 +function useTheme() {
29 + return {
30 + getTheme() {
31 + return DARK;
32 + },
33 + };
34 +}
35 +
36 +export const FIXTURE_ENTRYPOINT = {
37 + fn: Component,
38 + params: [{}],
39 +};
compiler/packages/snap/src/compiler.ts
+14
@@ -46,6 +46,7 @@ function makePluginOptions(
46 // TODO(@mofeiZ) rewrite snap fixtures to @validatePreserveExistingMemo:false
47 let validatePreserveExistingMemoizationGuarantees = false;
48 let enableChangeDetectionForDebugging = null;
49 + let customMacros = null;
50
51 if (firstLine.indexOf("@compilationMode(annotation)") !== -1) {
52 assert(
@@ -142,6 +143,18 @@ function makePluginOptions(
143 );
144 }
145
146 + const customMacrosMatch = /@customMacros\(([^)]+)\)/.exec(firstLine);
147 + if (
148 + customMacrosMatch &&
149 + customMacrosMatch.length > 1 &&
150 + customMacrosMatch[1].trim().length > 0
151 + ) {
152 + customMacros = customMacrosMatch[1]
153 + .split(" ")
154 + .map((s) => s.trim())
155 + .filter((s) => s.length > 0);
156 + }
157 +
158 let logs: Array<{ filename: string | null; event: LoggerEvent }> = [];
159 let logger: Logger | null = null;
160 if (firstLine.includes("@logger")) {
@@ -185,6 +198,7 @@ function makePluginOptions(
198 },
199 ],
200 ]),
201 + customMacros,
202 enableEmitFreeze,
203 enableEmitInstrumentForget,
204 enableEmitHookGuards,