[compiler][needs cleanup]: stop using function deps in PropagateScope…

…DepsHIR

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

@@ -7,7 +7,6 @@ import {
   Set_union,
   getOrInsertDefault,
 } from '../Utils/utils';
-import {collectOptionalChainSidemap} from './CollectOptionalChainDependencies';
 import {
   BasicBlock,
   BlockId,
@@ -21,7 +20,6 @@ import {
   ReactiveScopeDependency,
   ScopeId,
 } from './HIR';
-import {collectTemporariesSidemap} from './PropagateScopeDependenciesHIR';
 
 /**
  * Helper function for `PropagateScopeDependencies`. Uses control flow graph
@@ -353,15 +351,10 @@ function collectNonNullsInBlocks(
         !fn.env.config.enableTreatFunctionDepsAsConditional
       ) {
         const innerFn = instr.value.loweredFunc;
-        const innerTemporaries = collectTemporariesSidemap(
-          innerFn.func,
-          new Set(),
-        );
-        const innerOptionals = collectOptionalChainSidemap(innerFn.func);
         const innerHoistableMap = collectHoistablePropertyLoads(
           innerFn.func,
-          innerTemporaries,
-          innerOptionals.hoistableObjects,
+          context.temporaries,
+          context.hoistableFromOptionals,
           context.nestedFnImmutableContext ??
             new Set(
               innerFn.func.context

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

@@ -1,4 +1,5 @@
 import {CompilerError} from '..';
+import {getOrInsertDefault} from '../Utils/utils';
 import {assertNonNull} from './CollectHoistablePropertyLoads';
 import {
   BlockId,
@@ -22,25 +23,14 @@ export function collectOptionalChainSidemap(
   fn: HIRFunction,
 ): OptionalChainSidemap {
   const context: OptionalTraversalContext = {
+    currFn: fn,
     blocks: fn.body.blocks,
     seenOptionals: new Set(),
-    processedInstrsInOptional: new Set(),
+    processedInstrsInOptional: new Map(),
     temporariesReadInOptional: new Map(),
     hoistableObjects: new Map(),
   };
-  for (const [_, block] of fn.body.blocks) {
-    if (
-      block.terminal.kind === 'optional' &&
-      !context.seenOptionals.has(block.id)
-    ) {
-      traverseOptionalBlock(
-        block as TBasicBlock<OptionalTerminal>,
-        context,
-        null,
-      );
-    }
-  }
-
+  traverseFunction(fn, context);
   return {
     temporariesReadInOptional: context.temporariesReadInOptional,
     processedInstrsInOptional: context.processedInstrsInOptional,
@@ -97,7 +87,10 @@ export type OptionalChainSidemap = {
    *   $5 = MethodCall $2.$4() <--- here, we want to take a dep on $2 and $4!
    * ```
    */
-  processedInstrsInOptional: ReadonlySet<InstructionId>;
+  processedInstrsInOptional: ReadonlyMap<
+    HIRFunction,
+    ReadonlySet<InstructionId>
+  >;
   /**
    * Records optional chains for which we can safely evaluate non-optional
    * PropertyLoads. e.g. given `a?.b.c`, we can evaluate any load from `a?.b` at
@@ -115,16 +108,47 @@ export type OptionalChainSidemap = {
 };
 
 type OptionalTraversalContext = {
+  currFn: HIRFunction;
   blocks: ReadonlyMap<BlockId, BasicBlock>;
 
   // Track optional blocks to avoid outer calls into nested optionals
   seenOptionals: Set<BlockId>;
 
-  processedInstrsInOptional: Set<InstructionId>;
+  processedInstrsInOptional: Map<HIRFunction, Set<InstructionId>>;
   temporariesReadInOptional: Map<IdentifierId, ReactiveScopeDependency>;
   hoistableObjects: Map<BlockId, ReactiveScopeDependency>;
 };
 
+function traverseFunction(
+  fn: HIRFunction,
+  context: OptionalTraversalContext,
+): void {
+  for (const [_, block] of fn.body.blocks) {
+    for (const instr of block.instructions) {
+      if (
+        instr.value.kind === 'FunctionExpression' ||
+        instr.value.kind === 'ObjectMethod'
+      ) {
+        traverseFunction(instr.value.loweredFunc.func, {
+          ...context,
+          currFn: instr.value.loweredFunc.func,
+          blocks: instr.value.loweredFunc.func.body.blocks,
+          seenOptionals: new Set(),
+        });
+      }
+    }
+    if (
+      block.terminal.kind === 'optional' &&
+      !context.seenOptionals.has(block.id)
+    ) {
+      traverseOptionalBlock(
+        block as TBasicBlock<OptionalTerminal>,
+        context,
+        null,
+      );
+    }
+  }
+}
 /**
  * Match the consequent and alternate blocks of an optional.
  * @returns propertyload computed by the consequent block, or null if the
@@ -369,10 +393,13 @@ function traverseOptionalBlock(
       },
     ],
   };
-  context.processedInstrsInOptional.add(
-    matchConsequentResult.storeLocalInstrId,
+  const processedInstrsInOptional = getOrInsertDefault(
+    context.processedInstrsInOptional,
+    context.currFn,
+    new Set(),
   );
-  context.processedInstrsInOptional.add(test.id);
+  processedInstrsInOptional.add(matchConsequentResult.storeLocalInstrId);
+  processedInstrsInOptional.add(test.id);
   context.temporariesReadInOptional.set(
     matchConsequentResult.consequentId,
     load,

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

@@ -16,6 +16,8 @@ import {
   DeclarationId,
   areEqualPaths,
   IdentifierId,
+  BasicBlock,
+  BlockId,
 } from './HIR';
 import {
   collectHoistablePropertyLoads,
@@ -37,7 +39,11 @@ import {collectOptionalChainSidemap} from './CollectOptionalChainDependencies';
 export function propagateScopeDependenciesHIR(fn: HIRFunction): void {
   const usedOutsideDeclaringScope =
     findTemporariesUsedOutsideDeclaringScope(fn);
-  const temporaries = collectTemporariesSidemap(fn, usedOutsideDeclaringScope);
+  const temporaries = collectTemporariesSidemap(
+    fn,
+    usedOutsideDeclaringScope,
+    new Map(),
+  );
   const {
     temporariesReadInOptional,
     processedInstrsInOptional,
@@ -214,8 +220,9 @@ function findTemporariesUsedOutsideDeclaringScope(
 export function collectTemporariesSidemap(
   fn: HIRFunction,
   usedOutsideDeclaringScope: ReadonlySet<DeclarationId>,
+  temporaries: Map<IdentifierId, ReactiveScopeDependency>,
+  isInnerFn: boolean = false,
 ): ReadonlyMap<IdentifierId, ReactiveScopeDependency> {
-  const temporaries = new Map<IdentifierId, ReactiveScopeDependency>();
   for (const [_, block] of fn.body.blocks) {
     for (const instr of block.instructions) {
       const {value, lvalue} = instr;
@@ -224,23 +231,54 @@ export function collectTemporariesSidemap(
       );
 
       if (value.kind === 'PropertyLoad' && !usedOutside) {
-        const property = getProperty(
-          value.object,
-          value.property,
-          false,
-          temporaries,
-        );
-        temporaries.set(lvalue.identifier.id, property);
+        if (isInnerFn) {
+          const source = temporaries.get(value.object.identifier.id);
+          if (source) {
+            // only add inner function loads if their source is valid
+            const property = getProperty(
+              value.object,
+              value.property,
+              false,
+              temporaries,
+            );
+            temporaries.set(lvalue.identifier.id, property);
+          }
+        } else {
+          const property = getProperty(
+            value.object,
+            value.property,
+            false,
+            temporaries,
+          );
+          temporaries.set(lvalue.identifier.id, property);
+        }
       } else if (
         value.kind === 'LoadLocal' &&
         lvalue.identifier.name == null &&
         value.place.identifier.name !== null &&
         !usedOutside
       ) {
-        temporaries.set(lvalue.identifier.id, {
-          identifier: value.place.identifier,
-          path: [],
-        });
+        if (
+          !isInnerFn ||
+          fn.context.some(
+            context => context.identifier.id === value.place.identifier.id,
+          )
+        ) {
+          temporaries.set(lvalue.identifier.id, {
+            identifier: value.place.identifier,
+            path: [],
+          });
+        }
+      } else if (
+        value.kind === 'FunctionExpression' ||
+        value.kind === 'ObjectMethod'
+      ) {
+        collectTemporariesSidemap(
+          value.loweredFunc.func,
+          usedOutsideDeclaringScope,
+          temporaries,
+          true,
+        );
       }
     }
   }
@@ -310,6 +348,8 @@ class Context {
   #temporaries: ReadonlyMap<IdentifierId, ReactiveScopeDependency>;
   #temporariesUsedOutsideScope: ReadonlySet<DeclarationId>;
 
+  innerFnContext: HIRFunction | null = null;
+
   constructor(
     temporariesUsedOutsideScope: ReadonlySet<DeclarationId>,
     temporaries: ReadonlyMap<IdentifierId, ReactiveScopeDependency>,
@@ -374,6 +414,14 @@ class Context {
 
   // Checks if identifier is a valid dependency in the current scope
   #checkValidDependency(maybeDependency: ReactiveScopeDependency): boolean {
+    if (
+      this.innerFnContext != null &&
+      !this.innerFnContext.context.some(
+        context => context.identifier.id === maybeDependency.identifier.id,
+      )
+    ) {
+      return false;
+    }
     // ref.current access is not a valid dep
     if (
       isUseRefType(maybeDependency.identifier) &&
@@ -575,7 +623,10 @@ function collectDependencies(
   fn: HIRFunction,
   usedOutsideDeclaringScope: ReadonlySet<DeclarationId>,
   temporaries: ReadonlyMap<IdentifierId, ReactiveScopeDependency>,
-  processedInstrsInOptional: ReadonlySet<InstructionId>,
+  processedInstrsInOptional: ReadonlyMap<
+    HIRFunction,
+    ReadonlySet<InstructionId>
+  >,
 ): Map<ReactiveScope, Array<ReactiveScopeDependency>> {
   const context = new Context(usedOutsideDeclaringScope, temporaries);
 
@@ -594,8 +645,13 @@ function collectDependencies(
   }
 
   const scopeTraversal = new ScopeBlockTraversal();
+  // TODO: make this less hacky
+  const st: Array<[HIRFunction, BlockId, BasicBlock]> = [...fn.body.blocks].map(
+    ([id, block]) => [fn, id, block],
+  );
 
-  for (const [blockId, block] of fn.body.blocks) {
+  while (st.length != 0) {
+    const [currFn, blockId, block] = st.shift()!;
     scopeTraversal.recordScopes(block);
     const scopeBlockInfo = scopeTraversal.blockInfos.get(blockId);
     if (scopeBlockInfo?.kind === 'begin') {
@@ -604,6 +660,9 @@ function collectDependencies(
       context.exitScope(scopeBlockInfo.scope, scopeBlockInfo?.pruned);
     }
 
+    // TODO: make this less hacky
+    context.innerFnContext = currFn === fn ? null : currFn;
+
     // Record referenced optional chains in phis
     for (const phi of block.phis) {
       for (const operand of phi.operands) {
@@ -614,16 +673,37 @@ function collectDependencies(
       }
     }
     for (const instr of block.instructions) {
-      if (!processedInstrsInOptional.has(instr.id)) {
+      if (
+        instr.value.kind === 'FunctionExpression' ||
+        instr.value.kind === 'ObjectMethod'
+      ) {
+        /**
+         * Push the outermost nested fn context, as these are guaranteed to reference
+         * component / hook identifiers (i.e. eligible dependencies)
+         */
+        const innerFn = instr.value.loweredFunc.func;
+        const outermostNestedFnContext = currFn === fn ? innerFn : currFn;
+        const x: Array<[HIRFunction, BlockId, BasicBlock]> = [
+          ...innerFn.body.blocks.entries(),
+        ].map(([id, block]) => [outermostNestedFnContext, id, block]);
+        st.unshift(...x);
+
+        context.declare(instr.lvalue.identifier, {
+          id: instr.id,
+          scope: context.currentScope,
+        });
+      } else if (!processedInstrsInOptional.get(currFn)?.has(instr.id)) {
+        // instruction ids are function-relative
         handleInstruction(instr, context);
       }
     }
 
-    if (!processedInstrsInOptional.has(block.terminal.id)) {
+    if (!processedInstrsInOptional.get(currFn)?.has(block.terminal.id)) {
       for (const place of eachTerminalOperand(block.terminal)) {
         context.visitOperand(place);
       }
     }
   }
+
   return context.deps;
 }

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/capturing-func-mutate-2.expect.md

@@ -26,29 +26,20 @@ export const FIXTURE_ENTRYPOINT = {
 ```javascript
 import { c as _c } from "react/compiler-runtime";
 function component(a, b) {
-  const $ = _c(5);
-  let t0;
-  if ($[0] !== b) {
-    t0 = { b };
-    $[0] = b;
-    $[1] = t0;
-  } else {
-    t0 = $[1];
-  }
-  const y = t0;
+  const $ = _c(2);
+  const y = { b };
   let z;
-  if ($[2] !== a || $[3] !== y) {
+  if ($[0] !== a) {
     z = { a };
     const x = function () {
       z.a = 2;
     };
 
     x();
-    $[2] = a;
-    $[3] = y;
-    $[4] = z;
+    $[0] = a;
+    $[1] = z;
   } else {
-    z = $[4];
+    z = $[1];
   }
   return z;
 }