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]

.eslintrc.js

@@ -496,6 +496,7 @@ module.exports = {
         __IS_CHROME__: 'readonly',
         __IS_FIREFOX__: 'readonly',
         __IS_EDGE__: 'readonly',
+        __IS_NATIVE__: 'readonly',
         __IS_INTERNAL_VERSION__: 'readonly',
       },
     },

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

@@ -128,11 +128,11 @@ export function* run(
     code,
     useMemoCacheIdentifier,
   );
-  yield {
+  yield log({
     kind: 'debug',
     name: 'EnvironmentConfig',
     value: prettyFormat(env.config),
-  };
+  });
   const ast = yield* runWithEnvironment(func, env);
   return ast;
 }

compiler/packages/babel-plugin-react-compiler/src/Inference/DropManualMemoization.ts

@@ -335,6 +335,7 @@ function extractManualMemoizationArgs(
 export function dropManualMemoization(func: HIRFunction): void {
   const isValidationEnabled =
     func.env.config.validatePreserveExistingMemoizationGuarantees ||
+    func.env.config.validateNoSetStateInRender ||
     func.env.config.enablePreserveExistingMemoizationGuarantees;
   const sidemap: IdentifierSidemap = {
     functions: new Map(),

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

@@ -21,17 +21,16 @@ import {
   PropertyLoad,
   isUseContextHookType,
   makeBlockId,
-  makeIdentifierId,
-  makeIdentifierName,
   makeInstructionId,
   makeTemporary,
-  makeType,
   markInstructionIds,
   mergeConsecutiveBlocks,
+  promoteTemporary,
   removeUnnecessaryTryCatch,
   reversePostorderBlocks,
 } from '../HIR';
 import {
+  createTemporaryPlace,
   removeDeadDoWhileStatements,
   removeUnreachableForUpdates,
 } from '../HIR/HIRBuilder';
@@ -66,7 +65,7 @@ export function lowerContextAccess(fn: HIRFunction): void {
 
       const keys = getContextKeys(value);
       if (keys === null) {
-        continue;
+        return;
       }
 
       if (contextKeys.has(destructureId)) {
@@ -83,8 +82,10 @@ export function lowerContextAccess(fn: HIRFunction): void {
 
   if (contextAccess.size > 0) {
     for (const [, block] of fn.body.blocks) {
-      const nextInstructions: Array<Instruction> = [];
-      for (const instr of block.instructions) {
+      let nextInstructions: Array<Instruction> | null = null;
+
+      for (let i = 0; i < block.instructions.length; i++) {
+        const instr = block.instructions[i];
         const {lvalue, value} = instr;
         if (
           value.kind === 'CallExpression' &&
@@ -93,15 +94,22 @@ export function lowerContextAccess(fn: HIRFunction): void {
         ) {
           const keys = contextKeys.get(lvalue.identifier.id)!;
           const selectorFnInstr = emitSelectorFn(fn.env, keys);
+          if (nextInstructions === null) {
+            nextInstructions = block.instructions.slice(0, i);
+          }
           nextInstructions.push(selectorFnInstr);
 
           const selectorFn = selectorFnInstr.lvalue;
           value.args.push(selectorFn);
         }
 
-        nextInstructions.push(instr);
+        if (nextInstructions) {
+          nextInstructions.push(instr);
+        }
+      }
+      if (nextInstructions) {
+        block.instructions = nextInstructions;
       }
-      block.instructions = nextInstructions;
     }
     markInstructionIds(fn.body);
   }
@@ -113,26 +121,18 @@ function getContextKeys(value: Destructure): Array<string> | null {
 
   switch (pattern.kind) {
     case 'ArrayPattern': {
-      for (const place of pattern.items) {
-        if (place.kind !== 'Identifier') {
-          return null;
-        }
-
-        if (place.identifier.name === null) {
-          return null;
-        }
-
-        keys.push(place.identifier.name.value);
-      }
-      return keys;
+      return null;
     }
 
     case 'ObjectPattern': {
       for (const place of pattern.properties) {
+        debugger;
         if (
           place.kind !== 'ObjectProperty' ||
           place.type !== 'property' ||
-          place.key.kind !== 'identifier'
+          place.key.kind !== 'identifier' ||
+          place.place.identifier.name === null ||
+          place.place.identifier.name.kind !== 'named'
         ) {
           return null;
         }
@@ -153,13 +153,7 @@ function emitPropertyLoad(
     place: obj,
     loc: GeneratedSource,
   };
-  const object: Place = {
-    kind: 'Identifier',
-    identifier: makeTemporary(env.nextIdentifierId, GeneratedSource),
-    effect: Effect.Unknown,
-    reactive: false,
-    loc: GeneratedSource,
-  };
+  const object: Place = createTemporaryPlace(env, GeneratedSource);
   const loadLocalInstr: Instruction = {
     lvalue: object,
     value: loadObj,
@@ -173,13 +167,7 @@ function emitPropertyLoad(
     property,
     loc: GeneratedSource,
   };
-  const element: Place = {
-    kind: 'Identifier',
-    identifier: makeTemporary(env.nextIdentifierId, GeneratedSource),
-    effect: Effect.Unknown,
-    reactive: false,
-    loc: GeneratedSource,
-  };
+  const element: Place = createTemporaryPlace(env, GeneratedSource);
   const loadPropInstr: Instruction = {
     lvalue: element,
     value: loadProp,
@@ -193,20 +181,8 @@ function emitPropertyLoad(
 }
 
 function emitSelectorFn(env: Environment, keys: Array<string>): Instruction {
-  const obj: Place = {
-    kind: 'Identifier',
-    identifier: {
-      id: makeIdentifierId(env.nextIdentifierId),
-      name: makeIdentifierName('c'),
-      mutableRange: {start: makeInstructionId(0), end: makeInstructionId(0)},
-      scope: null,
-      type: makeType(),
-      loc: GeneratedSource,
-    },
-    effect: Effect.Unknown,
-    reactive: false,
-    loc: GeneratedSource,
-  };
+  const obj: Place = createTemporaryPlace(env, GeneratedSource);
+  promoteTemporary(obj.identifier);
   const instr: Array<Instruction> = [];
   const elements = [];
   for (const key of keys) {
@@ -251,11 +227,7 @@ function emitSelectorFn(env: Environment, keys: Array<string>): Instruction {
   };
 
   reversePostorderBlocks(fn.body);
-  removeUnreachableForUpdates(fn.body);
-  removeDeadDoWhileStatements(fn.body);
-  removeUnnecessaryTryCatch(fn.body);
   markInstructionIds(fn.body);
-  mergeConsecutiveBlocks(fn);
   enterSSA(fn);
   inferTypes(fn);
 

compiler/packages/babel-plugin-react-compiler/src/Validation/ValidateLocalsNotReassignedAfterRender.ts

@@ -8,9 +8,11 @@
 import {CompilerError, Effect} from '..';
 import {HIRFunction, IdentifierId, Place} from '../HIR';
 import {
+  eachInstructionLValue,
   eachInstructionValueOperand,
   eachTerminalOperand,
 } from '../HIR/visitors';
+import {getFunctionCallSignature} from '../Inference/InferReferenceEffects';
 
 /**
  * Validates that local variables cannot be reassigned after render.
@@ -131,23 +133,49 @@ function getContextReassignment(
           break;
         }
         default: {
-          for (const operand of eachInstructionValueOperand(value)) {
+          let operands = eachInstructionValueOperand(value);
+          // If we're calling a function that doesn't let its arguments escape, only test the callee
+          if (value.kind === 'CallExpression') {
+            const signature = getFunctionCallSignature(
+              fn.env,
+              value.callee.identifier.type,
+            );
+            if (signature?.noAlias) {
+              operands = [value.callee];
+            }
+          } else if (value.kind === 'MethodCall') {
+            const signature = getFunctionCallSignature(
+              fn.env,
+              value.property.identifier.type,
+            );
+            if (signature?.noAlias) {
+              operands = [value.receiver, value.property];
+            }
+          }
+          for (const operand of operands) {
             CompilerError.invariant(operand.effect !== Effect.Unknown, {
               reason: `Expected effects to be inferred prior to ValidateLocalsNotReassignedAfterRender`,
               loc: operand.loc,
             });
             const reassignment = reassigningFunctions.get(
               operand.identifier.id,
             );
-            if (
-              reassignment !== undefined &&
-              operand.effect === Effect.Freeze
-            ) {
+            if (reassignment !== undefined) {
               /*
                * Functions that reassign local variables are inherently mutable and are unsafe to pass
                * to a place that expects a frozen value. Propagate the reassignment upward.
                */
-              return reassignment;
+              if (operand.effect === Effect.Freeze) {
+                return reassignment;
+              } else {
+                /*
+                 * If the operand is not frozen but it does reassign, then the lvalues
+                 * of the instruction could also be reassigning
+                 */
+                for (const lval of eachInstructionLValue(instr)) {
+                  reassigningFunctions.set(lval.identifier.id, reassignment);
+                }
+              }
             }
           }
           break;