From bea8495ef6bf5f4a474bb57eb0b99d324a3b3063 Mon Sep 17 00:00:00 2001 From: Joe Savona Date: Tue, 3 Sep 2024 18:44:35 -0700 Subject: [PATCH] [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-poisoned] --- .../src/Inference/InferReferenceEffects.ts | 53 +++++++-- .../InferReactiveScopeVariables.ts | 4 +- .../ReactiveScopes/PruneNonEscapingScopes.ts | 31 ++++- .../src/TypeInference/InferTypes.ts | 21 +++- .../ValidateLocalsNotReassignedAfterRender.ts | 8 ++ ...and-local-variables-with-default.expect.md | 74 ++++++------ ...vider-tagged-template-expression.expect.md | 106 ++++++++++++++++++ ...ype-provider-tagged-template-expression.js | 24 ++++ .../sprout/shared-runtime-type-provider.ts | 8 ++ 9 files changed, 271 insertions(+), 58 deletions(-) create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/type-provider-tagged-template-expression.expect.md create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/type-provider-tagged-template-expression.js diff --git a/compiler/packages/babel-plugin-react-compiler/src/Inference/InferReferenceEffects.ts b/compiler/packages/babel-plugin-react-compiler/src/Inference/InferReferenceEffects.ts index 8aa82469bdec4..1604f4813967a 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/Inference/InferReferenceEffects.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/Inference/InferReferenceEffects.ts @@ -1180,18 +1180,6 @@ function inferBlock( }; break; } - case 'TaggedTemplateExpression': { - valueKind = { - kind: ValueKind.Mutable, - reason: new Set([ValueReason.Other]), - context: new Set(), - }; - effect = { - kind: Effect.ConditionallyMutate, - reason: ValueReason.Other, - }; - break; - } case 'TemplateLiteral': { /* * template literal (with no tag function) always produces @@ -1312,6 +1300,47 @@ function inferBlock( instr.lvalue.effect = Effect.Store; continue; } + case 'TaggedTemplateExpression': { + const operands = [...eachInstructionValueOperand(instrValue)]; + if (operands.length !== 1) { + // future-proofing to make sure we update this case when we support interpolation + CompilerError.throwTodo({ + reason: 'Support tagged template expressions with interpolations', + loc: instrValue.loc, + }); + } + const signature = getFunctionCallSignature( + env, + instrValue.tag.identifier.type, + ); + let calleeEffect = + signature?.calleeEffect ?? Effect.ConditionallyMutate; + const returnValueKind: AbstractValue = + signature !== null + ? { + kind: signature.returnValueKind, + reason: new Set([ + signature.returnValueReason ?? + ValueReason.KnownReturnSignature, + ]), + context: new Set(), + } + : { + kind: ValueKind.Mutable, + reason: new Set([ValueReason.Other]), + context: new Set(), + }; + state.referenceAndRecordEffects( + instrValue.tag, + calleeEffect, + ValueReason.Other, + functionEffects, + ); + state.initialize(instrValue, returnValueKind); + state.define(instr.lvalue, instrValue); + instr.lvalue.effect = Effect.ConditionallyMutate; + continue; + } case 'CallExpression': { const signature = getFunctionCallSignature( env, diff --git a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/InferReactiveScopeVariables.ts b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/InferReactiveScopeVariables.ts index 27aba91af2b1c..126772f591b41 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/InferReactiveScopeVariables.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/InferReactiveScopeVariables.ts @@ -227,6 +227,7 @@ function mayAllocate(env: Environment, instruction: Instruction): boolean { case 'StoreGlobal': { return false; } + case 'TaggedTemplateExpression': case 'CallExpression': case 'MethodCall': { return instruction.lvalue.identifier.type.kind !== 'Primitive'; @@ -241,8 +242,7 @@ function mayAllocate(env: Environment, instruction: Instruction): boolean { case 'ObjectExpression': case 'UnsupportedNode': case 'ObjectMethod': - case 'FunctionExpression': - case 'TaggedTemplateExpression': { + case 'FunctionExpression': { return true; } default: { diff --git a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneNonEscapingScopes.ts b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneNonEscapingScopes.ts index b2e91fa302728..8033d05e2b3e3 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneNonEscapingScopes.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneNonEscapingScopes.ts @@ -671,12 +671,37 @@ function computeMemoizationInputs( ], }; } + case 'TaggedTemplateExpression': { + const signature = getFunctionCallSignature( + env, + value.tag.identifier.type, + ); + let lvalues = []; + if (lvalue !== null) { + lvalues.push({place: lvalue, level: MemoizationLevel.Memoized}); + } + if (signature?.noAlias === true) { + return { + lvalues, + rvalues: [], + }; + } + const operands = [...eachReactiveValueOperand(value)]; + lvalues.push( + ...operands + .filter(operand => isMutableEffect(operand.effect, operand.loc)) + .map(place => ({place, level: MemoizationLevel.Memoized})), + ); + return { + lvalues, + rvalues: operands, + }; + } case 'CallExpression': { const signature = getFunctionCallSignature( env, value.callee.identifier.type, ); - const operands = [...eachReactiveValueOperand(value)]; let lvalues = []; if (lvalue !== null) { lvalues.push({place: lvalue, level: MemoizationLevel.Memoized}); @@ -687,6 +712,7 @@ function computeMemoizationInputs( rvalues: [], }; } + const operands = [...eachReactiveValueOperand(value)]; lvalues.push( ...operands .filter(operand => isMutableEffect(operand.effect, operand.loc)) @@ -702,7 +728,6 @@ function computeMemoizationInputs( env, value.property.identifier.type, ); - const operands = [...eachReactiveValueOperand(value)]; let lvalues = []; if (lvalue !== null) { lvalues.push({place: lvalue, level: MemoizationLevel.Memoized}); @@ -713,6 +738,7 @@ function computeMemoizationInputs( rvalues: [], }; } + const operands = [...eachReactiveValueOperand(value)]; lvalues.push( ...operands .filter(operand => isMutableEffect(operand.effect, operand.loc)) @@ -726,7 +752,6 @@ function computeMemoizationInputs( case 'RegExpLiteral': case 'ObjectMethod': case 'FunctionExpression': - case 'TaggedTemplateExpression': case 'ArrayExpression': case 'NewExpression': case 'ObjectExpression': diff --git a/compiler/packages/babel-plugin-react-compiler/src/TypeInference/InferTypes.ts b/compiler/packages/babel-plugin-react-compiler/src/TypeInference/InferTypes.ts index d9f7ffd5bf8b8..b460124ec71f3 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/TypeInference/InferTypes.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/TypeInference/InferTypes.ts @@ -250,6 +250,7 @@ function* generateInstructionTypes( } case 'CallExpression': { + const returnType = makeType(); /* * TODO: callee could be a hook or a function, so this type equation isn't correct. * We should change Hook to a subtype of Function or change unifier logic. @@ -258,8 +259,25 @@ function* generateInstructionTypes( yield equation(value.callee.identifier.type, { kind: 'Function', shapeId: null, - return: left, + return: returnType, }); + yield equation(left, returnType); + break; + } + + case 'TaggedTemplateExpression': { + const returnType = makeType(); + /* + * TODO: callee could be a hook or a function, so this type equation isn't correct. + * We should change Hook to a subtype of Function or change unifier logic. + * (see https://github.com/facebook/react-forget/pull/1427) + */ + yield equation(value.tag.identifier.type, { + kind: 'Function', + shapeId: null, + return: returnType, + }); + yield equation(left, returnType); break; } @@ -392,7 +410,6 @@ function* generateInstructionTypes( case 'MetaProperty': case 'ComputedStore': case 'ComputedLoad': - case 'TaggedTemplateExpression': case 'Await': case 'GetIterator': case 'IteratorNext': diff --git a/compiler/packages/babel-plugin-react-compiler/src/Validation/ValidateLocalsNotReassignedAfterRender.ts b/compiler/packages/babel-plugin-react-compiler/src/Validation/ValidateLocalsNotReassignedAfterRender.ts index 0ea1814349f7f..9c41ebcae19f6 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/Validation/ValidateLocalsNotReassignedAfterRender.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/Validation/ValidateLocalsNotReassignedAfterRender.ts @@ -161,6 +161,14 @@ function getContextReassignment( if (signature?.noAlias) { operands = [value.receiver, value.property]; } + } else if (value.kind === 'TaggedTemplateExpression') { + const signature = getFunctionCallSignature( + fn.env, + value.tag.identifier.type, + ); + if (signature?.noAlias) { + operands = [value.tag]; + } } for (const operand of operands) { CompilerError.invariant(operand.effect !== Effect.Unknown, { diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/destructuring-mixed-scope-and-local-variables-with-default.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/destructuring-mixed-scope-and-local-variables-with-default.expect.md index 5e8f199206f58..17dd0f835942d 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/destructuring-mixed-scope-and-local-variables-with-default.expect.md +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/destructuring-mixed-scope-and-local-variables-with-default.expect.md @@ -63,67 +63,63 @@ function useFragment(_arg1, _arg2) { } function Component(props) { - const $ = _c(9); - let t0; - if ($[0] === Symbol.for("react.memo_cache_sentinel")) { - t0 = graphql` + const $ = _c(8); + const post = useFragment( + graphql` fragment F on T { id } - `; - $[0] = t0; - } else { - t0 = $[0]; - } - const post = useFragment(t0, props.post); - let t1; - if ($[1] !== post) { + `, + props.post, + ); + let t0; + if ($[0] !== post) { const allUrls = []; - const { media: t2, comments: t3, urls: t4 } = post; - const media = t2 === undefined ? null : t2; + const { media: t1, comments: t2, urls: t3 } = post; + const media = t1 === undefined ? null : t1; + let t4; + if ($[2] !== t2) { + t4 = t2 === undefined ? [] : t2; + $[2] = t2; + $[3] = t4; + } else { + t4 = $[3]; + } + const comments = t4; let t5; - if ($[3] !== t3) { + if ($[4] !== t3) { t5 = t3 === undefined ? [] : t3; - $[3] = t3; - $[4] = t5; + $[4] = t3; + $[5] = t5; } else { - t5 = $[4]; + t5 = $[5]; } - const comments = t5; + const urls = t5; let t6; - if ($[5] !== t4) { - t6 = t4 === undefined ? [] : t4; - $[5] = t4; - $[6] = t6; - } else { - t6 = $[6]; - } - const urls = t6; - let t7; - if ($[7] !== comments.length) { - t7 = (e) => { + if ($[6] !== comments.length) { + t6 = (e) => { if (!comments.length) { return; } console.log(comments.length); }; - $[7] = comments.length; - $[8] = t7; + $[6] = comments.length; + $[7] = t6; } else { - t7 = $[8]; + t6 = $[7]; } - const onClick = t7; + const onClick = t6; allUrls.push(...urls); - t1 = ; - $[1] = post; - $[2] = t1; + t0 = ; + $[0] = post; + $[1] = t0; } else { - t1 = $[2]; + t0 = $[1]; } - return t1; + return t0; } export const FIXTURE_ENTRYPOINT = { diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/type-provider-tagged-template-expression.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/type-provider-tagged-template-expression.expect.md new file mode 100644 index 0000000000000..03bfef9fb2eff --- /dev/null +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/type-provider-tagged-template-expression.expect.md @@ -0,0 +1,106 @@ + +## Input + +```javascript +import {graphql} from 'shared-runtime'; + +export function Component({a, b}) { + const fragment = graphql` + fragment Foo on User { + name + } + `; + return
{fragment}
; +} + +export const FIXTURE_ENTRYPOINT = { + fn: Component, + params: [{a: 0, b: 0}], + sequentialRenders: [ + {a: 0, b: 0}, + {a: 1, b: 0}, + {a: 1, b: 1}, + {a: 1, b: 2}, + {a: 2, b: 2}, + {a: 3, b: 2}, + {a: 0, b: 0}, + ], +}; + +``` + +## Code + +```javascript +import { c as _c } from "react/compiler-runtime"; +import { graphql } from "shared-runtime"; + +export function Component(t0) { + const $ = _c(1); + const fragment = graphql` + fragment Foo on User { + name + } + `; + let t1; + if ($[0] === Symbol.for("react.memo_cache_sentinel")) { + t1 =
{fragment}
; + $[0] = t1; + } else { + t1 = $[0]; + } + return t1; +} + +export const FIXTURE_ENTRYPOINT = { + fn: Component, + params: [{ a: 0, b: 0 }], + sequentialRenders: [ + { a: 0, b: 0 }, + { a: 1, b: 0 }, + { a: 1, b: 1 }, + { a: 1, b: 2 }, + { a: 2, b: 2 }, + { a: 3, b: 2 }, + { a: 0, b: 0 }, + ], +}; + +``` + +### Eval output +(kind: ok)
+ fragment Foo on User { + name + } +
+
+ fragment Foo on User { + name + } +
+
+ fragment Foo on User { + name + } +
+
+ fragment Foo on User { + name + } +
+
+ fragment Foo on User { + name + } +
+
+ fragment Foo on User { + name + } +
+
+ fragment Foo on User { + name + } +
\ No newline at end of file diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/type-provider-tagged-template-expression.js b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/type-provider-tagged-template-expression.js new file mode 100644 index 0000000000000..872d6b8f6fda9 --- /dev/null +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/type-provider-tagged-template-expression.js @@ -0,0 +1,24 @@ +import {graphql} from 'shared-runtime'; + +export function Component({a, b}) { + const fragment = graphql` + fragment Foo on User { + name + } + `; + return
{fragment}
; +} + +export const FIXTURE_ENTRYPOINT = { + fn: Component, + params: [{a: 0, b: 0}], + sequentialRenders: [ + {a: 0, b: 0}, + {a: 1, b: 0}, + {a: 1, b: 1}, + {a: 1, b: 2}, + {a: 2, b: 2}, + {a: 3, b: 2}, + {a: 0, b: 0}, + ], +}; diff --git a/compiler/packages/snap/src/sprout/shared-runtime-type-provider.ts b/compiler/packages/snap/src/sprout/shared-runtime-type-provider.ts index 9a07902dbc2f5..ba450191b2088 100644 --- a/compiler/packages/snap/src/sprout/shared-runtime-type-provider.ts +++ b/compiler/packages/snap/src/sprout/shared-runtime-type-provider.ts @@ -32,6 +32,14 @@ export function makeSharedRuntimeTypeProvider({ returnType: {kind: 'type', name: 'Primitive'}, returnValueKind: ValueKindEnum.Primitive, }, + graphql: { + kind: 'function', + calleeEffect: EffectEnum.Read, + positionalParams: [], + restParam: EffectEnum.Read, + returnType: {kind: 'type', name: 'Primitive'}, + returnValueKind: ValueKindEnum.Primitive, + }, typedArrayPush: { kind: 'function', calleeEffect: EffectEnum.Read,