compiler: Promote pruned scope declarations to temporaries if used in…

… a later scope

There's a category of bug currently where pruned reactive scopes whose outputs are non-reactive can have their code end up inlining into another scope, moving the location of the instruction. Any value that had a scope assigned has to have its order of evaluation preserved, despite the fact that it got pruned, so naively we could just force every pruned scope to have its declarations promoted to named variables.

However, that ends up assigning names to _tons_ of scope declarations that don't really need to be promoted. For example, a scope with just a hook call ends up with:

```
const x = useFoo();

=>

scope {
 $t0 = Call read useFoo$ (...);
}
$t1 = StoreLocal 'x' = read $t0;
```

Where t0 doesn't need to be promoted since it's used immediately to assign to another value which is a non-temporary.

So the idea of this PR is that we can track outputs of pruned scopes which are directly referenced from inside a later scope. This fixes one of the two cases of the above pattern. We'll also likely have to consider values from pruned scopes as always reactive, i'll do that in the next PR.

ghstack-source-id: b37fb9a7cb1430b7c35ec5946269ce5a886a486a
Pull Request resolved: #29789

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

@@ -36,7 +36,7 @@ import {
   getHookKind,
   makeIdentifierName,
 } from "../HIR/HIR";
-import { printPlace } from "../HIR/PrintHIR";
+import { printIdentifier, printPlace } from "../HIR/PrintHIR";
 import { eachPatternOperand } from "../HIR/visitors";
 import { Err, Ok, Result } from "../Utils/Result";
 import { GuardKind } from "../Utils/RuntimeDiagnosticConstants";
@@ -524,8 +524,8 @@ function codegenReactiveScope(
     }
 
     CompilerError.invariant(identifier.name != null, {
-      reason: `Expected identifier '@${identifier.id}' to be named`,
-      description: null,
+      reason: `Expected scope declaration identifier to be named`,
+      description: `Declaration \`${printIdentifier(identifier)}\` is unnamed in scope @${scope.id}`,
       loc: null,
       suggestions: null,
     });

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

@@ -12,21 +12,20 @@ import {
   IdentifierId,
   InstructionId,
   Place,
+  PrunedReactiveScopeBlock,
   ReactiveFunction,
   ReactiveScopeBlock,
   ReactiveValue,
+  ScopeId,
   promoteTemporary,
   promoteTemporaryJsxTag,
 } from "../HIR/HIR";
 import { ReactiveFunctionVisitor, visitReactiveFunction } from "./visitors";
 
