[compiler] Fix issue with macro arguments being outlined

Summary:
Fixes issue documented by #30435. We change the pipeline order so that outlining comes after tracking macro operands, and any function that is referenced in a macro will now not be outlined.

ghstack-source-id: f731ad6
Pull Request resolved: #30587

Author: mvitousek

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

@@ -56,7 +56,7 @@ import {
   flattenReactiveLoops,
   flattenScopesWithHooksOrUse,
   inferReactiveScopeVariables,
-  memoizeFbtOperandsInSameScope,
+  memoizeFbtAndMacroOperandsInSameScope,
   mergeOverlappingReactiveScopes,
   mergeReactiveScopesThatInvalidateTogether,
   promoteUsedTemporaries,
@@ -243,8 +243,15 @@ function* runWithEnvironment(
   inferReactiveScopeVariables(hir);
   yield log({kind: 'hir', name: 'InferReactiveScopeVariables', value: hir});
 
+  const fbtOperands = memoizeFbtAndMacroOperandsInSameScope(hir);
+  yield log({
+    kind: 'hir',
+    name: 'MemoizeFbtAndMacroOperandsInSameScope',
+    value: hir,
+  });
+
   if (env.config.enableFunctionOutlining) {
-    outlineFunctions(hir);
+    outlineFunctions(hir, fbtOperands);
     yield log({kind: 'hir', name: 'OutlineFunctions', value: hir});
   }
 
@@ -262,13 +269,6 @@ function* runWithEnvironment(
     value: hir,
   });
 
-  const fbtOperands = memoizeFbtOperandsInSameScope(hir);
-  yield log({
-    kind: 'hir',
-    name: 'MemoizeFbtAndMacroOperandsInSameScope',
-    value: hir,
-  });
-
   if (env.config.enableReactiveScopesInHIR) {
     pruneUnusedLabelsHIR(hir);
     yield log({

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

@@ -485,6 +485,11 @@ export function parseConfigPragma(pragma: string): EnvironmentConfig {
       continue;
     }
 
+    if (key === 'customMacros' && val) {
+      maybeConfig[key] = [val];
+      continue;
+    }
+
     if (typeof defaultConfig[key as keyof EnvironmentConfig] !== 'boolean') {
       // skip parsing non-boolean properties
       continue;

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

@@ -5,27 +5,30 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-import {HIRFunction} from '../HIR';
+import {HIRFunction, IdentifierId} from '../HIR';
 
-export function outlineFunctions(fn: HIRFunction): void {
+export function outlineFunctions(
+  fn: HIRFunction,
+  fbtOperands: Set<IdentifierId>,
+): void {
   for (const [, block] of fn.body.blocks) {
     for (const instr of block.instructions) {
-      const {value} = instr;
+      const {value, lvalue} = instr;
 
       if (
         value.kind === 'FunctionExpression' ||
         value.kind === 'ObjectMethod'
       ) {
         // Recurse in case there are inner functions which can be outlined
-        outlineFunctions(value.loweredFunc.func);
+        outlineFunctions(value.loweredFunc.func, fbtOperands);
       }
-
       if (
         value.kind === 'FunctionExpression' &&
         value.loweredFunc.dependencies.length === 0 &&
         value.loweredFunc.func.context.length === 0 &&
         // TODO: handle outlining named functions
-        value.loweredFunc.func.id === null
+        value.loweredFunc.func.id === null &&
+        !fbtOperands.has(lvalue.identifier.id)
       ) {
         const loweredFunc = value.loweredFunc.func;
 

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

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

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

@@ -0,0 +1,64 @@
+
+## Input
+
+```javascript
+// @customMacros(idx)
+import idx from 'idx';
+
+function Component(props) {
+  // the lambda should not be outlined
+  const groupName = idx(props, _ => _.group.label);
+  return <div>{groupName}</div>;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{}],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime"; // @customMacros(idx)
+
+function Component(props) {
+  var _ref2;
+  const $ = _c(4);
+  let t0;
+  if ($[0] !== props) {
+    var _ref;
+
+    t0 =
+      (_ref = props) != null
+        ? (_ref = _ref.group) != null
+          ? _ref.label
+          : _ref
+        : _ref;
+    $[0] = props;
+    $[1] = t0;
+  } else {
+    t0 = $[1];
+  }
+  const groupName = t0;
+  let t1;
+  if ($[2] !== groupName) {
+    t1 = <div>{groupName}</div>;
+    $[2] = groupName;
+    $[3] = t1;
+  } else {
+    t1 = $[3];
+  }
+  return t1;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{}],
+};
+
+```
+      
+### Eval output
+(kind: ok) <div></div>

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

@@ -1,3 +1,4 @@
+// @customMacros(idx)
 import idx from 'idx';
 
 function Component(props) {