Update base for Update on "[compiler] Optimize emission in normal (non-value) blocks"

In #29863 I tried to find a clean way to share code for emitting instructions between value blocks and regular blocks. The catch is that value blocks have special meaning for their final instruction — that's the value of the block — so reordering can't change the last instruction. However, in finding a clean way to share code for these two categories of code, i also inadvertently reduced the effectiveness of the optimization.

This PR updates to use different strategies for these two kinds of blocks: value blocks use the code from #29863 where we first emit all non-reorderable instructions in their original order, _then_ try to emit reorderable values. The reason this is suboptimal, though, is that we want to move instructions closer to their dependencies so that they can invalidate (merge) together. Emitting the reorderable values last prevents this.

So for normal blocks, we now emit terminal operands first. This will invariably cause _some_ of the non-reorderable instructions to be emitted, but it will intersperse reoderable instructions in between, right after their dependencies. This maximizes our ability to merge scopes.

I think the complexity cost of two strategies is worth the benefit, though i still need to double-check all the output changes.

[ghstack-poisoned]

Author: josephsavona

.prettierrc.js

@@ -14,6 +14,12 @@ module.exports = {
   parser: 'flow',
   arrowParens: 'avoid',
   overrides: [
+    {
+      files: ['*.code-workspace'],
+      options: {
+        parser: 'json-stringify',
+      },
+    },
     {
       files: esNextPaths,
       options: {

compiler/packages/babel-plugin-react-compiler/package.json

@@ -1,6 +1,6 @@
 {
   "name": "babel-plugin-react-compiler",
-  "version": "0.0.0-experimental-938cd9a-20240601",
+  "version": "0.0.0-experimental-179941d-20240614",
   "description": "Babel plugin for React Compiler.",
   "main": "dist/index.js",
   "license": "MIT",

compiler/packages/babel-plugin-react-compiler/src/HIR/BuildHIR.ts

@@ -1668,6 +1668,15 @@ function lowerExpression(
       const left = lowerExpressionToTemporary(builder, leftPath);
       const right = lowerExpressionToTemporary(builder, expr.get("right"));
       const operator = expr.node.operator;
+      if (operator === "|>") {
+        builder.errors.push({
+          reason: `(BuildHIR::lowerExpression) Pipe operator not supported`,
+          severity: ErrorSeverity.Todo,
+          loc: leftPath.node.loc ?? null,
+          suggestions: null,
+        });
+        return { kind: "UnsupportedNode", node: exprNode, loc: exprLoc };
+      }
       return {
         kind: "BinaryExpression",
         operator,
@@ -1893,7 +1902,9 @@ function lowerExpression(
         );
       }
 
-      const operators: { [key: string]: t.BinaryExpression["operator"] } = {
+      const operators: {
+        [key: string]: Exclude<t.BinaryExpression["operator"], "|>">;
+      } = {
         "+=": "+",
         "-=": "-",
         "/=": "/",
@@ -2307,6 +2318,20 @@ function lowerExpression(
           });
           return { kind: "UnsupportedNode", node: expr.node, loc: exprLoc };
         }
+      } else if (expr.node.operator === "throw") {
+        builder.errors.push({
+          reason: `Throw expressions are not supported`,
+          severity: ErrorSeverity.InvalidJS,
+          loc: expr.node.loc ?? null,
+          suggestions: [
+            {
+              description: "Remove this line",
+              range: [expr.node.start!, expr.node.end!],
+              op: CompilerSuggestionOperation.Remove,
+            },
+          ],
+        });
+        return { kind: "UnsupportedNode", node: expr.node, loc: exprLoc };
       } else {
         return {
           kind: "UnaryExpression",

compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts

@@ -119,6 +119,18 @@ export type Hook = z.infer<typeof HookSchema>;
 const EnvironmentConfigSchema = z.object({
   customHooks: z.map(z.string(), HookSchema).optional().default(new Map()),
 
+  /**
+   * A list of functions which the application compiles as macros, where
+   * the compiler must ensure they are not compiled to rename the macro or separate the
+   * "function" from its argument.
+   *
+   * For example, Meta has some APIs such as `featureflag("name-of-feature-flag")` which
+   * are rewritten by a plugin. Assigning `featureflag` to a temporary would break the
+   * plugin since it looks specifically for the name of the function being invoked, not
+   * following aliases.
+   */
+  customMacros: z.nullable(z.array(z.string())).default(null),
+
   /**
    * Enable a check that resets the memoization cache when the source code of the file changes.
    * This is intended to support hot module reloading (HMR), where the same runtime component
@@ -416,6 +428,17 @@ export function parseConfigPragma(pragma: string): EnvironmentConfig {
       continue;
     }
 
+    if (
+      key === "enableChangeDetectionForDebugging" &&
+      (val === undefined || val === "true")
+    ) {
+      maybeConfig[key] = {
+        source: "react-compiler-runtime",
+        importSpecifierName: "$structuralCheck",
+      };
+      continue;
+    }
+
     if (typeof defaultConfig[key as keyof EnvironmentConfig] !== "boolean") {
       // skip parsing non-boolean properties
       continue;

compiler/packages/babel-plugin-react-compiler/src/HIR/HIR.ts

@@ -888,7 +888,7 @@ export type InstructionValue =
   | JSXText
   | {
       kind: "BinaryExpression";
-      operator: t.BinaryExpression["operator"];
+      operator: Exclude<t.BinaryExpression["operator"], "|>">;
       left: Place;
       right: Place;
       loc: SourceLocation;
@@ -903,7 +903,7 @@ export type InstructionValue =
   | MethodCall
   | {
       kind: "UnaryExpression";
-      operator: t.UnaryExpression["operator"];
+      operator: Exclude<t.UnaryExpression["operator"], "throw" | "delete">;
       value: Place;
       loc: SourceLocation;
     }

compiler/packages/babel-plugin-react-compiler/src/HIR/visitors.ts

@@ -856,7 +856,7 @@ export function mapTerminalSuccessors(
       const block = fn(terminal.block);
       const fallthrough = fn(terminal.fallthrough);
       return {
-        kind: "scope",
+        kind: terminal.kind,
         scope: terminal.scope,
         block,
         fallthrough,

compiler/packages/babel-plugin-react-compiler/src/Optimization/InstructionReordering.ts

@@ -14,7 +14,6 @@ import {
   IdentifierId,
   Instruction,
   InstructionId,
-  MutableRange,
   Place,
   isExpressionBlockKind,
   makeInstructionId,
@@ -26,7 +25,6 @@ import {
   eachInstructionValueLValue,
   eachInstructionValueOperand,
   eachTerminalOperand,
-  terminalFallthrough,
 } from "../HIR/visitors";
 import { mayAllocate } from "../ReactiveScopes/InferReactiveScopeVariables";
 import { getOrInsertWith } from "../Utils/utils";

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

@@ -43,7 +43,7 @@ import { Err, Ok, Result } from "../Utils/Result";
 import { GuardKind } from "../Utils/RuntimeDiagnosticConstants";
 import { assertExhaustive } from "../Utils/utils";
 import { buildReactiveFunction } from "./BuildReactiveFunction";
-import { SINGLE_CHILD_FBT_TAGS } from "./MemoizeFbtOperandsInSameScope";
+import { SINGLE_CHILD_FBT_TAGS } from "./MemoizeFbtAndMacroOperandsInSameScope";
 import { ReactiveFunctionVisitor, visitReactiveFunction } from "./visitors";
 
 export const MEMO_CACHE_SENTINEL = "react.memo_cache_sentinel";

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

@@ -14,8 +14,15 @@ import {
 } from "../HIR";
 import { eachReactiveValueOperand } from "./visitors";
 
-/*
- * This pass supports the `fbt` translation system (https://facebook.github.io/fbt/).
+/**
+ * This pass supports the
+ * This pass supports the `fbt` translation system (https://facebook.github.io/fbt/)
+ * as well as similar user-configurable macro-like APIs where it's important that
+ * the name of the function not be changed, and it's literal arguments not be turned
+ * into temporaries.
+ *
+ * ## FBT
+ *
  * FBT provides the `<fbt>` JSX element and `fbt()` calls (which take params in the
  * form of `<fbt:param>` children or `fbt.param()` arguments, respectively). These
  * tags/functions have restrictions on what types of syntax may appear as props/children/
@@ -26,13 +33,22 @@ import { eachReactiveValueOperand } from "./visitors";
  * operands to fbt tags/calls have the same scope as the tag/call itself.
  *
  * Note that this still allows the props/arguments of `<fbt:param>`/`fbt.param()`
- * to be independently memoized
+ * to be independently memoized.
+ *
+ * ## User-defined macro-like function
+ *
+ * Users can also specify their own functions to be treated similarly to fbt via the
+ * `customMacros` environment configuration.
  */
-export function memoizeFbtOperandsInSameScope(fn: HIRFunction): void {
+export function memoizeFbtAndMacroOperandsInSameScope(fn: HIRFunction): void {
+  const fbtMacroTags = new Set([
+    ...FBT_TAGS,
+    ...(fn.env.config.customMacros ?? []),
+  ]);
   const fbtValues: Set<IdentifierId> = new Set();
   while (true) {
     let size = fbtValues.size;
-    visit(fn, fbtValues);
+    visit(fn, fbtMacroTags, fbtValues);
     if (size === fbtValues.size) {
       break;
     }
@@ -50,7 +66,11 @@ export const SINGLE_CHILD_FBT_TAGS: Set<string> = new Set([
   "fbs:param",
 ]);
 
-function visit(fn: HIRFunction, fbtValues: Set<IdentifierId>): void {
+function visit(
+  fn: HIRFunction,
+  fbtMacroTags: Set<string>,
+  fbtValues: Set<IdentifierId>
+): void {
   for (const [, block] of fn.body.blocks) {
     for (const instruction of block.instructions) {
       const { lvalue, value } = instruction;
@@ -60,7 +80,7 @@ function visit(fn: HIRFunction, fbtValues: Set<IdentifierId>): void {
       if (
         value.kind === "Primitive" &&
         typeof value.value === "string" &&
-        FBT_TAGS.has(value.value)
+        fbtMacroTags.has(value.value)
       ) {
         /*
          * We don't distinguish between tag names and strings, so record
@@ -69,7 +89,7 @@ function visit(fn: HIRFunction, fbtValues: Set<IdentifierId>): void {
         fbtValues.add(lvalue.identifier.id);
       } else if (
         value.kind === "LoadGlobal" &&
-        FBT_TAGS.has(value.binding.name)
+        fbtMacroTags.has(value.binding.name)
       ) {
         // Record references to `fbt` as a global
         fbtValues.add(lvalue.identifier.id);
@@ -97,7 +117,7 @@ function visit(fn: HIRFunction, fbtValues: Set<IdentifierId>): void {
           fbtValues.add(operand.identifier.id);
         }
       } else if (
-        isFbtJsxExpression(fbtValues, value) ||
+        isFbtJsxExpression(fbtMacroTags, fbtValues, value) ||
         isFbtJsxChild(fbtValues, lvalue, value)
       ) {
         const fbtScope = lvalue.identifier.scope;
@@ -169,14 +189,15 @@ function isFbtCallExpression(
 }
 
 function isFbtJsxExpression(
+  fbtMacroTags: Set<string>,
   fbtValues: Set<IdentifierId>,
   value: ReactiveValue
 ): boolean {
   return (
     value.kind === "JsxExpression" &&
     ((value.tag.kind === "Identifier" &&
       fbtValues.has(value.tag.identifier.id)) ||
-      (value.tag.kind === "BuiltinTag" && FBT_TAGS.has(value.tag.name)))
+      (value.tag.kind === "BuiltinTag" && fbtMacroTags.has(value.tag.name)))
   );
 }
 

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

@@ -19,7 +19,7 @@ export { extractScopeDeclarationsFromDestructuring } from "./ExtractScopeDeclara
 export { flattenReactiveLoops } from "./FlattenReactiveLoops";
 export { flattenScopesWithHooksOrUse } from "./FlattenScopesWithHooksOrUse";
 export { inferReactiveScopeVariables } from "./InferReactiveScopeVariables";
-export { memoizeFbtOperandsInSameScope } from "./MemoizeFbtOperandsInSameScope";
+export { memoizeFbtAndMacroOperandsInSameScope as memoizeFbtOperandsInSameScope } from "./MemoizeFbtAndMacroOperandsInSameScope";
 export { mergeOverlappingReactiveScopes } from "./MergeOverlappingReactiveScopes";
 export { mergeReactiveScopesThatInvalidateTogether } from "./MergeReactiveScopesThatInvalidateTogether";
 export { printReactiveFunction } from "./PrintReactiveFunction";

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/meta-isms/repro-cx-assigned-to-temporary.expect.md

@@ -0,0 +1,100 @@
+
+## Input
+
+```javascript
+// @compilationMode(infer) @enableAssumeHooksFollowRulesOfReact:false @customMacros(cx)
+import { identity } from "shared-runtime";
+
+const DARK = "dark";
+
+function Component() {
+  const theme = useTheme();
+  return (
+    <div
+      className={cx({
+        "styles/light": true,
+        "styles/dark": theme.getTheme() === DARK,
+      })}
+    />
+  );
+}
+
+function cx(obj) {
+  const classes = [];
+  for (const [key, value] of Object.entries(obj)) {
+    if (value) {
+      classes.push(key);
+    }
+  }
+  return classes.join(" ");
+}
+
+function useTheme() {
+  return {
+    getTheme() {
+      return DARK;
+    },
+  };
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{}],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime"; // @compilationMode(infer) @enableAssumeHooksFollowRulesOfReact:false @customMacros(cx)
+import { identity } from "shared-runtime";
+
+const DARK = "dark";
+
+function Component() {
+  const $ = _c(2);
+  const theme = useTheme();
+
+  const t0 = cx({
+    "styles/light": true,
+    "styles/dark": theme.getTheme() === DARK,
+  });
+  let t1;
+  if ($[0] !== t0) {
+    t1 = <div className={t0} />;
+    $[0] = t0;
+    $[1] = t1;
+  } else {
+    t1 = $[1];
+  }
+  return t1;
+}
+
+function cx(obj) {
+  const classes = [];
+  for (const [key, value] of Object.entries(obj)) {
+    if (value) {
+      classes.push(key);
+    }
+  }
+  return classes.join(" ");
+}
+
+function useTheme() {
+  return {
+    getTheme() {
+      return DARK;
+    },
+  };
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{}],
+};
+
+```
+      
+### Eval output
+(kind: ok) <div class="styles/light styles/dark"></div>