Skip to content

Commit 6532c41

Browse files
committed
[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: d78d9ba Pull Request resolved: #30589
1 parent d50e024 commit 6532c41

File tree

8 files changed

+378
-17
lines changed

8 files changed

+378
-17
lines changed

compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts

Lines changed: 28 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,20 @@ export const InstrumentationSchema = z
6767

6868
export type ExternalFunction = z.infer<typeof ExternalFunctionSchema>;
6969

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+
7084
const HookSchema = z.object({
7185
/*
7286
* The effect of arguments to this hook. Describes whether the hook may or may
@@ -133,7 +147,7 @@ const EnvironmentConfigSchema = z.object({
133147
* plugin since it looks specifically for the name of the function being invoked, not
134148
* following aliases.
135149
*/
136-
customMacros: z.nullable(z.array(z.string())).default(null),
150+
customMacros: z.nullable(z.array(MacroSchema)).default(null),
137151

138152
/**
139153
* 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 {
490504
}
491505

492506
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+
}
494520
continue;
495521
}
496522

compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/MemoizeFbtAndMacroOperandsInSameScope.ts

Lines changed: 75 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ import {
1313
Place,
1414
ReactiveValue,
1515
} from '../HIR';
16+
import {Macro, MacroMethod} from '../HIR/Environment';
1617
import {eachReactiveValueOperand} from './visitors';
1718

1819
/**
@@ -43,15 +44,17 @@ import {eachReactiveValueOperand} from './visitors';
4344
export function memoizeFbtAndMacroOperandsInSameScope(
4445
fn: HIRFunction,
4546
): 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, []]),
4849
...(fn.env.config.customMacros ?? []),
4950
]);
5051
const fbtValues: Set<IdentifierId> = new Set();
52+
const macroMethods = new Map<IdentifierId, Array<Array<MacroMethod>>>();
5153
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) {
5558
break;
5659
}
5760
}
@@ -71,8 +74,9 @@ export const SINGLE_CHILD_FBT_TAGS: Set<string> = new Set([
7174

7275
function visit(
7376
fn: HIRFunction,
74-
fbtMacroTags: Set<string>,
77+
fbtMacroTags: Set<Macro>,
7578
fbtValues: Set<IdentifierId>,
79+
macroMethods: Map<IdentifierId, Array<Array<MacroMethod>>>,
7680
): void {
7781
for (const [, block] of fn.body.blocks) {
7882
for (const instruction of block.instructions) {
@@ -83,7 +87,7 @@ function visit(
8387
if (
8488
value.kind === 'Primitive' &&
8589
typeof value.value === 'string' &&
86-
fbtMacroTags.has(value.value)
90+
matchesExactTag(value.value, fbtMacroTags)
8791
) {
8892
/*
8993
* We don't distinguish between tag names and strings, so record
@@ -92,10 +96,38 @@ function visit(
9296
fbtValues.add(lvalue.identifier.id);
9397
} else if (
9498
value.kind === 'LoadGlobal' &&
95-
fbtMacroTags.has(value.binding.name)
99+
matchesExactTag(value.binding.name, fbtMacroTags)
96100
) {
97101
// Record references to `fbt` as a global
98102
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+
}
99131
} else if (isFbtCallExpression(fbtValues, value)) {
100132
const fbtScope = lvalue.identifier.scope;
101133
if (fbtScope === null) {
@@ -167,25 +199,57 @@ function visit(
167199
}
168200
}
169201

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+
170231
function isFbtCallExpression(
171232
fbtValues: Set<IdentifierId>,
172233
value: ReactiveValue,
173234
): boolean {
174235
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))
176239
);
177240
}
178241

179242
function isFbtJsxExpression(
180-
fbtMacroTags: Set<string>,
243+
fbtMacroTags: Set<Macro>,
181244
fbtValues: Set<IdentifierId>,
182245
value: ReactiveValue,
183246
): boolean {
184247
return (
185248
value.kind === 'JsxExpression' &&
186249
((value.tag.kind === 'Identifier' &&
187250
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)))
189253
);
190254
}
191255

Original file line numberDiff line numberDiff line change
@@ -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+
Original file line numberDiff line numberDiff line change
@@ -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+
}

0 commit comments

Comments
 (0)