compiler: More precise escape analysis

Our current escape analysis currently works by determing values which are returned, finding the scope those values are produced in, and then walking the dependencies of those scopes. Notably, when we traverse into scopes, we _force_ their values to be memoized — even in cases where we can prove it isn't necessary.

The idea of this PR is to change escape analysis to be purely based on the flow of values. So rather than assume all scope deps need to be force-memoized, we look at how those values are used.

Note: the main motivation for this PR is the follow-up, which allows us to infer dependencies for pruned scopes. Without this change, inferring deps for pruned scopes causes the escape analysis pass to consider values as escaping which shouldn't be.

ghstack-source-id: d44bd38fd4
Pull Request resolved: https://github.com/facebook/react/pull/29782

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

@@ -65,12 +65,19 @@ export type ReactiveScopeBlock = {
   instructions: ReactiveBlock;
 };
 
+export type PrunedReactiveScopeBlock = {
+  kind: "pruned-scope";
+  scope: ReactiveScope;
+  instructions: ReactiveBlock;
+};
+
 export type ReactiveBlock = Array<ReactiveStatement>;
 
 export type ReactiveStatement =
   | ReactiveInstructionStatement
   | ReactiveTerminalStatement
-  | ReactiveScopeBlock;
+  | ReactiveScopeBlock
+  | PrunedReactiveScopeBlock;
 
 export type ReactiveInstructionStatement = {
   kind: "instruction";

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

@@ -183,6 +183,7 @@ function visitBlock(context: Context, block: ReactiveBlock): void {
         context.append(stmt, stmt.label);
         break;
       }
+      case "pruned-scope":
       case "scope": {
         CompilerError.invariant(false, {
           reason: "Expected the function to not have scopes already assigned",

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

@@ -400,6 +400,11 @@ function codegenBlockNoReset(
         }
         break;
       }
+      case "pruned-scope": {
+        const scopeBlock = codegenBlockNoReset(cx, item.instructions);
+        statements.push(...scopeBlock.body);
+        break;
+      }
       case "scope": {
         const temp = new Map(cx.temp);
         codegenReactiveScope(cx, statements, item.scope, item.instructions);

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

@@ -34,7 +34,14 @@ class Transform extends ReactiveFunctionTransform<boolean> {
   ): Transformed<ReactiveStatement> {
     this.visitScope(scope, isWithinLoop);
     if (isWithinLoop) {
-      return { kind: "replace-many", value: scope.instructions };
+      return {
+        kind: "replace",
+        value: {
+          kind: "pruned-scope",
+          scope: scope.scope,
+          instructions: scope.instructions,
+        },
+      };
     } else {
       return { kind: "keep" };
     }

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

@@ -66,7 +66,14 @@ class Transform extends ReactiveFunctionTransform<State> {
     this.visitScope(scope, innerState);
     outerState.hasHook ||= innerState.hasHook;
     if (innerState.hasHook) {
-      return { kind: "replace-many", value: scope.instructions };
+      return {
+        kind: "replace",
+        value: {
+          kind: "pruned-scope",
+          scope: scope.scope,
+          instructions: scope.instructions,
+        },
+      };
     } else {
       return { kind: "keep" };
     }

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

@@ -174,6 +174,16 @@ class Transform extends ReactiveFunctionTransform<ReactiveScopeDependencies | nu
           }
           break;
         }
+        case "pruned-scope": {
+          // For now we don't merge across pruned scopes
+          if (current !== null) {
+            log(
+              `Reset scope @${current.block.scope.id} from pruned scope @${instr.scope.id}`
+            );
+            reset();
+          }
+          break;
+        }
         case "instruction": {
           switch (instr.instruction.value.kind) {
             case "ComputedLoad":

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

@@ -7,6 +7,7 @@
 
 import { CompilerError } from "../CompilerError";
 import {
+  PrunedReactiveScopeBlock,
   ReactiveFunction,
   ReactiveScope,
   ReactiveScopeBlock,
@@ -83,6 +84,15 @@ export function writeReactiveBlock(
   writer.writeLine("}");
 }
 
+export function writePrunedScope(
+  writer: Writer,
+  block: PrunedReactiveScopeBlock
+): void {
+  writer.writeLine(`<pruned> ${printReactiveScopeSummary(block.scope)} {`);
+  writeReactiveInstructions(writer, block.instructions);
+  writer.writeLine("}");
+}
+
 export function printDependency(dependency: ReactiveScopeDependency): string {
   const identifier =
     printIdentifier(dependency.identifier) +
@@ -133,6 +143,10 @@ function writeReactiveInstruction(
       writeReactiveBlock(writer, instr);
       break;
     }
+    case "pruned-scope": {
+      writePrunedScope(writer, instr);
+      break;
+    }
     case "terminal": {
       if (instr.label !== null) {
         writer.write(`bb${instr.label.id}: `);

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

@@ -107,7 +107,14 @@ class Transform extends ReactiveFunctionTransform<boolean> {
             this.unmemoizedValues.add(identifier);
           }
         }
-        return { kind: "replace-many", value: scopeBlock.instructions };
+        return {
+          kind: "replace",
+          value: {
+            kind: "pruned-scope",
+            scope: scopeBlock.scope,
+            instructions: scopeBlock.instructions,
+          },
+        };
       }
     }
     return { kind: "keep" };

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

@@ -14,6 +14,7 @@ import {
   Place,
   ReactiveFunction,
   ReactiveInstruction,
+  ReactiveScope,
   ReactiveScopeBlock,
   ReactiveStatement,
   ReactiveTerminal,
@@ -122,18 +123,12 @@ export function pruneNonEscapingScopes(fn: ReactiveFunction): void {
   }
   visitReactiveFunction(fn, new CollectDependenciesVisitor(fn.env), state);
 
-  // log(() => prettyFormat(state));
-
   /*
    * Then walk outward from the returned values and find all captured operands.
    * This forms the set of identifiers which should be memoized.
    */
   const memoized = computeMemoizedIdentifiers(state);
 
-  // log(() => prettyFormat(memoized));
-
-  // log(() => printReactiveFunction(fn));
-
   // Prune scopes that do not declare/reassign any escaping values
   visitReactiveFunction(fn, new PruneScopesTransform(), memoized);
 }
@@ -200,7 +195,9 @@ type IdentifierNode = {
 
 // A scope node describing its dependencies
 type ScopeNode = {
-  dependencies: Array<IdentifierId>;
+  declarations: Array<IdentifierId>;
+  inputs: Set<IdentifierId>;
+  innerScopes: Set<ScopeId>;
   seen: boolean;
 };
 
@@ -216,6 +213,7 @@ class State {
   identifiers: Map<IdentifierId, IdentifierNode> = new Map();
   scopes: Map<ScopeId, ScopeNode> = new Map();
   escapingValues: Set<IdentifierId> = new Set();
+  currentScope: ReactiveScope | null = null;
 
   constructor(env: Environment) {
     this.env = env;
@@ -247,11 +245,18 @@ class State {
       let node = this.scopes.get(scope.id);
       if (node === undefined) {
         node = {
-          dependencies: [...scope.dependencies].map((dep) => dep.identifier.id),
+          declarations: [...scope.declarations]
+            .map(([id]) => id)
+            .concat(
+              [...scope.reassignments].map((identifier) => identifier.id)
+            ),
+          inputs: new Set(),
+          innerScopes: new Set(),
           seen: false,
         };
         this.scopes.set(scope.id, node);
       }
+      node.inputs.add(identifier);
       const identifierNode = this.identifiers.get(identifier);
       CompilerError.invariant(identifierNode !== undefined, {
         reason: "Expected identifier to be initialized",
@@ -262,6 +267,22 @@ class State {
       identifierNode.scopes.add(scope.id);
     }
   }
+
+  enterScope(scope: ReactiveScope, fn: () => void): void {
+    const parentScope = this.currentScope;
+    this.currentScope = scope;
+    fn();
+    this.currentScope = parentScope;
+
+    if (
+      parentScope !== null &&
+      scope.declarations.size === 0 &&
+      scope.reassignments.size === 0
+    ) {
+      const parentNode = this.scopes.get(parentScope.id)!;
+      parentNode.innerScopes.add(scope.id);
+    }
+  }
 }
 
 /*
@@ -273,7 +294,7 @@ function computeMemoizedIdentifiers(state: State): Set<IdentifierId> {
   const memoized = new Set<IdentifierId>();
 
   // Visit an identifier, optionally forcing it to be memoized
-  function visit(id: IdentifierId, forceMemoize: boolean = false): boolean {
+  function visit(id: IdentifierId): boolean {
     const node = state.identifiers.get(id);
     CompilerError.invariant(node !== undefined, {
       reason: `Expected a node for all identifiers, none found for \`${id}\``,
@@ -301,21 +322,19 @@ function computeMemoizedIdentifiers(state: State): Set<IdentifierId> {
 
     if (
       node.level === MemoizationLevel.Memoized ||
-      (node.level === MemoizationLevel.Conditional &&
-        (hasMemoizedDependency || forceMemoize)) ||
-      (node.level === MemoizationLevel.Unmemoized && forceMemoize)
+      (node.level === MemoizationLevel.Conditional && hasMemoizedDependency)
     ) {
       node.memoized = true;
       memoized.add(id);
       for (const scope of node.scopes) {
-        forceMemoizeScopeDependencies(scope);
+        memoizeScopeDependencies(scope);
       }
     }
     return node.memoized;
   }
 
   // Force all the scope's optionally-memoizeable dependencies (not "Never") to be memoized
-  function forceMemoizeScopeDependencies(id: ScopeId): void {
+  function memoizeScopeDependencies(id: ScopeId): void {
     const node = state.scopes.get(id);
     CompilerError.invariant(node !== undefined, {
       reason: "Expected a node for all scopes",
@@ -328,8 +347,14 @@ function computeMemoizedIdentifiers(state: State): Set<IdentifierId> {
     }
     node.seen = true;
 
-    for (const dep of node.dependencies) {
-      visit(dep, true);
+    for (const dep of node.declarations) {
+      visit(dep);
+    }
+    for (const dep of node.inputs) {
+      visit(dep);
+    }
+    for (const scope of node.innerScopes) {
+      memoizeScopeDependencies(scope);
     }
     return;
   }
@@ -814,6 +839,12 @@ class CollectDependenciesVisitor extends ReactiveFunctionVisitor<State> {
     };
   }
 
+  override visitScope(scopeBlock: ReactiveScopeBlock, state: State): void {
+    state.enterScope(scopeBlock.scope, () => {
+      this.traverseScope(scopeBlock, state);
+    });
+  }
+
   override visitInstruction(
     instruction: ReactiveInstruction,
     state: State
@@ -950,7 +981,14 @@ class PruneScopesTransform extends ReactiveFunctionTransform<
       return { kind: "keep" };
     } else {
       this.prunedScopes.add(scopeBlock.scope.id);
-      return { kind: "replace-many", value: scopeBlock.instructions };
+      return {
+        kind: "replace",
+        value: {
+          kind: "pruned-scope",
+          scope: scopeBlock.scope,
+          instructions: scopeBlock.instructions,
+        },
+      };
     }
   }
 

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

@@ -51,7 +51,14 @@ class Transform extends ReactiveFunctionTransform<State> {
          */
         !hasOwnDeclaration(scopeBlock))
     ) {
-      return { kind: "replace-many", value: scopeBlock.instructions };
+      return {
+        kind: "replace",
+        value: {
+          kind: "pruned-scope",
+          scope: scopeBlock.scope,
+          instructions: scopeBlock.instructions,
+        },
+      };
     } else {
       return { kind: "keep" };
     }

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

@@ -12,6 +12,7 @@ import {
   IdentifierName,
   InstructionId,
   Place,
+  PrunedReactiveScopeBlock,
   ReactiveBlock,
   ReactiveFunction,
   ReactiveScopeBlock,
@@ -84,6 +85,13 @@ class Visitor extends ReactiveFunctionVisitor<Scopes> {
     });
   }
 
+  override visitPrunedScope(
+    scopeBlock: PrunedReactiveScopeBlock,
+    state: Scopes
+  ): void {
+    this.traverseBlock(scopeBlock.instructions, state);
+  }
+
   override visitScope(scope: ReactiveScopeBlock, state: Scopes): void {
     for (const [_, declaration] of scope.scope.declarations) {
       state.visit(declaration.identifier);

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

@@ -9,6 +9,7 @@ import {
   HIRFunction,
   InstructionId,
   Place,
+  PrunedReactiveScopeBlock,
   ReactiveBlock,
   ReactiveFunction,
   ReactiveInstruction,
@@ -196,6 +197,16 @@ export class ReactiveFunctionVisitor<TState = void> {
     this.visitBlock(scope.instructions, state);
   }
 
+  visitPrunedScope(scopeBlock: PrunedReactiveScopeBlock, state: TState): void {
+    this.traversePrunedScope(scopeBlock, state);
+  }
+  traversePrunedScope(
+    scopeBlock: PrunedReactiveScopeBlock,
+    state: TState
+  ): void {
+    this.visitBlock(scopeBlock.instructions, state);
+  }
+
   visitBlock(block: ReactiveBlock, state: TState): void {
     this.traverseBlock(block, state);
   }
@@ -210,6 +221,10 @@ export class ReactiveFunctionVisitor<TState = void> {
           this.visitScope(instr, state);
           break;
         }
+        case "pruned-scope": {
+          this.visitPrunedScope(instr, state);
+          break;
+        }
         case "terminal": {
           this.visitTerminal(instr, state);
           break;
@@ -273,6 +288,10 @@ export class ReactiveFunctionTransform<
           transformed = this.transformScope(instr, state);
           break;
         }
+        case "pruned-scope": {
+          transformed = this.transformPrunedScope(instr, state);
+          break;
+        }
         case "terminal": {
           transformed = this.transformTerminal(instr, state);
           break;
@@ -339,6 +358,14 @@ export class ReactiveFunctionTransform<
     return { kind: "keep" };
   }
 
+  transformPrunedScope(
+    scope: PrunedReactiveScopeBlock,
+    state: TState
+  ): Transformed<ReactiveStatement> {
+    this.visitPrunedScope(scope, state);
+    return { kind: "keep" };
+  }
+
   transformValue(
     id: InstructionId,
     value: ReactiveValue,

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

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

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/method-call-computed.expect.md

@@ -23,7 +23,7 @@ function foo(a, b, c) {
 ```javascript
 import { c as _c } from "react/compiler-runtime";
 function foo(a, b, c) {
-  const $ = _c(8);
+  const $ = _c(6);
   let t0;
   if ($[0] !== a) {
     t0 = makeObject(a);
@@ -33,26 +33,18 @@ function foo(a, b, c) {
     t0 = $[1];
   }
   const x = t0;
+  const y = makeObject(a);
   let t1;
-  if ($[2] !== a) {
-    t1 = makeObject(a);
-    $[2] = a;
-    $[3] = t1;
+  if ($[2] !== x || $[3] !== y.method || $[4] !== b) {
+    t1 = x[y.method](b);
+    $[2] = x;
+    $[3] = y.method;
+    $[4] = b;
+    $[5] = t1;
   } else {
-    t1 = $[3];
+    t1 = $[5];
   }
-  const y = t1;
-  let t2;
-  if ($[4] !== x || $[5] !== y.method || $[6] !== b) {
-    t2 = x[y.method](b);
-    $[4] = x;
-    $[5] = y.method;
-    $[6] = b;
-    $[7] = t2;
-  } else {
-    t2 = $[7];
-  }
-  const z = t2;
+  const z = t1;
   return z;
 }
 

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/original-reactive-scopes-fork/mutate-outer-scope-within-value-block.expect.md

@@ -68,26 +68,20 @@ import { CONST_TRUE, identity, shallowCopy } from "shared-runtime";
  * should be merged.
  */
 function useFoo(t0) {
-  const $ = _c(3);
+  const $ = _c(2);
   const { input } = t0;
   const arr = shallowCopy(input);
+
+  const cond = identity(false);
   let t1;
-  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
-    t1 = identity(false);
-    $[0] = t1;
+  if ($[0] !== arr) {
+    t1 = cond ? { val: CONST_TRUE } : mutate(arr);
+    $[0] = arr;
+    $[1] = t1;
   } else {
-    t1 = $[0];
+    t1 = $[1];
   }
-  const cond = t1;
-  let t2;
-  if ($[1] !== arr) {
-    t2 = cond ? { val: CONST_TRUE } : mutate(arr);
-    $[1] = arr;
-    $[2] = t2;
-  } else {
-    t2 = $[2];
-  }
-  return t2;
+  return t1;
 }
 
 export const FIXTURE_ENTRYPOINT = {

compiler/packages/snap/src/runner-watch.ts

@@ -153,8 +153,8 @@ function subscribeFilterFile(
     } else if (
       events.findIndex((event) => event.path.includes(FILTER_FILENAME)) !== -1
     ) {
+      state.filter = await readTestFilter();
       if (state.mode.filter) {
-        state.filter = await readTestFilter();
         state.mode.action = RunnerAction.Test;
         onChange(state);
       }
@@ -218,7 +218,7 @@ export async function makeWatchRunner(
       action: RunnerAction.Test,
       filter: filterMode,
     },
-    filter: filterMode ? await readTestFilter() : null,
+    filter: await readTestFilter(),
   };
 
   subscribeTsc(state, onChange);