Memoize fbt children in same scope

Fixes another special-case rule of fbt that i wasn't aware of: apparently 
`<fbt:param>` elements don't have to appear as direct children of `<fbt>`, they 
can be nested, and in this case they must appear as direct children of the fbt 
and not via an identifier indirection. This PR recursively extends the scope of 
FBT operands to make this work.

compiler/packages/babel-plugin-react-forget/src/Entrypoint/Pipeline.ts

@@ -55,7 +55,11 @@ import {
 } from "../ReactiveScopes";
 import { eliminateRedundantPhi, enterSSA, leaveSSA } from "../SSA";
 import { inferTypes } from "../TypeInference";
-import { logHIRFunction, logReactiveFunction } from "../Utils/logger";
+import {
+  logCodegenFunction,
+  logHIRFunction,
+  logReactiveFunction,
+} from "../Utils/logger";
 import { assertExhaustive } from "../Utils/utils";
 import {
   validateFrozenLambdas,
@@ -319,6 +323,7 @@ export function compileFn(
 export function log(value: CompilerPipelineValue): CompilerPipelineValue {
   switch (value.kind) {
     case "ast": {
+      logCodegenFunction(value.name, value.value);
       break;
     }
     case "hir": {

compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/MemoizeFbtOperandsInSameScope.ts

@@ -8,6 +8,7 @@
 import {
   IdentifierId,
   makeInstructionId,
+  Place,
   ReactiveFunction,
   ReactiveInstruction,
   ReactiveValue,
@@ -33,9 +34,18 @@ import {
  * to be independently memoized
  */
 export function memoizeFbtOperandsInSameScope(fn: ReactiveFunction): void {
-  visitReactiveFunction(fn, new Transform(), undefined);
+  const transform = new Transform();
+  while (true) {
+    let size = transform.fbtValues.size;
+    visitReactiveFunction(fn, transform, undefined);
+    if (size === transform.fbtValues.size) {
+      break;
+    }
+  }
 }
 
+const FBT_TAGS: Set<string> = new Set(["fbt", "fbt:param"]);
+
 class Transform extends ReactiveFunctionVisitor<void> {
   // Values that represent *potential* references of `fbt` as a JSX tag name
   // or as a callee.
@@ -52,19 +62,15 @@ class Transform extends ReactiveFunctionVisitor<void> {
     if (
       value.kind === "Primitive" &&
       typeof value.value === "string" &&
-      value.value === "fbt"
+      FBT_TAGS.has(value.value)
     ) {
       // We don't distinguish between tag names and strings, so record
       // all `fbt` string literals in case they are used as a jsx tag.
       this.fbtValues.add(lvalue.identifier.id);
-    } else if (value.kind === "LoadGlobal" && value.name === "fbt") {
+    } else if (value.kind === "LoadGlobal" && FBT_TAGS.has(value.name)) {
       // Record references to `fbt` as a global
       this.fbtValues.add(lvalue.identifier.id);
-    } else if (
-      isFbtJsxExpression(this.fbtValues, value) ||
-      (value.kind === "CallExpression" &&
-        this.fbtValues.has(value.callee.identifier.id))
-    ) {
+    } else if (isFbtCallExpression(this.fbtValues, value)) {
       const fbtScope = lvalue.identifier.scope;
       if (fbtScope === null) {
         return;
@@ -81,10 +87,43 @@ class Transform extends ReactiveFunctionVisitor<void> {
           Math.min(fbtScope.range.start, operand.identifier.mutableRange.start)
         );
       }
+    } else if (
+      isFbtJsxExpression(this.fbtValues, value) ||
+      isFbtJsxChild(this.fbtValues, lvalue, value)
+    ) {
+      const fbtScope = lvalue.identifier.scope;
+      if (fbtScope === null) {
+        return;
+      }
+
+      // 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
+        fbtScope.range.start = makeInstructionId(
+          Math.min(fbtScope.range.start, operand.identifier.mutableRange.start)
+        );
+
+        // NOTE: we add the operands as fbt values so that they are also
+        // grouped with this expression
+        this.fbtValues.add(operand.identifier.id);
+      }
     }
   }
 }
 
+function isFbtCallExpression(
+  fbtValues: Set<IdentifierId>,
+  value: ReactiveValue
+): boolean {
+  return (
+    value.kind === "CallExpression" && fbtValues.has(value.callee.identifier.id)
+  );
+}
+
 function isFbtJsxExpression(
   fbtValues: Set<IdentifierId>,
   value: ReactiveValue
@@ -93,6 +132,18 @@ function isFbtJsxExpression(
     value.kind === "JsxExpression" &&
     ((value.tag.kind === "Identifier" &&
       fbtValues.has(value.tag.identifier.id)) ||
-      (value.tag.kind === "BuiltinTag" && value.tag.name === "fbt"))
+      (value.tag.kind === "BuiltinTag" && FBT_TAGS.has(value.tag.name)))
+  );
+}
+
+function isFbtJsxChild(
+  fbtValues: Set<IdentifierId>,
+  lvalue: Place | null,
+  value: ReactiveValue
+): boolean {
+  return (
+    (value.kind === "JsxExpression" || value.kind === "JsxFragment") &&
+    lvalue !== null &&
+    fbtValues.has(lvalue.identifier.id)
   );
 }

compiler/packages/babel-plugin-react-forget/src/Utils/logger.ts

@@ -5,10 +5,13 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+import generate from "@babel/generator";
+import * as t from "@babel/types";
 import chalk from "chalk";
+import { format } from "prettier";
 import { HIR, HIRFunction, ReactiveFunction } from "../HIR/HIR";
 import { printFunction, printHIR } from "../HIR/PrintHIR";
-import { printReactiveFunction } from "../ReactiveScopes";
+import { CodegenFunction, printReactiveFunction } from "../ReactiveScopes";
 
 let ENABLED: boolean = false;
 
@@ -30,6 +33,34 @@ export function logHIR(step: string, ir: HIR): void {
   }
 }
 
+export function logCodegenFunction(step: string, fn: CodegenFunction): void {
+  if (ENABLED) {
+    let printed: string | null = null;
+    try {
+      const node = t.functionDeclaration(
+        fn.id,
+        fn.params,
+        fn.body,
+        fn.generator,
+        fn.async
+      );
+      const ast = generate(node);
+      printed = format(ast.code);
+    } catch (e) {
+      console.log("Error formatting AST: " + e.message);
+    }
+    if (printed === null) {
+      return;
+    }
+    if (printed !== lastLogged) {
+      lastLogged = printed;
+      process.stdout.write(`${chalk.green(step)}:\n${printed}\n\n`);
+    } else {
+      process.stdout.write(`${chalk.blue(step)}: (no change)\n\n`);
+    }
+  }
+}
+
 export function logHIRFunction(step: string, fn: HIRFunction): void {
   if (ENABLED) {
     const printed = printFunction(fn);

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/fbt-params-complex-param-value.expect.md

@@ -21,28 +21,21 @@ import { unstable_useMemoCache as useMemoCache } from "react";
 import fbt from "fbt";
 
 function Component(props) {
-  const $ = useMemoCache(4);
+  const $ = useMemoCache(2);
   const c_0 = $[0] !== props.name;
-  let t1;
+  let t0;
   if (c_0) {
-    const c_2 = $[2] !== props.name;
-    let t0;
-    if (c_2) {
-      t0 = capitalize(props.name);
-      $[2] = props.name;
-      $[3] = t0;
-    } else {
-      t0 = $[3];
-    }
-    t1 = fbt._("Hello {user name}", [fbt._param("user name", t0)], {
-      hk: "2zEDKF",
-    });
+    t0 = fbt._(
+      "Hello {user name}",
+      [fbt._param("user name", capitalize(props.name))],
+      { hk: "2zEDKF" }
+    );
     $[0] = props.name;
-    $[1] = t1;
+    $[1] = t0;
   } else {
-    t1 = $[1];
+    t0 = $[1];
   }
-  return t1;
+  return t0;
 }
 
 ```

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/fbtparam-with-jsx-element-content.expect.md

@@ -0,0 +1,84 @@
+
+## Input
+
+```javascript
+// @debug
+import fbt from "fbt";
+
+function Component({ name, data, icon }) {
+  return (
+    <Text type="body4">
+      <fbt desc="Lorem ipsum">
+        <fbt:param name="item author">
+          <Text type="h4">{name}</Text>
+        </fbt:param>
+        <fbt:param name="icon">{icon}</fbt:param>
+        <Text type="h4">
+          <fbt:param name="item details">{data}</fbt:param>
+        </Text>
+      </fbt>
+    </Text>
+  );
+}
+
+```
+
+## Code
+
+```javascript
+import { unstable_useMemoCache as useMemoCache } from "react"; // @debug
+import fbt from "fbt";
+
+function Component(t39) {
+  const $ = useMemoCache(6);
+  const { name, data, icon } = t39;
+  const c_0 = $[0] !== name;
+  const c_1 = $[1] !== icon;
+  const c_2 = $[2] !== data;
+  let t0;
+  if (c_0 || c_1 || c_2) {
+    t0 = fbt._(
+      "{item author}{icon}{=m2}",
+      [
+        fbt._param(
+          "item author",
+
+          <Text type="h4">{name}</Text>
+        ),
+        fbt._param(
+          "icon",
+
+          icon
+        ),
+        fbt._implicitParam(
+          "=m2",
+          <Text type="h4">
+            {fbt._("{item details}", [fbt._param("item details", data)], {
+              hk: "4jLfVq",
+            })}
+          </Text>
+        ),
+      ],
+      { hk: "2HLm2j" }
+    );
+    $[0] = name;
+    $[1] = icon;
+    $[2] = data;
+    $[3] = t0;
+  } else {
+    t0 = $[3];
+  }
+  const c_4 = $[4] !== t0;
+  let t1;
+  if (c_4) {
+    t1 = <Text type="body4">{t0}</Text>;
+    $[4] = t0;
+    $[5] = t1;
+  } else {
+    t1 = $[5];
+  }
+  return t1;
+}
+
+```
+      

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/fbtparam-with-jsx-element-content.js

@@ -0,0 +1,18 @@
+// @debug
+import fbt from "fbt";
+
+function Component({ name, data, icon }) {
+  return (
+    <Text type="body4">
+      <fbt desc="Lorem ipsum">
+        <fbt:param name="item author">
+          <Text type="h4">{name}</Text>
+        </fbt:param>
+        <fbt:param name="icon">{icon}</fbt:param>
+        <Text type="h4">
+          <fbt:param name="item details">{data}</fbt:param>
+        </Text>
+      </fbt>
+    </Text>
+  );
+}

compiler/packages/sprout/src/SproutTodoFilter.ts

@@ -454,6 +454,7 @@ const skipFilter = new Set([
   "infer-function-expression-React-memo-gating",
   "infer-skip-components-without-hooks-or-jsx",
   "class-component-with-render-helper",
+  "fbtparam-with-jsx-element-content",
 ]);
 
 export default skipFilter;