[compiler] Allow macro methods

Summary:
Builds support for macros that are invoked as methods rather than just function calls or jsx.

We now record macros as a schema that represents arbitrary member expressions including wildcards (so we can support, e.g., myMacro.*.foo.bar). When examining PropertyLoads in the macro memoization stage, we build up a map of partially-satisfied macro patterns until we determine that the pattern has been fully satisfied, at which point we treat the result of the PropertyLoad as a macro value.

ghstack-source-id: d78d9ba7041968c861ffa110fb7882b339a0e257
Pull Request resolved: #30589

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

@@ -67,6 +67,20 @@ export const InstrumentationSchema = z
 
 export type ExternalFunction = z.infer<typeof ExternalFunctionSchema>;
 
+export const MacroMethodSchema = z.union([
+  z.object({type: z.literal('wildcard')}),
+  z.object({type: z.literal('name'), name: z.string()}),
+]);
+
+// Would like to change this to drop the string option, but breaks compatibility with existing configs
+export const MacroSchema = z.union([
+  z.string(),
+  z.tuple([z.string(), z.array(MacroMethodSchema)]),
+]);
+
+export type Macro = z.infer<typeof MacroSchema>;
+export type MacroMethod = z.infer<typeof MacroMethodSchema>;
+
 const HookSchema = z.object({
   /*
    * The effect of arguments to this hook. Describes whether the hook may or may
@@ -133,7 +147,7 @@ const EnvironmentConfigSchema = z.object({
    * 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),
+  customMacros: z.nullable(z.array(MacroSchema)).default(null),
 
   /**
    * Enable a check that resets the memoization cache when the source code of the file changes.
@@ -490,7 +504,19 @@ export function parseConfigPragma(pragma: string): EnvironmentConfig {
     }
 
     if (key === 'customMacros' && val) {
-      maybeConfig[key] = [val];
+      const valSplit = val.split('.');
+      if (valSplit.length > 0) {
+        const props = [];
+        for (const elt of valSplit.slice(1)) {
+          if (elt === '*') {
+            props.push({type: 'wildcard'});
+          } else if (elt.length > 0) {
+            props.push({type: 'name', name: elt});
+          }
+        }
+        console.log([valSplit[0], props.map(x => x.name ?? '*').join('.')]);
+        maybeConfig[key] = [[valSplit[0], props]];
+      }
       continue;
     }
 

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

@@ -13,6 +13,7 @@ import {
   Place,
   ReactiveValue,
 } from '../HIR';
+import {Macro, MacroMethod} from '../HIR/Environment';
 import {eachReactiveValueOperand} from './visitors';
 
 /**
@@ -43,15 +44,17 @@ import {eachReactiveValueOperand} from './visitors';
 export function memoizeFbtAndMacroOperandsInSameScope(
   fn: HIRFunction,
 ): Set<IdentifierId> {
-  const fbtMacroTags = new Set([
-    ...FBT_TAGS,
+  const fbtMacroTags = new Set<Macro>([
+    ...Array.from(FBT_TAGS).map((tag): Macro => [tag, []]),
     ...(fn.env.config.customMacros ?? []),
   ]);
   const fbtValues: Set<IdentifierId> = new Set();
+  const macroMethods = new Map<IdentifierId, Array<Array<MacroMethod>>>();
   while (true) {
-    let size = fbtValues.size;
-    visit(fn, fbtMacroTags, fbtValues);
-    if (size === fbtValues.size) {
+    let vsize = fbtValues.size;
+    let msize = macroMethods.size;
+    visit(fn, fbtMacroTags, fbtValues, macroMethods);
+    if (vsize === fbtValues.size && msize === macroMethods.size) {
       break;
     }
   }
@@ -71,8 +74,9 @@ export const SINGLE_CHILD_FBT_TAGS: Set<string> = new Set([
 
 function visit(
   fn: HIRFunction,
-  fbtMacroTags: Set<string>,
+  fbtMacroTags: Set<Macro>,
   fbtValues: Set<IdentifierId>,
+  macroMethods: Map<IdentifierId, Array<Array<MacroMethod>>>,
 ): void {
   for (const [, block] of fn.body.blocks) {
     for (const instruction of block.instructions) {
@@ -83,7 +87,7 @@ function visit(
       if (
         value.kind === 'Primitive' &&
         typeof value.value === 'string' &&
-        fbtMacroTags.has(value.value)
+        matchesExactTag(value.value, fbtMacroTags)
       ) {
         /*
          * We don't distinguish between tag names and strings, so record
@@ -92,10 +96,38 @@ function visit(
         fbtValues.add(lvalue.identifier.id);
       } else if (
         value.kind === 'LoadGlobal' &&
-        fbtMacroTags.has(value.binding.name)
+        matchesExactTag(value.binding.name, fbtMacroTags)
       ) {
         // Record references to `fbt` as a global
         fbtValues.add(lvalue.identifier.id);
+      } else if (
+        value.kind === 'LoadGlobal' &&
+        matchTagRoot(value.binding.name, fbtMacroTags) !== null
+      ) {
+        const methods = matchTagRoot(value.binding.name, fbtMacroTags)!;
+        macroMethods.set(lvalue.identifier.id, methods);
+      } else if (
+        value.kind === 'PropertyLoad' &&
+        macroMethods.has(value.object.identifier.id)
+      ) {
+        const methods = macroMethods.get(value.object.identifier.id)!;
+        const newMethods = [];
+        for (const method of methods) {
+          if (
+            method.length > 0 &&
+            (method[0].type === 'wildcard' ||
+              (method[0].type === 'name' && method[0].name === value.property))
+          ) {
+            if (method.length > 1) {
+              newMethods.push(method.slice(1));
+            } else {
+              fbtValues.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) {
@@ -167,25 +199,57 @@ function visit(
   }
 }
 
+function matchesExactTag(s: string, tags: Set<Macro>): boolean {
+  return Array.from(tags).some(macro =>
+    typeof macro === 'string'
+      ? s === macro
+      : macro[1].length === 0 && macro[0] === s,
+  );
+}
+
+function matchTagRoot(
+  s: string,
+  tags: Set<Macro>,
+): Array<Array<MacroMethod>> | null {
+  const methods: Array<Array<MacroMethod>> = [];
+  for (const macro of tags) {
+    if (typeof macro === 'string') {
+      continue;
+    }
+    const [tag, rest] = macro;
+    if (tag === s && rest.length > 0) {
+      methods.push(rest);
+    }
+  }
+  if (methods.length > 0) {
+    return methods;
+  } else {
+    return null;
+  }
+}
+
 function isFbtCallExpression(
   fbtValues: Set<IdentifierId>,
   value: ReactiveValue,
 ): boolean {
   return (
-    value.kind === 'CallExpression' && fbtValues.has(value.callee.identifier.id)
+    (value.kind === 'CallExpression' &&
+      fbtValues.has(value.callee.identifier.id)) ||
+    (value.kind === 'MethodCall' && fbtValues.has(value.property.identifier.id))
   );
 }
 
 function isFbtJsxExpression(
-  fbtMacroTags: Set<string>,
+  fbtMacroTags: Set<Macro>,
   fbtValues: Set<IdentifierId>,
   value: ReactiveValue,
 ): boolean {
   return (
     value.kind === 'JsxExpression' &&
     ((value.tag.kind === 'Identifier' &&
       fbtValues.has(value.tag.identifier.id)) ||
-      (value.tag.kind === 'BuiltinTag' && fbtMacroTags.has(value.tag.name)))
+      (value.tag.kind === 'BuiltinTag' &&
+        matchesExactTag(value.tag.name, fbtMacroTags)))
   );
 }
 

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

@@ -0,0 +1,122 @@
+
+## Input
+
+```javascript
+// @customMacros(idx.*.b)
+
+function Component(props) {
+  // outlined
+  const groupName1 = idx(props, _ => _.group.label);
+  // outlined
+  const groupName2 = idx.a(props, _ => _.group.label);
+  // not outlined
+  const groupName3 = idx.a.b(props, _ => _.group.label);
+  // not outlined
+  const groupName4 = idx.hello_world.b(props, _ => _.group.label);
+  // outlined
+  const groupName5 = idx.hello_world.b.c(props, _ => _.group.label);
+  return (
+    <div>
+      {groupName1}
+      {groupName2}
+      {groupName3}
+      {groupName4}
+      {groupName5}
+    </div>
+  );
+}
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime"; // @customMacros(idx.*.b)
+
+function Component(props) {
+  const $ = _c(16);
+  let t0;
+  if ($[0] !== props) {
+    t0 = idx(props, _temp);
+    $[0] = props;
+    $[1] = t0;
+  } else {
+    t0 = $[1];
+  }
+  const groupName1 = t0;
+  let t1;
+  if ($[2] !== props) {
+    t1 = idx.a(props, _temp2);
+    $[2] = props;
+    $[3] = t1;
+  } else {
+    t1 = $[3];
+  }
+  const groupName2 = t1;
+  let t2;
+  if ($[4] !== props) {
+    t2 = idx.a.b(props, (__1) => __1.group.label);
+    $[4] = props;
+    $[5] = t2;
+  } else {
+    t2 = $[5];
+  }
+  const groupName3 = t2;
+  let t3;
+  if ($[6] !== props) {
+    t3 = idx.hello_world.b(props, (__2) => __2.group.label);
+    $[6] = props;
+    $[7] = t3;
+  } else {
+    t3 = $[7];
+  }
+  const groupName4 = t3;
+  let t4;
+  if ($[8] !== props) {
+    t4 = idx.hello_world.b.c(props, _temp3);
+    $[8] = props;
+    $[9] = t4;
+  } else {
+    t4 = $[9];
+  }
+  const groupName5 = t4;
+  let t5;
+  if (
+    $[10] !== groupName1 ||
+    $[11] !== groupName2 ||
+    $[12] !== groupName3 ||
+    $[13] !== groupName4 ||
+    $[14] !== groupName5
+  ) {
+    t5 = (
+      <div>
+        {groupName1}
+        {groupName2}
+        {groupName3}
+        {groupName4}
+        {groupName5}
+      </div>
+    );
+    $[10] = groupName1;
+    $[11] = groupName2;
+    $[12] = groupName3;
+    $[13] = groupName4;
+    $[14] = groupName5;
+    $[15] = t5;
+  } else {
+    t5 = $[15];
+  }
+  return t5;
+}
+function _temp3(__3) {
+  return __3.group.label;
+}
+function _temp2(__0) {
+  return __0.group.label;
+}
+function _temp(_) {
+  return _.group.label;
+}
+
+```
+      

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

@@ -0,0 +1,23 @@
+// @customMacros(idx.*.b)
+
+function Component(props) {
+  // outlined
+  const groupName1 = idx(props, _ => _.group.label);
+  // outlined
+  const groupName2 = idx.a(props, _ => _.group.label);
+  // not outlined
+  const groupName3 = idx.a.b(props, _ => _.group.label);
+  // not outlined
+  const groupName4 = idx.hello_world.b(props, _ => _.group.label);
+  // outlined
+  const groupName5 = idx.hello_world.b.c(props, _ => _.group.label);
+  return (
+    <div>
+      {groupName1}
+      {groupName2}
+      {groupName3}
+      {groupName4}
+      {groupName5}
+    </div>
+  );
+}