@samitouri / QOS-React / commits / f820f5a8b6

[compiler] Type inference for tagged template literals

At Meta we have a pattern of using tagged template literals for features that are compiled away: ``` // Relay: graphql`...graphql text...` ``` In many cases these tags produce a primitive value, and we can get even more optimal output if we can tell the compiler about these types. The new moduleTypeProvider gives us the ability to declare such types, this PR extends the compiler to use this type information for TaggedTemplateExpression values. ghstack-source-id: 3cd6511b7f4e708bcb86f3f3fde5773bc51c7197 Pull Request resolved: https://github.com/facebook/react/pull/30869

Joe Savona committed Sep 4, 2024 at 13:26 UTC f820f5a8b6c8b106ba3756f3e60a5a4017eb5080
9 files changed +271 -58
compiler/packages/babel-plugin-react-compiler/src/Inference/InferReferenceEffects.ts
+41 -12
@@ -1180,18 +1180,6 @@ function inferBlock(
1180 };
1181 break;
1182 }
1183 - case 'TaggedTemplateExpression': {
1184 - valueKind = {
1185 - kind: ValueKind.Mutable,
1186 - reason: new Set([ValueReason.Other]),
1187 - context: new Set(),
1188 - };
1189 - effect = {
1190 - kind: Effect.ConditionallyMutate,
1191 - reason: ValueReason.Other,
1192 - };
1193 - break;
1194 - }
1183 case 'TemplateLiteral': {
1184 /*
1185 * template literal (with no tag function) always produces
@@ -1312,6 +1300,47 @@ function inferBlock(
1300 instr.lvalue.effect = Effect.Store;
1301 continue;
1302 }
1303 + case 'TaggedTemplateExpression': {
1304 + const operands = [...eachInstructionValueOperand(instrValue)];
1305 + if (operands.length !== 1) {
1306 + // future-proofing to make sure we update this case when we support interpolation
1307 + CompilerError.throwTodo({
1308 + reason: 'Support tagged template expressions with interpolations',
1309 + loc: instrValue.loc,
1310 + });
1311 + }
1312 + const signature = getFunctionCallSignature(
1313 + env,
1314 + instrValue.tag.identifier.type,
1315 + );
1316 + let calleeEffect =
1317 + signature?.calleeEffect ?? Effect.ConditionallyMutate;
1318 + const returnValueKind: AbstractValue =
1319 + signature !== null
1320 + ? {
1321 + kind: signature.returnValueKind,
1322 + reason: new Set([
1323 + signature.returnValueReason ??
1324 + ValueReason.KnownReturnSignature,
1325 + ]),
1326 + context: new Set(),
1327 + }
1328 + : {
1329 + kind: ValueKind.Mutable,
1330 + reason: new Set([ValueReason.Other]),
1331 + context: new Set(),
1332 + };
1333 + state.referenceAndRecordEffects(
1334 + instrValue.tag,
1335 + calleeEffect,
1336 + ValueReason.Other,
1337 + functionEffects,
1338 + );
1339 + state.initialize(instrValue, returnValueKind);
1340 + state.define(instr.lvalue, instrValue);
1341 + instr.lvalue.effect = Effect.ConditionallyMutate;
1342 + continue;
1343 + }
1344 case 'CallExpression': {
1345 const signature = getFunctionCallSignature(
1346 env,
compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/InferReactiveScopeVariables.ts
+2 -2
@@ -227,6 +227,7 @@ function mayAllocate(env: Environment, instruction: Instruction): boolean {
227 case 'StoreGlobal': {
228 return false;
229 }
230 + case 'TaggedTemplateExpression':
231 case 'CallExpression':
232 case 'MethodCall': {
233 return instruction.lvalue.identifier.type.kind !== 'Primitive';
@@ -241,8 +242,7 @@ function mayAllocate(env: Environment, instruction: Instruction): boolean {
242 case 'ObjectExpression':
243 case 'UnsupportedNode':
244 case 'ObjectMethod':
244 - case 'FunctionExpression':
245 - case 'TaggedTemplateExpression': {
245 + case 'FunctionExpression': {
246 return true;
247 }
248 default: {
compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneNonEscapingScopes.ts
+28 -3
@@ -671,12 +671,37 @@ function computeMemoizationInputs(
671 ],
672 };
673 }
674 + case 'TaggedTemplateExpression': {
675 + const signature = getFunctionCallSignature(
676 + env,
677 + value.tag.identifier.type,
678 + );
679 + let lvalues = [];
680 + if (lvalue !== null) {
681 + lvalues.push({place: lvalue, level: MemoizationLevel.Memoized});
682 + }
683 + if (signature?.noAlias === true) {
684 + return {
685 + lvalues,
686 + rvalues: [],
687 + };
688 + }
689 + const operands = [...eachReactiveValueOperand(value)];
690 + lvalues.push(
691 + ...operands
692 + .filter(operand => isMutableEffect(operand.effect, operand.loc))
693 + .map(place => ({place, level: MemoizationLevel.Memoized})),
694 + );
695 + return {
696 + lvalues,
697 + rvalues: operands,
698 + };
699 + }
700 case 'CallExpression': {
701 const signature = getFunctionCallSignature(
702 env,
703 value.callee.identifier.type,
704 );
679 - const operands = [...eachReactiveValueOperand(value)];
705 let lvalues = [];
706 if (lvalue !== null) {
707 lvalues.push({place: lvalue, level: MemoizationLevel.Memoized});
@@ -687,6 +712,7 @@ function computeMemoizationInputs(
712 rvalues: [],
713 };
714 }
715 + const operands = [...eachReactiveValueOperand(value)];
716 lvalues.push(
717 ...operands
718 .filter(operand => isMutableEffect(operand.effect, operand.loc))
@@ -702,7 +728,6 @@ function computeMemoizationInputs(
728 env,
729 value.property.identifier.type,
730 );
705 - const operands = [...eachReactiveValueOperand(value)];
731 let lvalues = [];
732 if (lvalue !== null) {
733 lvalues.push({place: lvalue, level: MemoizationLevel.Memoized});
@@ -713,6 +738,7 @@ function computeMemoizationInputs(
738 rvalues: [],
739 };
740 }
741 + const operands = [...eachReactiveValueOperand(value)];
742 lvalues.push(
743 ...operands
744 .filter(operand => isMutableEffect(operand.effect, operand.loc))
@@ -726,7 +752,6 @@ function computeMemoizationInputs(
752 case 'RegExpLiteral':
753 case 'ObjectMethod':
754 case 'FunctionExpression':
729 - case 'TaggedTemplateExpression':
755 case 'ArrayExpression':
756 case 'NewExpression':
757 case 'ObjectExpression':
compiler/packages/babel-plugin-react-compiler/src/TypeInference/InferTypes.ts
+19 -2
@@ -250,6 +250,7 @@ function* generateInstructionTypes(
250 }
251
252 case 'CallExpression': {
253 + const returnType = makeType();
254 /*
255 * TODO: callee could be a hook or a function, so this type equation isn't correct.
256 * We should change Hook to a subtype of Function or change unifier logic.
@@ -258,8 +259,25 @@ function* generateInstructionTypes(
259 yield equation(value.callee.identifier.type, {
260 kind: 'Function',
261 shapeId: null,
261 - return: left,
262 + return: returnType,
263 });
264 + yield equation(left, returnType);
265 + break;
266 + }
267 +
268 + case 'TaggedTemplateExpression': {
269 + const returnType = makeType();
270 + /*
271 + * TODO: callee could be a hook or a function, so this type equation isn't correct.
272 + * We should change Hook to a subtype of Function or change unifier logic.
273 + * (see https://github.com/facebook/react-forget/pull/1427)
274 + */
275 + yield equation(value.tag.identifier.type, {
276 + kind: 'Function',
277 + shapeId: null,
278 + return: returnType,
279 + });
280 + yield equation(left, returnType);
281 break;
282 }
283
@@ -392,7 +410,6 @@ function* generateInstructionTypes(
410 case 'MetaProperty':
411 case 'ComputedStore':
412 case 'ComputedLoad':
395 - case 'TaggedTemplateExpression':
413 case 'Await':
414 case 'GetIterator':
415 case 'IteratorNext':
compiler/packages/babel-plugin-react-compiler/src/Validation/ValidateLocalsNotReassignedAfterRender.ts
+8
@@ -161,6 +161,14 @@ function getContextReassignment(
161 if (signature?.noAlias) {
162 operands = [value.receiver, value.property];
163 }
164 + } else if (value.kind === 'TaggedTemplateExpression') {
165 + const signature = getFunctionCallSignature(
166 + fn.env,
167 + value.tag.identifier.type,
168 + );
169 + if (signature?.noAlias) {
170 + operands = [value.tag];
171 + }
172 }
173 for (const operand of operands) {
174 CompilerError.invariant(operand.effect !== Effect.Unknown, {
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/destructuring-mixed-scope-and-local-variables-with-default.expect.md
+35 -39
@@ -63,67 +63,63 @@ function useFragment(_arg1, _arg2) {
63 }
64
65 function Component(props) {
66 - const $ = _c(9);
67 - let t0;
68 - if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
69 - t0 = graphql`
66 + const $ = _c(8);
67 + const post = useFragment(
68 + graphql`
69 fragment F on T {
70 id
71 }
73 - `;
74 - $[0] = t0;
75 - } else {
76 - t0 = $[0];
77 - }
78 - const post = useFragment(t0, props.post);
79 - let t1;
80 - if ($[1] !== post) {
72 + `,
73 + props.post,
74 + );
75 + let t0;
76 + if ($[0] !== post) {
77 const allUrls = [];
78
83 - const { media: t2, comments: t3, urls: t4 } = post;
84 - const media = t2 === undefined ? null : t2;
79 + const { media: t1, comments: t2, urls: t3 } = post;
80 + const media = t1 === undefined ? null : t1;
81 + let t4;
82 + if ($[2] !== t2) {
83 + t4 = t2 === undefined ? [] : t2;
84 + $[2] = t2;
85 + $[3] = t4;
86 + } else {
87 + t4 = $[3];
88 + }
89 + const comments = t4;
90 let t5;
86 - if ($[3] !== t3) {
91 + if ($[4] !== t3) {
92 t5 = t3 === undefined ? [] : t3;
88 - $[3] = t3;
89 - $[4] = t5;
93 + $[4] = t3;
94 + $[5] = t5;
95 } else {
91 - t5 = $[4];
96 + t5 = $[5];
97 }
93 - const comments = t5;
98 + const urls = t5;
99 let t6;
95 - if ($[5] !== t4) {
96 - t6 = t4 === undefined ? [] : t4;
97 - $[5] = t4;
98 - $[6] = t6;
99 - } else {
100 - t6 = $[6];
101 - }
102 - const urls = t6;
103 - let t7;
104 - if ($[7] !== comments.length) {
105 - t7 = (e) => {
100 + if ($[6] !== comments.length) {
101 + t6 = (e) => {
102 if (!comments.length) {
103 return;
104 }
105
106 console.log(comments.length);
107 };
112 - $[7] = comments.length;
113 - $[8] = t7;
108 + $[6] = comments.length;
109 + $[7] = t6;
110 } else {
115 - t7 = $[8];
111 + t6 = $[7];
112 }
117 - const onClick = t7;
113 + const onClick = t6;
114
115 allUrls.push(...urls);
120 - t1 = <Stringify media={media} allUrls={allUrls} onClick={onClick} />;
121 - $[1] = post;
122 - $[2] = t1;
116 + t0 = <Stringify media={media} allUrls={allUrls} onClick={onClick} />;
117 + $[0] = post;
118 + $[1] = t0;
119 } else {
124 - t1 = $[2];
120 + t0 = $[1];
121 }
126 - return t1;
122 + return t0;
123 }
124
125 export const FIXTURE_ENTRYPOINT = {
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/type-provider-tagged-template-expression.expect.md new
+106
@@ -0,0 +1,106 @@
1 +
2 +## Input
3 +
4 +```javascript
5 +import {graphql} from 'shared-runtime';
6 +
7 +export function Component({a, b}) {
8 + const fragment = graphql`
9 + fragment Foo on User {
10 + name
11 + }
12 + `;
13 + return <div>{fragment}</div>;
14 +}
15 +
16 +export const FIXTURE_ENTRYPOINT = {
17 + fn: Component,
18 + params: [{a: 0, b: 0}],
19 + sequentialRenders: [
20 + {a: 0, b: 0},
21 + {a: 1, b: 0},
22 + {a: 1, b: 1},
23 + {a: 1, b: 2},
24 + {a: 2, b: 2},
25 + {a: 3, b: 2},
26 + {a: 0, b: 0},
27 + ],
28 +};
29 +
30 +```
31 +
32 +## Code
33 +
34 +```javascript
35 +import { c as _c } from "react/compiler-runtime";
36 +import { graphql } from "shared-runtime";
37 +
38 +export function Component(t0) {
39 + const $ = _c(1);
40 + const fragment = graphql`
41 + fragment Foo on User {
42 + name
43 + }
44 + `;
45 + let t1;
46 + if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
47 + t1 = <div>{fragment}</div>;
48 + $[0] = t1;
49 + } else {
50 + t1 = $[0];
51 + }
52 + return t1;
53 +}
54 +
55 +export const FIXTURE_ENTRYPOINT = {
56 + fn: Component,
57 + params: [{ a: 0, b: 0 }],
58 + sequentialRenders: [
59 + { a: 0, b: 0 },
60 + { a: 1, b: 0 },
61 + { a: 1, b: 1 },
62 + { a: 1, b: 2 },
63 + { a: 2, b: 2 },
64 + { a: 3, b: 2 },
65 + { a: 0, b: 0 },
66 + ],
67 +};
68 +
69 +```
70 +
71 +### Eval output
72 +(kind: ok) <div>
73 + fragment Foo on User {
74 + name
75 + }
76 + </div>
77 +<div>
78 + fragment Foo on User {
79 + name
80 + }
81 + </div>
82 +<div>
83 + fragment Foo on User {
84 + name
85 + }
86 + </div>
87 +<div>
88 + fragment Foo on User {
89 + name
90 + }
91 + </div>
92 +<div>
93 + fragment Foo on User {
94 + name
95 + }
96 + </div>
97 +<div>
98 + fragment Foo on User {
99 + name
100 + }
101 + </div>
102 +<div>
103 + fragment Foo on User {
104 + name
105 + }
106 + </div>
\ No newline at end of file
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/type-provider-tagged-template-expression.js new
+24
@@ -0,0 +1,24 @@
1 +import {graphql} from 'shared-runtime';
2 +
3 +export function Component({a, b}) {
4 + const fragment = graphql`
5 + fragment Foo on User {
6 + name
7 + }
8 + `;
9 + return <div>{fragment}</div>;
10 +}
11 +
12 +export const FIXTURE_ENTRYPOINT = {
13 + fn: Component,
14 + params: [{a: 0, b: 0}],
15 + sequentialRenders: [
16 + {a: 0, b: 0},
17 + {a: 1, b: 0},
18 + {a: 1, b: 1},
19 + {a: 1, b: 2},
20 + {a: 2, b: 2},
21 + {a: 3, b: 2},
22 + {a: 0, b: 0},
23 + ],
24 +};
compiler/packages/snap/src/sprout/shared-runtime-type-provider.ts
+8
@@ -32,6 +32,14 @@ export function makeSharedRuntimeTypeProvider({
32 returnType: {kind: 'type', name: 'Primitive'},
33 returnValueKind: ValueKindEnum.Primitive,
34 },
35 + graphql: {
36 + kind: 'function',
37 + calleeEffect: EffectEnum.Read,
38 + positionalParams: [],
39 + restParam: EffectEnum.Read,
40 + returnType: {kind: 'type', name: 'Primitive'},
41 + returnValueKind: ValueKindEnum.Primitive,
42 + },
43 typedArrayPush: {
44 kind: 'function',
45 calleeEffect: EffectEnum.Read,