[compiler] Provide support for custom fbt-like macro functions

ghstack-source-id: 6c05e7aa20ba8cbd19c932b052ce4c744ed5f07b
Pull Request resolved: #29893

Author: josephsavona

Diff for: 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

Diff for: 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";

Diff for: compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/MemoizeFbtOperandsInSameScope.ts renamed to compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/MemoizeFbtAndMacroOperandsInSameScope.ts

@@ -5,6 +5,7 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+import prettyFormat from "pretty-format";
 import {
   HIRFunction,
   IdentifierId,
@@ -14,8 +15,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 +34,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 +67,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 +81,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 +90,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);
@@ -96,7 +117,7 @@ function visit(fn: HIRFunction, fbtValues: Set<IdentifierId>): void {
           );
         }
       } else if (
-        isFbtJsxExpression(fbtValues, value) ||
+        isFbtJsxExpression(fbtMacroTags, fbtValues, value) ||
         isFbtJsxChild(fbtValues, lvalue, value)
       ) {
         const fbtScope = lvalue.identifier.scope;
@@ -141,14 +162,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)))
   );
 }
 

Diff for: 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";

Diff for: 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>

Diff for: compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/meta-isms/repro-cx-assigned-to-temporary.js

@@ -0,0 +1,39 @@
+// @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: [{}],
+};

Diff for: compiler/packages/snap/src/compiler.ts

@@ -46,6 +46,7 @@ function makePluginOptions(
   // TODO(@mofeiZ) rewrite snap fixtures to @validatePreserveExistingMemo:false
   let validatePreserveExistingMemoizationGuarantees = false;
   let enableChangeDetectionForDebugging = null;
+  let customMacros = null;
 
   if (firstLine.indexOf("@compilationMode(annotation)") !== -1) {
     assert(
@@ -142,6 +143,18 @@ function makePluginOptions(
     );
   }
 
+  const customMacrosMatch = /@customMacros\(([^)]+)\)/.exec(firstLine);
+  if (
+    customMacrosMatch &&
+    customMacrosMatch.length > 1 &&
+    customMacrosMatch[1].trim().length > 0
+  ) {
+    customMacros = customMacrosMatch[1]
+      .split(" ")
+      .map((s) => s.trim())
+      .filter((s) => s.length > 0);
+  }
+
   let logs: Array<{ filename: string | null; event: LoggerEvent }> = [];
   let logger: Logger | null = null;
   if (firstLine.includes("@logger")) {
@@ -185,6 +198,7 @@ function makePluginOptions(
           },
         ],
       ]),
+      customMacros,
       enableEmitFreeze,
       enableEmitInstrumentForget,
       enableEmitHookGuards,