Update on "[compiler] Add lowerContextAccess pass"

*This is only for internal profiling, not intended to ship.*

This pass is intended to be used with #30407.

This pass synthesizes selector functions by collecting immediately
destructured context acesses. We bailout for other types of context
access.

This pass lowers context access to use a selector function by passing
the synthesized selector function as the second argument.

[ghstack-poisoned]

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

@@ -56,7 +56,7 @@ import {
   flattenReactiveLoops,
   flattenScopesWithHooksOrUse,
   inferReactiveScopeVariables,
-  memoizeFbtOperandsInSameScope,
+  memoizeFbtAndMacroOperandsInSameScope,
   mergeOverlappingReactiveScopes,
   mergeReactiveScopesThatInvalidateTogether,
   promoteUsedTemporaries,
@@ -248,8 +248,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});
   }
 
@@ -267,13 +274,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/Inference/InferReferenceEffects.ts

@@ -26,6 +26,7 @@ import {
   Type,
   ValueKind,
   ValueReason,
+  getHookKind,
   isArrayType,
   isMutableEffect,
   isObjectType,
@@ -48,7 +49,6 @@ import {
   eachTerminalSuccessor,
 } from '../HIR/visitors';
 import {assertExhaustive} from '../Utils/utils';
-import {isEffectHook} from '../Validation/ValidateMemoizedEffectDependencies';
 
 const UndefinedValue: InstructionValue = {
   kind: 'Primitive',
@@ -1151,7 +1151,7 @@ function inferBlock(
             );
             functionEffects.push(
               ...propEffects.filter(
-                propEffect => propEffect.kind !== 'GlobalMutation',
+                effect => !isEffectSafeOutsideRender(effect),
               ),
             );
           }
@@ -1330,7 +1330,7 @@ function inferBlock(
                 context: new Set(),
               };
         let hasCaptureArgument = false;
-        let isUseEffect = isEffectHook(instrValue.callee.identifier);
+        let isHook = getHookKind(env, instrValue.callee.identifier) != null;
         for (let i = 0; i < instrValue.args.length; i++) {
           const argumentEffects: Array<FunctionEffect> = [];
           const arg = instrValue.args[i];
@@ -1356,8 +1356,7 @@ function inferBlock(
            */
           functionEffects.push(
             ...argumentEffects.filter(
-              argEffect =>
-                !isUseEffect || i !== 0 || argEffect.kind !== 'GlobalMutation',
+              argEffect => !isHook || !isEffectSafeOutsideRender(argEffect),
             ),
           );
           hasCaptureArgument ||= place.effect === Effect.Capture;
@@ -1455,7 +1454,7 @@ function inferBlock(
         const effects =
           signature !== null ? getFunctionEffects(instrValue, signature) : null;
         let hasCaptureArgument = false;
-        let isUseEffect = isEffectHook(instrValue.property.identifier);
+        let isHook = getHookKind(env, instrValue.property.identifier) != null;
         for (let i = 0; i < instrValue.args.length; i++) {
           const argumentEffects: Array<FunctionEffect> = [];
           const arg = instrValue.args[i];
@@ -1485,8 +1484,7 @@ function inferBlock(
            */
           functionEffects.push(
             ...argumentEffects.filter(
-              argEffect =>
-                !isUseEffect || i !== 0 || argEffect.kind !== 'GlobalMutation',
+              argEffect => !isHook || !isEffectSafeOutsideRender(argEffect),
             ),
           );
           hasCaptureArgument ||= place.effect === Effect.Capture;
@@ -2010,11 +2008,15 @@ function inferBlock(
     } else {
       effect = Effect.Read;
     }
+    const propEffects: Array<FunctionEffect> = [];
     state.referenceAndRecordEffects(
       operand,
       effect,
       ValueReason.Other,
-      functionEffects,
+      propEffects,
+    );
+    functionEffects.push(
+      ...propEffects.filter(effect => !isEffectSafeOutsideRender(effect)),
     );
   }
 }
@@ -2128,6 +2130,10 @@ function areArgumentsImmutableAndNonMutating(
   return true;
 }
 
+function isEffectSafeOutsideRender(effect: FunctionEffect): boolean {
+  return effect.kind === 'GlobalMutation';
+}
+
 function getWriteErrorReason(abstractValue: AbstractValue): string {
   if (abstractValue.reason.has(ValueReason.Global)) {
     return 'Writing to a variable defined outside a component or hook is not allowed. Consider using an effect';

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.reassign-global-fn-arg.expect.md

@@ -0,0 +1,36 @@
+
+## Input
+
+```javascript
+let b = 1;
+
+export default function MyApp() {
+  const fn = () => {
+    b = 2;
+  };
+  return foo(fn);
+}
+
+function foo(fn) {}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: MyApp,
+  params: [],
+};
+
+```
+
+
+## Error
+
+```
+  3 | export default function MyApp() {
+  4 |   const fn = () => {
+> 5 |     b = 2;
+    |     ^ InvalidReact: Unexpected reassignment of a variable which was defined outside of the component. Components and hooks should be pure and side-effect free, but variable reassignment is a form of side-effect. If this variable is used in rendering, use useState instead. (https://react.dev/reference/rules/components-and-hooks-must-be-pure#side-effects-must-run-outside-of-render) (5:5)
+  6 |   };
+  7 |   return foo(fn);
+  8 | }
+```
+          
+      

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.reassign-global-fn-arg.js

@@ -0,0 +1,15 @@
+let b = 1;
+
+export default function MyApp() {
+  const fn = () => {
+    b = 2;
+  };
+  return foo(fn);
+}
+
+function foo(fn) {}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: MyApp,
+  params: [],
+};

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/idx-no-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/error.idx-outlining.js → compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/idx-no-outlining.js

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

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/reassign-global-hook-arg.expect.md

@@ -0,0 +1,46 @@
+
+## Input
+
+```javascript
+let b = 1;
+
+export default function MyApp() {
+  const fn = () => {
+    b = 2;
+  };
+  return useFoo(fn);
+}
+
+function useFoo(fn) {}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: MyApp,
+  params: [],
+};
+
+```
+
+## Code
+
+```javascript
+let b = 1;
+
+export default function MyApp() {
+  const fn = _temp;
+  return useFoo(fn);
+}
+function _temp() {
+  b = 2;
+}
+
+function useFoo(fn) {}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: MyApp,
+  params: [],
+};
+
+```
+      
+### Eval output
+(kind: ok)