[compiler] Type inference for tagged template literals

josephsavona · josephsavona · commit 857c1dc67348 · 2024-09-03T18:44:39.000-07:00
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: 00aceed Pull Request resolved: #30869
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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 = <Stringify media={media} allUrls={allUrls} onClick={onClick} />;
-    $[1] = post;
-    $[2] = t1;
+    t0 = <Stringify media={media} allUrls={allUrls} onClick={onClick} />;
+    $[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
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
diff --git a/compiler/packages/snap/src/sprout/shared-runtime-type-provider.ts b/compiler/packages/snap/src/sprout/shared-runtime-type-provider.ts
