[compiler] Fix fbt for the ∞th time (#34865)

We now do a single pass over the HIR, building up two data structures:
* One tracks values that are known macro tags or macro calls.
* One tracks operands of macro-related instructions so that we can later
group them.

After building up these data structures, we do a pass over the latter
structure. For each macro call instruction, we recursively traverse its
operands to ensure they're in the same scope. Thus, something like
`fbt('hello' + fbt.param(foo(), "..."))` will correctly merge the fbt
call, the `+` binary expression, the `fbt.param()` call, and `foo()`
into a single scope.

---
[//]: # (BEGIN SAPLING FOOTER)
Stack created with [Sapling](https://sapling-scm.com). Best reviewed
with [ReviewStack](https://reviewstack.dev/facebook/react/pull/34865).
* #34855
* __->__ #34865

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

@@ -7,14 +7,17 @@
 
 import {
   HIRFunction,
+  Identifier,
   IdentifierId,
+  InstructionValue,
   makeInstructionId,
   MutableRange,
   Place,
-  ReactiveValue,
+  ReactiveScope,
 } from '../HIR';
 import {Macro, MacroMethod} from '../HIR/Environment';
-import {eachReactiveValueOperand} from './visitors';
+import {eachInstructionValueOperand} from '../HIR/visitors';
+import {Iterable_some} from '../Utils/utils';
 
 /**
  * This pass supports the `fbt` translation system (https://facebook.github.io/fbt/)
@@ -48,24 +51,49 @@ export function memoizeFbtAndMacroOperandsInSameScope(
     ...Array.from(FBT_TAGS).map((tag): Macro => [tag, []]),
     ...(fn.env.config.customMacros ?? []),
   ]);
-  const fbtValues: Set<IdentifierId> = new Set();
+  /**
+   * Set of all identifiers that load fbt or other macro functions or their nested
+   * properties, as well as values known to be the results of invoking macros
+   */
+  const macroTagsCalls: Set<IdentifierId> = new Set();
+  /**
+   * Mapping of lvalue => list of operands for all expressions where either
+   * the lvalue is a known fbt/macro call and/or the operands transitively
+   * contain fbt/macro calls.
+   *
+   * This is the key data structure that powers the scope merging: we start
+   * at the lvalues and merge operands into the lvalue's scope.
+   */
+  const macroValues: Map<Identifier, Array<Identifier>> = new Map();
+  // Tracks methods loaded from macros, like fbt.param or idx.foo
   const macroMethods = new Map<IdentifierId, Array<Array<MacroMethod>>>();
-  while (true) {
-    let vsize = fbtValues.size;
-    let msize = macroMethods.size;
-    visit(fn, fbtMacroTags, fbtValues, macroMethods);
-    if (vsize === fbtValues.size && msize === macroMethods.size) {
-      break;
+
+  visit(fn, fbtMacroTags, macroTagsCalls, macroMethods, macroValues);
+
+  for (const root of macroValues.keys()) {
+    const scope = root.scope;
+    if (scope == null) {
+      continue;
     }
+    // Merge the operands into the same scope if this is a known macro invocation
+    if (!macroTagsCalls.has(root.id)) {
+      continue;
+    }
+    mergeScopes(root, scope, macroValues, macroTagsCalls);
   }
-  return fbtValues;
+
+  return macroTagsCalls;
 }
 
 export const FBT_TAGS: Set<string> = new Set([
   'fbt',
   'fbt:param',
+  'fbt:enum',
+  'fbt:plural',
   'fbs',
   'fbs:param',
+  'fbs:enum',
+  'fbs:plural',
 ]);
 export const SINGLE_CHILD_FBT_TAGS: Set<string> = new Set([
   'fbt:param',
@@ -75,10 +103,22 @@ export const SINGLE_CHILD_FBT_TAGS: Set<string> = new Set([
 function visit(
   fn: HIRFunction,
   fbtMacroTags: Set<Macro>,
-  fbtValues: Set<IdentifierId>,
+  macroTagsCalls: Set<IdentifierId>,
   macroMethods: Map<IdentifierId, Array<Array<MacroMethod>>>,
+  macroValues: Map<Identifier, Array<Identifier>>,
 ): void {
   for (const [, block] of fn.body.blocks) {
+    for (const phi of block.phis) {
+      const macroOperands: Array<Identifier> = [];
+      for (const operand of phi.operands.values()) {
+        if (macroValues.has(operand.identifier)) {
+          macroOperands.push(operand.identifier);
+        }
+      }
+      if (macroOperands.length !== 0) {
+        macroValues.set(phi.place.identifier, macroOperands);
+      }
+    }
     for (const instruction of block.instructions) {
       const {lvalue, value} = instruction;
       if (lvalue === null) {
@@ -93,13 +133,13 @@ function visit(
          * We don't distinguish between tag names and strings, so record
          * all `fbt` string literals in case they are used as a jsx tag.
          */
-        fbtValues.add(lvalue.identifier.id);
+        macroTagsCalls.add(lvalue.identifier.id);
       } else if (
         value.kind === 'LoadGlobal' &&
         matchesExactTag(value.binding.name, fbtMacroTags)
       ) {
         // Record references to `fbt` as a global
-        fbtValues.add(lvalue.identifier.id);
+        macroTagsCalls.add(lvalue.identifier.id);
       } else if (
         value.kind === 'LoadGlobal' &&
         matchTagRoot(value.binding.name, fbtMacroTags) !== null
@@ -121,84 +161,66 @@ function visit(
             if (method.length > 1) {
               newMethods.push(method.slice(1));
             } else {
-              fbtValues.add(lvalue.identifier.id);
+              macroTagsCalls.add(lvalue.identifier.id);
             }
           }
         }
         if (newMethods.length > 0) {
           macroMethods.set(lvalue.identifier.id, newMethods);
         }
-      } else if (isFbtCallExpression(fbtValues, value)) {
-        const fbtScope = lvalue.identifier.scope;
-        if (fbtScope === null) {
-          continue;
-        }
-
-        /*
-         * if the JSX element's tag was `fbt`, mark all its operands
-         * to ensure that they end up in the same scope as the jsx element
-         * itself.
-         */
-        for (const operand of eachReactiveValueOperand(value)) {
-          operand.identifier.scope = fbtScope;
-
-          // Expand the jsx element's range to account for its operands
-          expandFbtScopeRange(fbtScope.range, operand.identifier.mutableRange);
-          fbtValues.add(operand.identifier.id);
-        }
       } else if (
-        isFbtJsxExpression(fbtMacroTags, fbtValues, value) ||
-        isFbtJsxChild(fbtValues, lvalue, value)
+        value.kind === 'PropertyLoad' &&
+        macroTagsCalls.has(value.object.identifier.id)
       ) {
-        const fbtScope = lvalue.identifier.scope;
-        if (fbtScope === null) {
-          continue;
-        }
-
-        /*
-         * if the JSX element's tag was `fbt`, mark all its operands
-         * to ensure that they end up in the same scope as the jsx element
-         * itself.
-         */
-        for (const operand of eachReactiveValueOperand(value)) {
-          operand.identifier.scope = fbtScope;
-
-          // Expand the jsx element's range to account for its operands
-          expandFbtScopeRange(fbtScope.range, operand.identifier.mutableRange);
-
-          /*
-           * NOTE: we add the operands as fbt values so that they are also
-           * grouped with this expression
-           */
-          fbtValues.add(operand.identifier.id);
-        }
-      } else if (fbtValues.has(lvalue.identifier.id)) {
-        const fbtScope = lvalue.identifier.scope;
-        if (fbtScope === null) {
-          return;
-        }
-
-        for (const operand of eachReactiveValueOperand(value)) {
-          if (
-            operand.identifier.name !== null &&
-            operand.identifier.name.kind === 'named'
-          ) {
-            /*
-             * named identifiers were already locals, we only have to force temporaries
-             * into the same scope
-             */
-            continue;
+        macroTagsCalls.add(lvalue.identifier.id);
+      } else if (
+        isFbtJsxExpression(fbtMacroTags, macroTagsCalls, value) ||
+        isFbtJsxChild(macroTagsCalls, lvalue, value) ||
+        isFbtCallExpression(macroTagsCalls, value)
+      ) {
+        macroTagsCalls.add(lvalue.identifier.id);
+        macroValues.set(
+          lvalue.identifier,
+          Array.from(
+            eachInstructionValueOperand(value),
+            operand => operand.identifier,
+          ),
+        );
+      } else if (
+        Iterable_some(eachInstructionValueOperand(value), operand =>
+          macroValues.has(operand.identifier),
+        )
+      ) {
+        const macroOperands: Array<Identifier> = [];
+        for (const operand of eachInstructionValueOperand(value)) {
+          if (macroValues.has(operand.identifier)) {
+            macroOperands.push(operand.identifier);
           }
-          operand.identifier.scope = fbtScope;
-
-          // Expand the jsx element's range to account for its operands
-          expandFbtScopeRange(fbtScope.range, operand.identifier.mutableRange);
         }
+        macroValues.set(lvalue.identifier, macroOperands);
       }
     }
   }
 }
 
+function mergeScopes(
+  root: Identifier,
+  scope: ReactiveScope,
+  macroValues: Map<Identifier, Array<Identifier>>,
+  macroTagsCalls: Set<IdentifierId>,
+): void {
+  const operands = macroValues.get(root);
+  if (operands == null) {
+    return;
+  }
+  for (const operand of operands) {
+    operand.scope = scope;
+    expandFbtScopeRange(scope.range, operand.mutableRange);
+    macroTagsCalls.add(operand.id);
+    mergeScopes(operand, scope, macroValues, macroTagsCalls);
+  }
+}
+
 function matchesExactTag(s: string, tags: Set<Macro>): boolean {
   return Array.from(tags).some(macro =>
     typeof macro === 'string'
@@ -229,39 +251,40 @@ function matchTagRoot(
 }
 
 function isFbtCallExpression(
-  fbtValues: Set<IdentifierId>,
-  value: ReactiveValue,
+  macroTagsCalls: Set<IdentifierId>,
+  value: InstructionValue,
 ): boolean {
   return (
     (value.kind === 'CallExpression' &&
-      fbtValues.has(value.callee.identifier.id)) ||
-    (value.kind === 'MethodCall' && fbtValues.has(value.property.identifier.id))
+      macroTagsCalls.has(value.callee.identifier.id)) ||
+    (value.kind === 'MethodCall' &&
+      macroTagsCalls.has(value.property.identifier.id))
   );
 }
 
 function isFbtJsxExpression(
   fbtMacroTags: Set<Macro>,
-  fbtValues: Set<IdentifierId>,
-  value: ReactiveValue,
+  macroTagsCalls: Set<IdentifierId>,
+  value: InstructionValue,
 ): boolean {
   return (
     value.kind === 'JsxExpression' &&
     ((value.tag.kind === 'Identifier' &&
-      fbtValues.has(value.tag.identifier.id)) ||
+      macroTagsCalls.has(value.tag.identifier.id)) ||
       (value.tag.kind === 'BuiltinTag' &&
         matchesExactTag(value.tag.name, fbtMacroTags)))
   );
 }
 
 function isFbtJsxChild(
-  fbtValues: Set<IdentifierId>,
+  macroTagsCalls: Set<IdentifierId>,
   lvalue: Place | null,
-  value: ReactiveValue,
+  value: InstructionValue,
 ): boolean {
   return (
     (value.kind === 'JsxExpression' || value.kind === 'JsxFragment') &&
     lvalue !== null &&
-    fbtValues.has(lvalue.identifier.id)
+    macroTagsCalls.has(lvalue.identifier.id)
   );
 }
 

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/fbt/error.todo-fbt-param-nested-fbt.expect.md

@@ -1,56 +0,0 @@
-
-## Input
-
-```javascript
-import fbt from 'fbt';
-import {Stringify} from 'shared-runtime';
-
-/**
- * MemoizeFbtAndMacroOperands needs to account for nested fbt calls.
- * Expected fixture `fbt-param-call-arguments` to succeed but it failed with error:
- *   /fbt-param-call-arguments.ts: Line 19 Column 11: fbt: unsupported babel node: Identifier
- *   ---
- *   t3
- *   ---
- */
-function Component({firstname, lastname}) {
-  'use memo';
-  return (
-    <Stringify>
-      {fbt(
-        [
-          'Name: ',
-          fbt.param('firstname', <Stringify key={0} name={firstname} />),
-          ', ',
-          fbt.param(
-            'lastname',
-            <Stringify key={0} name={lastname}>
-              {fbt('(inner fbt)', 'Inner fbt value')}
-            </Stringify>
-          ),
-        ],
-        'Name'
-      )}
-    </Stringify>
-  );
-}
-
-export const FIXTURE_ENTRYPOINT = {
-  fn: Component,
-  params: [{firstname: 'first', lastname: 'last'}],
-  sequentialRenders: [{firstname: 'first', lastname: 'last'}],
-};
-
-```
-
-
-## Error
-
-```
-Line 19 Column 11: fbt: unsupported babel node: Identifier
----
-t3
----
-```
-          
-      

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/fbt/fbt-preserve-whitespace-two-subtrees.expect.md

@@ -37,27 +37,31 @@ import { c as _c } from "react/compiler-runtime";
 import fbt from "fbt";
 
 function Foo(t0) {
-  const $ = _c(3);
+  const $ = _c(7);
   const { name1, name2 } = t0;
   let t1;
   if ($[0] !== name1 || $[1] !== name2) {
+    let t2;
+    if ($[3] !== name1) {
+      t2 = <b>{name1}</b>;
+      $[3] = name1;
+      $[4] = t2;
+    } else {
+      t2 = $[4];
+    }
+    let t3;
+    if ($[5] !== name2) {
+      t3 = <b>{name2}</b>;
+      $[5] = name2;
+      $[6] = t3;
+    } else {
+      t3 = $[6];
+    }
     t1 = fbt._(
       "{user1} and {user2} accepted your PR!",
       [
-        fbt._param(
-          "user1",
-
-          <span key={name1}>
-            <b>{name1}</b>
-          </span>,
-        ),
-        fbt._param(
-          "user2",
-
-          <span key={name2}>
-            <b>{name2}</b>
-          </span>,
-        ),
+        fbt._param("user1", <span key={name1}>{t2}</span>),
+        fbt._param("user2", <span key={name2}>{t3}</span>),
       ],
       { hk: "2PxMie" },
     );

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/fbt/repro-fbt-param-nested-fbt.expect.md

@@ -0,0 +1,111 @@
+
+## Input
+
+```javascript
+import fbt from 'fbt';
+import {Stringify} from 'shared-runtime';
+
+/**
+ * MemoizeFbtAndMacroOperands needs to account for nested fbt calls.
+ * Expected fixture `fbt-param-call-arguments` to succeed but it failed with error:
+ *   /fbt-param-call-arguments.ts: Line 19 Column 11: fbt: unsupported babel node: Identifier
+ *   ---
+ *   t3
+ *   ---
+ */
+function Component({firstname, lastname}) {
+  'use memo';
+  return (
+    <Stringify>
+      {fbt(
+        [
+          'Name: ',
+          fbt.param('firstname', <Stringify key={0} name={firstname} />),
+          ', ',
+          fbt.param(
+            'lastname',
+            <Stringify key={0} name={lastname}>
+              {fbt('(inner fbt)', 'Inner fbt value')}
+            </Stringify>
+          ),
+        ],
+        'Name'
+      )}
+    </Stringify>
+  );
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{firstname: 'first', lastname: 'last'}],
+  sequentialRenders: [{firstname: 'first', lastname: 'last'}],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime";
+import fbt from "fbt";
+import { Stringify } from "shared-runtime";
+
+/**
+ * MemoizeFbtAndMacroOperands needs to account for nested fbt calls.
+ * Expected fixture `fbt-param-call-arguments` to succeed but it failed with error:
+ *   /fbt-param-call-arguments.ts: Line 19 Column 11: fbt: unsupported babel node: Identifier
+ *   ---
+ *   t3
+ *   ---
+ */
+function Component(t0) {
+  "use memo";
+  const $ = _c(5);
+  const { firstname, lastname } = t0;
+  let t1;
+  if ($[0] !== firstname || $[1] !== lastname) {
+    t1 = fbt._(
+      "Name: {firstname}, {lastname}",
+      [
+        fbt._param(
+          "firstname",
+
+          <Stringify key={0} name={firstname} />,
+        ),
+        fbt._param(
+          "lastname",
+
+          <Stringify key={0} name={lastname}>
+            {fbt._("(inner fbt)", null, { hk: "36qNwF" })}
+          </Stringify>,
+        ),
+      ],
+      { hk: "3AiIf8" },
+    );
+    $[0] = firstname;
+    $[1] = lastname;
+    $[2] = t1;
+  } else {
+    t1 = $[2];
+  }
+  let t2;
+  if ($[3] !== t1) {
+    t2 = <Stringify>{t1}</Stringify>;
+    $[3] = t1;
+    $[4] = t2;
+  } else {
+    t2 = $[4];
+  }
+  return t2;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ firstname: "first", lastname: "last" }],
+  sequentialRenders: [{ firstname: "first", lastname: "last" }],
+};
+
+```
+      
+### Eval output
+(kind: ok) <div>{"children":"Name: , "}</div>

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/fbt/repro-separately-memoized-fbt-param.expect.md

@@ -0,0 +1,78 @@
+
+## Input
+
+```javascript
+import {fbt} from 'fbt';
+import {useState} from 'react';
+
+const MIN = 10;
+
+function Component() {
+  const [count, setCount] = useState(0);
+
+  return fbt(
+    'Expected at least ' +
+      fbt.param('min', MIN, {number: true}) +
+      ' items, but got ' +
+      fbt.param('count', count, {number: true}) +
+      ' items.',
+    'Error description'
+  );
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{}],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime";
+import { fbt } from "fbt";
+import { useState } from "react";
+
+const MIN = 10;
+
+function Component() {
+  const $ = _c(2);
+  const [count] = useState(0);
+  let t0;
+  if ($[0] !== count) {
+    t0 = fbt._(
+      { "*": { "*": "Expected at least {min} items, but got {count} items." } },
+      [
+        fbt._param(
+          "min",
+
+          MIN,
+          [0],
+        ),
+        fbt._param(
+          "count",
+
+          count,
+          [0],
+        ),
+      ],
+      { hk: "36gbz8" },
+    );
+    $[0] = count;
+    $[1] = t0;
+  } else {
+    t0 = $[1];
+  }
+  return t0;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{}],
+};
+
+```
+      
+### Eval output
+(kind: ok) Expected at least 10 items, but got 0 items.

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/fbt/repro-separately-memoized-fbt-param.js

@@ -0,0 +1,22 @@
+import {fbt} from 'fbt';
+import {useState} from 'react';
+
+const MIN = 10;
+
+function Component() {
+  const [count, setCount] = useState(0);
+
+  return fbt(
+    'Expected at least ' +
+      fbt.param('min', MIN, {number: true}) +
+      ' items, but got ' +
+      fbt.param('count', count, {number: true}) +
+      ' items.',
+    'Error description'
+  );
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{}],
+};

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/idx-method-no-outlining-wildcard.expect.md

@@ -73,7 +73,7 @@ function Component(props) {
   const groupName4 = t3;
   let t4;
   if ($[8] !== props) {
-    t4 = idx.hello_world.b.c(props, _temp3);
+    t4 = idx.hello_world.b.c(props, (__3) => __3.group.label);
     $[8] = props;
     $[9] = t4;
   } else {
@@ -108,9 +108,6 @@ function Component(props) {
   }
   return t5;
 }
-function _temp3(__3) {
-  return __3.group.label;
-}
 function _temp2(__0) {
   return __0.group.label;
 }

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/idx-method-no-outlining.expect.md

@@ -49,7 +49,7 @@ function Component(props) {
   const groupName2 = t1;
   let t2;
   if ($[4] !== props) {
-    t2 = idx.a.b(props, _temp2);
+    t2 = idx.a.b(props, (__1) => __1.group.label);
     $[4] = props;
     $[5] = t2;
   } else {
@@ -74,9 +74,6 @@ function Component(props) {
   }
   return t3;
 }
-function _temp2(__1) {
-  return __1.group.label;
-}
 function _temp(_) {
   return _.group.label;
 }