[compiler] Allow macro methods
Summary: Builds support for macros that are invoked as methods rather than just function calls or jsx. We now record macros as a schema that represents arbitrary member expressions including wildcards (so we can support, e.g., myMacro.*.foo.bar). When examining PropertyLoads in the macro memoization stage, we build up a map of partially-satisfied macro patterns until we determine that the pattern has been fully satisfied, at which point we treat the result of the PropertyLoad as a macro value. ghstack-source-id: d78d9ba7041968c861ffa110fb7882b339a0e257 Pull Request resolved: https://github.com/facebook/react/pull/30589
Mike Vitousek committed
Aug 12, 2024 at 12:55 UTC
6532c412566f7223f82b1460023180c25e15fca3
8 files changed
+378
-17
compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts
+28
-2
@@ -67,6 +67,20 @@ export const InstrumentationSchema = z
67
68
export type ExternalFunction = z.infer<typeof ExternalFunctionSchema>;
69
70
+export const MacroMethodSchema = z.union([
71
+ z.object({type: z.literal('wildcard')}),
72
+ z.object({type: z.literal('name'), name: z.string()}),
73
+]);
74
+
75
+// Would like to change this to drop the string option, but breaks compatibility with existing configs
76
+export const MacroSchema = z.union([
77
+ z.string(),
78
+ z.tuple([z.string(), z.array(MacroMethodSchema)]),
79
+]);
80
+
81
+export type Macro = z.infer<typeof MacroSchema>;
82
+export type MacroMethod = z.infer<typeof MacroMethodSchema>;
83
+
84
const HookSchema = z.object({
85
/*
86
* The effect of arguments to this hook. Describes whether the hook may or may
@@ -133,7 +147,7 @@ const EnvironmentConfigSchema = z.object({
147
* plugin since it looks specifically for the name of the function being invoked, not
148
* following aliases.
149
*/
136
- customMacros: z.nullable(z.array(z.string())).default(null),
150
+ customMacros: z.nullable(z.array(MacroSchema)).default(null),
151
152
/**
153
* Enable a check that resets the memoization cache when the source code of the file changes.
@@ -490,7 +504,19 @@ export function parseConfigPragma(pragma: string): EnvironmentConfig {
504
}
505
506
if (key === 'customMacros' && val) {
493
- maybeConfig[key] = [val];
507
+ const valSplit = val.split('.');
508
+ if (valSplit.length > 0) {
509
+ const props = [];
510
+ for (const elt of valSplit.slice(1)) {
511
+ if (elt === '*') {
512
+ props.push({type: 'wildcard'});
513
+ } else if (elt.length > 0) {
514
+ props.push({type: 'name', name: elt});
515
+ }
516
+ }
517
+ console.log([valSplit[0], props.map(x => x.name ?? '*').join('.')]);
518
+ maybeConfig[key] = [[valSplit[0], props]];
519
+ }
520
continue;
521
}
522
compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/MemoizeFbtAndMacroOperandsInSameScope.ts
+75
-11
@@ -13,6 +13,7 @@ import {
13
Place,
14
ReactiveValue,
15
} from '../HIR';
16
+import {Macro, MacroMethod} from '../HIR/Environment';
17
import {eachReactiveValueOperand} from './visitors';
18
19
/**
@@ -43,15 +44,17 @@ import {eachReactiveValueOperand} from './visitors';
44
export function memoizeFbtAndMacroOperandsInSameScope(
45
fn: HIRFunction,
46
): Set<IdentifierId> {
46
- const fbtMacroTags = new Set([
47
- ...FBT_TAGS,
47
+ const fbtMacroTags = new Set<Macro>([
48
+ ...Array.from(FBT_TAGS).map((tag): Macro => [tag, []]),
49
...(fn.env.config.customMacros ?? []),
50
]);
51
const fbtValues: Set<IdentifierId> = new Set();
52
+ const macroMethods = new Map<IdentifierId, Array<Array<MacroMethod>>>();
53
while (true) {
52
- let size = fbtValues.size;
53
- visit(fn, fbtMacroTags, fbtValues);
54
- if (size === fbtValues.size) {
54
+ let vsize = fbtValues.size;
55
+ let msize = macroMethods.size;
56
+ visit(fn, fbtMacroTags, fbtValues, macroMethods);
57
+ if (vsize === fbtValues.size && msize === macroMethods.size) {
58
break;
59
}
60
}
@@ -71,8 +74,9 @@ export const SINGLE_CHILD_FBT_TAGS: Set<string> = new Set([
74
75
function visit(
76
fn: HIRFunction,
74
- fbtMacroTags: Set<string>,
77
+ fbtMacroTags: Set<Macro>,
78
fbtValues: Set<IdentifierId>,
79
+ macroMethods: Map<IdentifierId, Array<Array<MacroMethod>>>,
80
): void {
81
for (const [, block] of fn.body.blocks) {
82
for (const instruction of block.instructions) {
@@ -83,7 +87,7 @@ function visit(
87
if (
88
value.kind === 'Primitive' &&
89
typeof value.value === 'string' &&
86
- fbtMacroTags.has(value.value)
90
+ matchesExactTag(value.value, fbtMacroTags)
91
) {
92
/*
93
* We don't distinguish between tag names and strings, so record
@@ -92,10 +96,38 @@ function visit(
96
fbtValues.add(lvalue.identifier.id);
97
} else if (
98
value.kind === 'LoadGlobal' &&
95
- fbtMacroTags.has(value.binding.name)
99
+ matchesExactTag(value.binding.name, fbtMacroTags)
100
) {
101
// Record references to `fbt` as a global
102
fbtValues.add(lvalue.identifier.id);
103
+ } else if (
104
+ value.kind === 'LoadGlobal' &&
105
+ matchTagRoot(value.binding.name, fbtMacroTags) !== null
106
+ ) {
107
+ const methods = matchTagRoot(value.binding.name, fbtMacroTags)!;
108
+ macroMethods.set(lvalue.identifier.id, methods);
109
+ } else if (
110
+ value.kind === 'PropertyLoad' &&
111
+ macroMethods.has(value.object.identifier.id)
112
+ ) {
113
+ const methods = macroMethods.get(value.object.identifier.id)!;
114
+ const newMethods = [];
115
+ for (const method of methods) {
116
+ if (
117
+ method.length > 0 &&
118
+ (method[0].type === 'wildcard' ||
119
+ (method[0].type === 'name' && method[0].name === value.property))
120
+ ) {
121
+ if (method.length > 1) {
122
+ newMethods.push(method.slice(1));
123
+ } else {
124
+ fbtValues.add(lvalue.identifier.id);
125
+ }
126
+ }
127
+ }
128
+ if (newMethods.length > 0) {
129
+ macroMethods.set(lvalue.identifier.id, newMethods);
130
+ }
131
} else if (isFbtCallExpression(fbtValues, value)) {
132
const fbtScope = lvalue.identifier.scope;
133
if (fbtScope === null) {
@@ -167,17 +199,48 @@ function visit(
199
}
200
}
201
202
+function matchesExactTag(s: string, tags: Set<Macro>): boolean {
203
+ return Array.from(tags).some(macro =>
204
+ typeof macro === 'string'
205
+ ? s === macro
206
+ : macro[1].length === 0 && macro[0] === s,
207
+ );
208
+}
209
+
210
+function matchTagRoot(
211
+ s: string,
212
+ tags: Set<Macro>,
213
+): Array<Array<MacroMethod>> | null {
214
+ const methods: Array<Array<MacroMethod>> = [];
215
+ for (const macro of tags) {
216
+ if (typeof macro === 'string') {
217
+ continue;
218
+ }
219
+ const [tag, rest] = macro;
220
+ if (tag === s && rest.length > 0) {
221
+ methods.push(rest);
222
+ }
223
+ }
224
+ if (methods.length > 0) {
225
+ return methods;
226
+ } else {
227
+ return null;
228
+ }
229
+}
230
+
231
function isFbtCallExpression(
232
fbtValues: Set<IdentifierId>,
233
value: ReactiveValue,
234
): boolean {
235
return (
175
- value.kind === 'CallExpression' && fbtValues.has(value.callee.identifier.id)
236
+ (value.kind === 'CallExpression' &&
237
+ fbtValues.has(value.callee.identifier.id)) ||
238
+ (value.kind === 'MethodCall' && fbtValues.has(value.property.identifier.id))
239
);
240
}
241
242
function isFbtJsxExpression(
180
- fbtMacroTags: Set<string>,
243
+ fbtMacroTags: Set<Macro>,
244
fbtValues: Set<IdentifierId>,
245
value: ReactiveValue,
246
): boolean {
@@ -185,7 +248,8 @@ function isFbtJsxExpression(
248
value.kind === 'JsxExpression' &&
249
((value.tag.kind === 'Identifier' &&
250
fbtValues.has(value.tag.identifier.id)) ||
188
- (value.tag.kind === 'BuiltinTag' && fbtMacroTags.has(value.tag.name)))
251
+ (value.tag.kind === 'BuiltinTag' &&
252
+ matchesExactTag(value.tag.name, fbtMacroTags)))
253
);
254
}
255
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/idx-method-no-outlining-wildcard.expect.md
new
+122
@@ -0,0 +1,122 @@
1
+
2
+## Input
3
+
4
+```javascript
5
+// @customMacros(idx.*.b)
6
+
7
+function Component(props) {
8
+ // outlined
9
+ const groupName1 = idx(props, _ => _.group.label);
10
+ // outlined
11
+ const groupName2 = idx.a(props, _ => _.group.label);
12
+ // not outlined
13
+ const groupName3 = idx.a.b(props, _ => _.group.label);
14
+ // not outlined
15
+ const groupName4 = idx.hello_world.b(props, _ => _.group.label);
16
+ // outlined
17
+ const groupName5 = idx.hello_world.b.c(props, _ => _.group.label);
18
+ return (
19
+ <div>
20
+ {groupName1}
21
+ {groupName2}
22
+ {groupName3}
23
+ {groupName4}
24
+ {groupName5}
25
+ </div>
26
+ );
27
+}
28
+
29
+```
30
+
31
+## Code
32
+
33
+```javascript
34
+import { c as _c } from "react/compiler-runtime"; // @customMacros(idx.*.b)
35
+
36
+function Component(props) {
37
+ const $ = _c(16);
38
+ let t0;
39
+ if ($[0] !== props) {
40
+ t0 = idx(props, _temp);
41
+ $[0] = props;
42
+ $[1] = t0;
43
+ } else {
44
+ t0 = $[1];
45
+ }
46
+ const groupName1 = t0;
47
+ let t1;
48
+ if ($[2] !== props) {
49
+ t1 = idx.a(props, _temp2);
50
+ $[2] = props;
51
+ $[3] = t1;
52
+ } else {
53
+ t1 = $[3];
54
+ }
55
+ const groupName2 = t1;
56
+ let t2;
57
+ if ($[4] !== props) {
58
+ t2 = idx.a.b(props, (__1) => __1.group.label);
59
+ $[4] = props;
60
+ $[5] = t2;
61
+ } else {
62
+ t2 = $[5];
63
+ }
64
+ const groupName3 = t2;
65
+ let t3;
66
+ if ($[6] !== props) {
67
+ t3 = idx.hello_world.b(props, (__2) => __2.group.label);
68
+ $[6] = props;
69
+ $[7] = t3;
70
+ } else {
71
+ t3 = $[7];
72
+ }
73
+ const groupName4 = t3;
74
+ let t4;
75
+ if ($[8] !== props) {
76
+ t4 = idx.hello_world.b.c(props, _temp3);
77
+ $[8] = props;
78
+ $[9] = t4;
79
+ } else {
80
+ t4 = $[9];
81
+ }
82
+ const groupName5 = t4;
83
+ let t5;
84
+ if (
85
+ $[10] !== groupName1 ||
86
+ $[11] !== groupName2 ||
87
+ $[12] !== groupName3 ||
88
+ $[13] !== groupName4 ||
89
+ $[14] !== groupName5
90
+ ) {
91
+ t5 = (
92
+ <div>
93
+ {groupName1}
94
+ {groupName2}
95
+ {groupName3}
96
+ {groupName4}
97
+ {groupName5}
98
+ </div>
99
+ );
100
+ $[10] = groupName1;
101
+ $[11] = groupName2;
102
+ $[12] = groupName3;
103
+ $[13] = groupName4;
104
+ $[14] = groupName5;
105
+ $[15] = t5;
106
+ } else {
107
+ t5 = $[15];
108
+ }
109
+ return t5;
110
+}
111
+function _temp3(__3) {
112
+ return __3.group.label;
113
+}
114
+function _temp2(__0) {
115
+ return __0.group.label;
116
+}
117
+function _temp(_) {
118
+ return _.group.label;
119
+}
120
+
121
+```
122
+
\ No newline at end of file
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/idx-method-no-outlining-wildcard.js
new
+23
@@ -0,0 +1,23 @@
1
+// @customMacros(idx.*.b)
2
+
3
+function Component(props) {
4
+ // outlined
5
+ const groupName1 = idx(props, _ => _.group.label);
6
+ // outlined
7
+ const groupName2 = idx.a(props, _ => _.group.label);
8
+ // not outlined
9
+ const groupName3 = idx.a.b(props, _ => _.group.label);
10
+ // not outlined
11
+ const groupName4 = idx.hello_world.b(props, _ => _.group.label);
12
+ // outlined
13
+ const groupName5 = idx.hello_world.b.c(props, _ => _.group.label);
14
+ return (
15
+ <div>
16
+ {groupName1}
17
+ {groupName2}
18
+ {groupName3}
19
+ {groupName4}
20
+ {groupName5}
21
+ </div>
22
+ );
23
+}
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/idx-method-no-outlining.expect.md
new
+85
@@ -0,0 +1,85 @@
1
+
2
+## Input
3
+
4
+```javascript
5
+// @customMacros(idx.a)
6
+
7
+function Component(props) {
8
+ // outlined
9
+ const groupName1 = idx(props, _ => _.group.label);
10
+ // not outlined
11
+ const groupName2 = idx.a(props, _ => _.group.label);
12
+ // outlined
13
+ const groupName3 = idx.a.b(props, _ => _.group.label);
14
+ return (
15
+ <div>
16
+ {groupName1}
17
+ {groupName2}
18
+ {groupName3}
19
+ </div>
20
+ );
21
+}
22
+
23
+```
24
+
25
+## Code
26
+
27
+```javascript
28
+import { c as _c } from "react/compiler-runtime"; // @customMacros(idx.a)
29
+
30
+function Component(props) {
31
+ const $ = _c(10);
32
+ let t0;
33
+ if ($[0] !== props) {
34
+ t0 = idx(props, _temp);
35
+ $[0] = props;
36
+ $[1] = t0;
37
+ } else {
38
+ t0 = $[1];
39
+ }
40
+ const groupName1 = t0;
41
+ let t1;
42
+ if ($[2] !== props) {
43
+ t1 = idx.a(props, (__0) => __0.group.label);
44
+ $[2] = props;
45
+ $[3] = t1;
46
+ } else {
47
+ t1 = $[3];
48
+ }
49
+ const groupName2 = t1;
50
+ let t2;
51
+ if ($[4] !== props) {
52
+ t2 = idx.a.b(props, _temp2);
53
+ $[4] = props;
54
+ $[5] = t2;
55
+ } else {
56
+ t2 = $[5];
57
+ }
58
+ const groupName3 = t2;
59
+ let t3;
60
+ if ($[6] !== groupName1 || $[7] !== groupName2 || $[8] !== groupName3) {
61
+ t3 = (
62
+ <div>
63
+ {groupName1}
64
+ {groupName2}
65
+ {groupName3}
66
+ </div>
67
+ );
68
+ $[6] = groupName1;
69
+ $[7] = groupName2;
70
+ $[8] = groupName3;
71
+ $[9] = t3;
72
+ } else {
73
+ t3 = $[9];
74
+ }
75
+ return t3;
76
+}
77
+function _temp2(__1) {
78
+ return __1.group.label;
79
+}
80
+function _temp(_) {
81
+ return _.group.label;
82
+}
83
+
84
+```
85
+
\ No newline at end of file
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/idx-method-no-outlining.js
new
+17
@@ -0,0 +1,17 @@
1
+// @customMacros(idx.a)
2
+
3
+function Component(props) {
4
+ // outlined
5
+ const groupName1 = idx(props, _ => _.group.label);
6
+ // not outlined
7
+ const groupName2 = idx.a(props, _ => _.group.label);
8
+ // outlined
9
+ const groupName3 = idx.a.b(props, _ => _.group.label);
10
+ return (
11
+ <div>
12
+ {groupName1}
13
+ {groupName2}
14
+ {groupName3}
15
+ </div>
16
+ );
17
+}
compiler/packages/snap/src/SproutTodoFilter.ts
+4
-1
@@ -493,13 +493,16 @@ const skipFilter = new Set([
493
494
// 'react-compiler-runtime' not yet supported
495
'flag-enable-emit-hook-guards',
496
-
496
'fast-refresh-refresh-on-const-changes-dev',
497
'useState-pruned-dependency-change-detect',
498
'useState-unpruned-dependency',
499
'useState-and-other-hook-unpruned-dependency',
500
'change-detect-reassign',
501
502
+ // Depends on external functions
503
+ 'idx-method-no-outlining-wildcard',
504
+ 'idx-method-no-outlining',
505
+
506
// needs to be executed as a module
507
'meta-property',
508
compiler/packages/snap/src/compiler.ts
+24
-3
@@ -20,7 +20,11 @@ import type {
20
PluginOptions,
21
} from 'babel-plugin-react-compiler/src/Entrypoint';
22
import type {Effect, ValueKind} from 'babel-plugin-react-compiler/src/HIR';
23
-import type {parseConfigPragma as ParseConfigPragma} from 'babel-plugin-react-compiler/src/HIR/Environment';
23
+import type {
24
+ Macro,
25
+ MacroMethod,
26
+ parseConfigPragma as ParseConfigPragma,
27
+} from 'babel-plugin-react-compiler/src/HIR/Environment';
28
import * as HermesParser from 'hermes-parser';
29
import invariant from 'invariant';
30
import path from 'path';
@@ -46,7 +50,7 @@ function makePluginOptions(
50
// TODO(@mofeiZ) rewrite snap fixtures to @validatePreserveExistingMemo:false
51
let validatePreserveExistingMemoizationGuarantees = false;
52
let enableChangeDetectionForDebugging = null;
49
- let customMacros = null;
53
+ let customMacros: null | Array<Macro> = null;
54
let validateBlocklistedImports = null;
55
56
if (firstLine.indexOf('@compilationMode(annotation)') !== -1) {
@@ -150,10 +154,27 @@ function makePluginOptions(
154
customMacrosMatch.length > 1 &&
155
customMacrosMatch[1].trim().length > 0
156
) {
153
- customMacros = customMacrosMatch[1]
157
+ const customMacrosStrs = customMacrosMatch[1]
158
.split(' ')
159
.map(s => s.trim())
160
.filter(s => s.length > 0);
161
+ if (customMacrosStrs.length > 0) {
162
+ customMacros = [];
163
+ for (const customMacroStr of customMacrosStrs) {
164
+ const props: Array<MacroMethod> = [];
165
+ const customMacroSplit = customMacroStr.split('.');
166
+ if (customMacroSplit.length > 0) {
167
+ for (const elt of customMacroSplit.slice(1)) {
168
+ if (elt === '*') {
169
+ props.push({type: 'wildcard'});
170
+ } else if (elt.length > 0) {
171
+ props.push({type: 'name', name: elt});
172
+ }
173
+ }
174
+ customMacros.push([customMacroSplit[0], props]);
175
+ }
176
+ }
177
+ }
178
}
179
180
const validateBlocklistedImportsMatch =