-type VisitorState = {
-  tags: JsxExpressionTags;
-};
-class Visitor extends ReactiveFunctionVisitor<VisitorState> {
-  override visitScope(block: ReactiveScopeBlock, state: VisitorState): void {
-    this.traverseScope(block, state);
-    for (const dep of block.scope.dependencies) {
+class Visitor extends ReactiveFunctionVisitor<State> {
+  override visitScope(scopeBlock: ReactiveScopeBlock, state: State): void {
+    this.traverseScope(scopeBlock, state);
+    for (const dep of scopeBlock.scope.dependencies) {
       const { identifier } = dep;
       if (identifier.name == null) {
         promoteIdentifier(identifier, state);
@@ -39,14 +38,29 @@ class Visitor extends ReactiveFunctionVisitor<VisitorState> {
      * Many of our current test fixtures do not return a value, so
      * it is better for now to promote (and memoize) every output.
      */
-    for (const [, declaration] of block.scope.declarations) {
+    for (const [, declaration] of scopeBlock.scope.declarations) {
       if (declaration.identifier.name == null) {
         promoteIdentifier(declaration.identifier, state);
       }
     }
   }
 
-  override visitParam(place: Place, state: VisitorState): void {
+  override visitPrunedScope(
+    scopeBlock: PrunedReactiveScopeBlock,
+    state: State
+  ): void {
+    this.traversePrunedScope(scopeBlock, state);
+    for (const [, declaration] of scopeBlock.scope.declarations) {
+      if (
+        declaration.identifier.name == null &&
+        state.pruned.get(declaration.identifier.id)?.usedOutsideScope === true
+      ) {
+        promoteIdentifier(declaration.identifier, state);
+      }
+    }
+  }
+
+  override visitParam(place: Place, state: State): void {
     if (place.identifier.name === null) {
       promoteIdentifier(place.identifier, state);
     }
@@ -55,7 +69,7 @@ class Visitor extends ReactiveFunctionVisitor<VisitorState> {
   override visitValue(
     id: InstructionId,
     value: ReactiveValue,
-    state: VisitorState
+    state: State
   ): void {
     this.traverseValue(id, value, state);
     if (value.kind === "FunctionExpression" || value.kind === "ObjectMethod") {
@@ -67,7 +81,7 @@ class Visitor extends ReactiveFunctionVisitor<VisitorState> {
     _id: InstructionId,
     _dependencies: Array<Place>,
     fn: ReactiveFunction,
-    state: VisitorState
+    state: State
   ): void {
     for (const operand of fn.params) {
       const place = operand.kind === "Identifier" ? operand : operand.place;
@@ -80,25 +94,65 @@ class Visitor extends ReactiveFunctionVisitor<VisitorState> {
 }
 
 type JsxExpressionTags = Set<IdentifierId>;
-class CollectJsxTagsVisitor extends ReactiveFunctionVisitor<JsxExpressionTags> {
+type State = {
+  tags: JsxExpressionTags;
+  pruned: Map<
+    IdentifierId,
+    { activeScopes: Array<ScopeId>; usedOutsideScope: boolean }
+  >; // true if referenced within another scope, false if only accessed outside of scopes
+};
+
+class CollectPromotableTemporaries extends ReactiveFunctionVisitor<State> {
+  activeScopes: Array<ScopeId> = [];
+
+  override visitPlace(_id: InstructionId, place: Place, state: State): void {
+    if (
+      this.activeScopes.length !== 0 &&
+      state.pruned.has(place.identifier.id)
+    ) {
+      const prunedPlace = state.pruned.get(place.identifier.id)!;
+      if (prunedPlace.activeScopes.indexOf(this.activeScopes.at(-1)!) === -1) {
+        prunedPlace.usedOutsideScope = true;
+      }
+    }
+  }
+
   override visitValue(
     id: InstructionId,
     value: ReactiveValue,
-    state: JsxExpressionTags
+    state: State
   ): void {
     this.traverseValue(id, value, state);
     if (value.kind === "JsxExpression" && value.tag.kind === "Identifier") {
-      state.add(value.tag.identifier.id);
+      state.tags.add(value.tag.identifier.id);
     }
   }
+
+  override visitPrunedScope(
+    scopeBlock: PrunedReactiveScopeBlock,
+    state: State
+  ): void {
+    for (const [id] of scopeBlock.scope.declarations) {
+      state.pruned.set(id, {
+        activeScopes: [...this.activeScopes],
+        usedOutsideScope: false,
+      });
+    }
+  }
+
+  override visitScope(scopeBlock: ReactiveScopeBlock, state: State): void {
+    this.activeScopes.push(scopeBlock.scope.id);
+    this.traverseScope(scopeBlock, state);
+    this.activeScopes.pop();
+  }
 }
 
 export function promoteUsedTemporaries(fn: ReactiveFunction): void {
-  const tags: JsxExpressionTags = new Set();
-  visitReactiveFunction(fn, new CollectJsxTagsVisitor(), tags);
-  const state: VisitorState = {
-    tags,
+  const state: State = {
+    tags: new Set(),
+    pruned: new Map(),
   };
+  visitReactiveFunction(fn, new CollectPromotableTemporaries(), state);
   for (const operand of fn.params) {
     const place = operand.kind === "Identifier" ? operand : operand.place;
     if (place.identifier.name === null) {
@@ -108,7 +162,7 @@ export function promoteUsedTemporaries(fn: ReactiveFunction): void {
   visitReactiveFunction(fn, new Visitor(), state);
 }
 
-function promoteIdentifier(identifier: Identifier, state: VisitorState): void {
+function promoteIdentifier(identifier: Identifier, state: State): void {
   CompilerError.invariant(identifier.name === null, {
     reason:
       "promoteTemporary: Expected to be called only for temporary variables",

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

@@ -18,6 +18,7 @@ import {
   isUseRefType,
   makeInstructionId,
   Place,
+  PrunedReactiveScopeBlock,
   ReactiveFunction,
   ReactiveInstruction,
   ReactiveScope,
@@ -687,6 +688,19 @@ class PropagationVisitor extends ReactiveFunctionVisitor<Context> {
     scope.scope.dependencies = scopeDependencies;
   }
 
+  override visitPrunedScope(
+    scopeBlock: PrunedReactiveScopeBlock,
+    context: Context
+  ): void {
+    /*
+     * NOTE: we explicitly throw away the deps, we only enter() the scope to record its
+     * declarations
+     */
+    const _scopeDepdencies = context.enter(scopeBlock.scope, () => {
+      this.visitBlock(scopeBlock.instructions, context);
+    });
+  }
+
   override visitInstruction(
     instruction: ReactiveInstruction,
     context: Context

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/bug-invalid-reactivity-value-block.expect.md

@@ -71,14 +71,15 @@ function Foo() {
   useNoAlias();
 
   const shouldCaptureObj = obj != null && CONST_TRUE;
-  let t0;
+  const t0 = shouldCaptureObj ? identity(obj) : null;
+  let t1;
   if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
-    t0 = [shouldCaptureObj ? identity(obj) : null, obj];
-    $[0] = t0;
+    t1 = [t0, obj];
+    $[0] = t1;
   } else {
-    t0 = $[0];
+    t1 = $[0];
   }
-  const result = t0;
+  const result = t1;
 
   useNoAlias(result, obj);
   if (shouldCaptureObj && result[0] !== obj) {