[compiler] Improve setState-in-effects rule to account for ref-gated conditionals

Conditionally calling setState in an effect is sometimes necessary, but should generally follow the pattern of using a "previous vaue" ref to manually compare and ensure that the setState is idempotent. See fixture for an example.

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

@@ -672,9 +672,14 @@ export const EnvironmentConfigSchema = z.object({
   validateNoDynamicallyCreatedComponentsOrHooks: z.boolean().default(false),
 
   /**
-   * When enabled, allows setState calls in effects when the value being set is
-   * derived from a ref. This is useful for patterns where initial layout measurements
-   * from refs need to be stored in state during mount.
+   * When enabled, allows setState calls in effects based on valid patterns involving refs:
+   * - Allow setState where the value being set is derived from a ref. This is useful where
+   *   state needs to take into account layer information, and a layout effect reads layout
+   *   data from a ref and sets state.
+   * - Allow conditionally calling setState after manually comparing previous/new values
+   *   for changes via a ref. Relying on effect deps is insufficient for non-primitive values,
+   *   so a ref is generally required to manually track previous values and compare prev/next
+   *   for meaningful changes before setting state.
    */
   enableAllowSetStateFromRefsInEffects: z.boolean().default(true),
 

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

@@ -11,6 +11,7 @@ import {
   BasicBlock,
   BlockId,
   Instruction,
+  InstructionKind,
   InstructionValue,
   makeInstructionId,
   Pattern,
@@ -32,6 +33,32 @@ export function* eachInstructionLValue(
   yield* eachInstructionValueLValue(instr.value);
 }
 
+export function* eachInstructionLValueWithKind(
+  instr: ReactiveInstruction,
+): Iterable<[Place, InstructionKind]> {
+  switch (instr.value.kind) {
+    case 'DeclareContext':
+    case 'StoreContext':
+    case 'DeclareLocal':
+    case 'StoreLocal': {
+      yield [instr.value.lvalue.place, instr.value.lvalue.kind];
+      break;
+    }
+    case 'Destructure': {
+      const kind = instr.value.lvalue.kind;
+      for (const place of eachPatternOperand(instr.value.lvalue.pattern)) {
+        yield [place, kind];
+      }
+      break;
+    }
+    case 'PostfixUpdate':
+    case 'PrefixUpdate': {
+      yield [instr.value.lvalue, InstructionKind.Reassign];
+      break;
+    }
+  }
+}
+
 export function* eachInstructionValueLValue(
   value: ReactiveValue,
 ): Iterable<Place> {

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

@@ -0,0 +1,114 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+import {BlockId, computePostDominatorTree, HIRFunction, Place} from '../HIR';
+import {PostDominator} from '../HIR/Dominator';
+
+export type ControlDominators = (id: BlockId) => boolean;
+
+/**
+ * Returns an object that lazily calculates whether particular blocks are controlled
+ * by values of interest. Which values matter are up to the caller.
+ */
+export function createControlDominators(
+  fn: HIRFunction,
+  isControlVariable: (place: Place) => boolean,
+): ControlDominators {
+  const postDominators = computePostDominatorTree(fn, {
+    includeThrowsAsExitNode: false,
+  });
+  const postDominatorFrontierCache = new Map<BlockId, Set<BlockId>>();
+
+  function isControlledBlock(id: BlockId): boolean {
+    let controlBlocks = postDominatorFrontierCache.get(id);
+    if (controlBlocks === undefined) {
+      controlBlocks = postDominatorFrontier(fn, postDominators, id);
+      postDominatorFrontierCache.set(id, controlBlocks);
+    }
+    for (const blockId of controlBlocks) {
+      const controlBlock = fn.body.blocks.get(blockId)!;
+      switch (controlBlock.terminal.kind) {
+        case 'if':
+        case 'branch': {
+          if (isControlVariable(controlBlock.terminal.test)) {
+            return true;
+          }
+          break;
+        }
+        case 'switch': {
+          if (isControlVariable(controlBlock.terminal.test)) {
+            return true;
+          }
+          for (const case_ of controlBlock.terminal.cases) {
+            if (case_.test !== null && isControlVariable(case_.test)) {
+              return true;
+            }
+          }
+          break;
+        }
+      }
+    }
+    return false;
+  }
+
+  return isControlledBlock;
+}
+
+/*
+ * Computes the post-dominator frontier of @param block. These are immediate successors of nodes that
+ * post-dominate @param targetId and from which execution may not reach @param block. Intuitively, these
+ * are the earliest blocks from which execution branches such that it may or may not reach the target block.
+ */
+function postDominatorFrontier(
+  fn: HIRFunction,
+  postDominators: PostDominator<BlockId>,
+  targetId: BlockId,
+): Set<BlockId> {
+  const visited = new Set<BlockId>();
+  const frontier = new Set<BlockId>();
+  const targetPostDominators = postDominatorsOf(fn, postDominators, targetId);
+  for (const blockId of [...targetPostDominators, targetId]) {
+    if (visited.has(blockId)) {
+      continue;
+    }
+    visited.add(blockId);
+    const block = fn.body.blocks.get(blockId)!;
+    for (const pred of block.preds) {
+      if (!targetPostDominators.has(pred)) {
+        // The predecessor does not always reach this block, we found an item on the frontier!
+        frontier.add(pred);
+      }
+    }
+  }
+  return frontier;
+}
+
+function postDominatorsOf(
+  fn: HIRFunction,
+  postDominators: PostDominator<BlockId>,
+  targetId: BlockId,
+): Set<BlockId> {
+  const result = new Set<BlockId>();
+  const visited = new Set<BlockId>();
+  const queue = [targetId];
+  while (queue.length) {
+    const currentId = queue.shift()!;
+    if (visited.has(currentId)) {
+      continue;
+    }
+    visited.add(currentId);
+    const current = fn.body.blocks.get(currentId)!;
+    for (const pred of current.preds) {
+      const predPostDominator = postDominators.get(pred) ?? pred;
+      if (predPostDominator === targetId || result.has(predPostDominator)) {
+        result.add(pred);
+      }
+      queue.push(pred);
+    }
+  }
+  return result;
+}

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

@@ -7,7 +7,6 @@
 
 import {CompilerError} from '..';
 import {
-  BlockId,
   Effect,
   Environment,
   HIRFunction,
@@ -15,14 +14,12 @@ import {
   IdentifierId,
   Instruction,
   Place,
-  computePostDominatorTree,
   evaluatesToStableTypeOrContainer,
   getHookKind,
   isStableType,
   isStableTypeContainer,
   isUseOperator,
 } from '../HIR';
-import {PostDominator} from '../HIR/Dominator';
 import {
   eachInstructionLValue,
   eachInstructionOperand,
@@ -35,6 +32,7 @@ import {
 } from '../ReactiveScopes/InferReactiveScopeVariables';
 import DisjointSet from '../Utils/DisjointSet';
 import {assertExhaustive} from '../Utils/utils';
+import {createControlDominators} from './ControlDominators';
 
 /**
  * Side map to track and propagate sources of stability (i.e. hook calls such as
@@ -212,45 +210,9 @@ export function inferReactivePlaces(fn: HIRFunction): void {
     reactiveIdentifiers.markReactive(place);
   }
 
-  const postDominators = computePostDominatorTree(fn, {
-    includeThrowsAsExitNode: false,
-  });
-  const postDominatorFrontierCache = new Map<BlockId, Set<BlockId>>();
-
-  function isReactiveControlledBlock(id: BlockId): boolean {
-    let controlBlocks = postDominatorFrontierCache.get(id);
-    if (controlBlocks === undefined) {
-      controlBlocks = postDominatorFrontier(fn, postDominators, id);
-      postDominatorFrontierCache.set(id, controlBlocks);
-    }
-    for (const blockId of controlBlocks) {
-      const controlBlock = fn.body.blocks.get(blockId)!;
-      switch (controlBlock.terminal.kind) {
-        case 'if':
-        case 'branch': {
-          if (reactiveIdentifiers.isReactive(controlBlock.terminal.test)) {
-            return true;
-          }
-          break;
-        }
-        case 'switch': {
-          if (reactiveIdentifiers.isReactive(controlBlock.terminal.test)) {
-            return true;
-          }
-          for (const case_ of controlBlock.terminal.cases) {
-            if (
-              case_.test !== null &&
-              reactiveIdentifiers.isReactive(case_.test)
-            ) {
-              return true;
-            }
-          }
-          break;
-        }
-      }
-    }
-    return false;
-  }
+  const isReactiveControlledBlock = createControlDominators(fn, place =>
+    reactiveIdentifiers.isReactive(place),
+  );
 
   do {
     for (const [, block] of fn.body.blocks) {
@@ -411,61 +373,6 @@ export function inferReactivePlaces(fn: HIRFunction): void {
   propagateReactivityToInnerFunctions(fn, true);
 }
 
-/*
- * Computes the post-dominator frontier of @param block. These are immediate successors of nodes that
- * post-dominate @param targetId and from which execution may not reach @param block. Intuitively, these
- * are the earliest blocks from which execution branches such that it may or may not reach the target block.
- */
-function postDominatorFrontier(
-  fn: HIRFunction,
-  postDominators: PostDominator<BlockId>,
-  targetId: BlockId,
-): Set<BlockId> {
-  const visited = new Set<BlockId>();
-  const frontier = new Set<BlockId>();
-  const targetPostDominators = postDominatorsOf(fn, postDominators, targetId);
-  for (const blockId of [...targetPostDominators, targetId]) {
-    if (visited.has(blockId)) {
-      continue;
-    }
-    visited.add(blockId);
-    const block = fn.body.blocks.get(blockId)!;
-    for (const pred of block.preds) {
-      if (!targetPostDominators.has(pred)) {
-        // The predecessor does not always reach this block, we found an item on the frontier!
-        frontier.add(pred);
-      }
-    }
-  }
-  return frontier;
-}
-
-function postDominatorsOf(
-  fn: HIRFunction,
-  postDominators: PostDominator<BlockId>,
-  targetId: BlockId,
-): Set<BlockId> {
-  const result = new Set<BlockId>();
-  const visited = new Set<BlockId>();
-  const queue = [targetId];
-  while (queue.length) {
-    const currentId = queue.shift()!;
-    if (visited.has(currentId)) {
-      continue;
-    }
-    visited.add(currentId);
-    const current = fn.body.blocks.get(currentId)!;
-    for (const pred of current.preds) {
-      const predPostDominator = postDominators.get(pred) ?? pred;
-      if (predPostDominator === targetId || result.has(predPostDominator)) {
-        result.add(pred);
-      }
-      queue.push(pred);
-    }
-  }
-  return result;
-}
-
 class ReactivityMap {
   hasChanges: boolean = false;
   reactive: Set<IdentifierId> = new Set();

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

@@ -1359,8 +1359,6 @@ function codegenInstructionNullable(
       value = null;
     } else {
       lvalue = instr.value.lvalue.pattern;
-      let hasReassign = false;
-      let hasDeclaration = false;
       for (const place of eachPatternOperand(lvalue)) {
         if (
           kind !== InstructionKind.Reassign &&
@@ -1368,26 +1366,6 @@ function codegenInstructionNullable(
         ) {
           cx.temp.set(place.identifier.declarationId, null);
         }
-        const isDeclared = cx.hasDeclared(place.identifier);
-        hasReassign ||= isDeclared;
-        hasDeclaration ||= !isDeclared;
-      }
-      if (hasReassign && hasDeclaration) {
-        CompilerError.invariant(false, {
-          reason:
-            'Encountered a destructuring operation where some identifiers are already declared (reassignments) but others are not (declarations)',
-          description: null,
-          details: [
-            {
-              kind: 'error',
-              loc: instr.loc,
-              message: null,
-            },
-          ],
-          suggestions: null,
-        });
-      } else if (hasReassign) {
-        kind = InstructionKind.Reassign;
       }
       value = codegenPlaceToExpression(cx, instr.value.value);
     }

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

@@ -19,7 +19,11 @@ import {
   promoteTemporary,
 } from '../HIR';
 import {clonePlaceToTemporary} from '../HIR/HIRBuilder';
-import {eachPatternOperand, mapPatternOperands} from '../HIR/visitors';
+import {
+  eachInstructionLValueWithKind,
+  eachPatternOperand,
+  mapPatternOperands,
+} from '../HIR/visitors';
 import {
   ReactiveFunctionTransform,
   Transformed,
@@ -113,6 +117,9 @@ class Visitor extends ReactiveFunctionTransform<State> {
   ): Transformed<ReactiveStatement> {
     this.visitInstruction(instruction, state);
 
+    let instructionsToProcess: Array<ReactiveInstruction> = [instruction];
+    let result: Transformed<ReactiveStatement> = {kind: 'keep'};
+
     if (instruction.value.kind === 'Destructure') {
       const transformed = transformDestructuring(
         state,
@@ -120,7 +127,8 @@ class Visitor extends ReactiveFunctionTransform<State> {
         instruction.value,
       );
       if (transformed) {
-        return {
+        instructionsToProcess = transformed;
+        result = {
           kind: 'replace-many',
           value: transformed.map(instruction => ({
             kind: 'instruction',
@@ -129,7 +137,17 @@ class Visitor extends ReactiveFunctionTransform<State> {
         };
       }
     }
-    return {kind: 'keep'};
+
+    // Update state.declared with declarations from the instruction(s)
+    for (const instr of instructionsToProcess) {
+      for (const [place, kind] of eachInstructionLValueWithKind(instr)) {
+        if (kind !== InstructionKind.Reassign) {
+          state.declared.add(place.identifier.declarationId);
+        }
+      }
+    }
+
+    return result;
   }
 }
 
@@ -144,10 +162,13 @@ function transformDestructuring(
     const isDeclared = state.declared.has(place.identifier.declarationId);
     if (isDeclared) {
       reassigned.add(place.identifier.id);
+    } else {
+      hasDeclaration = true;
     }
-    hasDeclaration ||= !isDeclared;
   }
-  if (reassigned.size === 0 || !hasDeclaration) {
+  if (!hasDeclaration) {
+    // all reassignments
+    destructure.lvalue.kind = InstructionKind.Reassign;
     return null;
   }
   /*

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

@@ -21,13 +21,17 @@ import {
   isUseRefType,
   isRefValueType,
   Place,
+  Effect,
+  BlockId,
 } from '../HIR';
 import {
   eachInstructionLValue,
   eachInstructionValueOperand,
 } from '../HIR/visitors';
+import {createControlDominators} from '../Inference/ControlDominators';
+import {isMutable} from '../ReactiveScopes/InferReactiveScopeVariables';
 import {Result} from '../Utils/Result';
-import {Iterable_some} from '../Utils/utils';
+import {assertExhaustive, Iterable_some} from '../Utils/utils';
 
 /**
  * Validates against calling setState in the body of an effect (useEffect and friends),
@@ -140,6 +144,8 @@ function getSetStateCall(
   setStateFunctions: Map<IdentifierId, Place>,
   env: Environment,
 ): Place | null {
+  const enableAllowSetStateFromRefsInEffects =
+    env.config.enableAllowSetStateFromRefsInEffects;
   const refDerivedValues: Set<IdentifierId> = new Set();
 
   const isDerivedFromRef = (place: Place): boolean => {
@@ -150,9 +156,38 @@ function getSetStateCall(
     );
   };
 
+  const isRefControlledBlock: (id: BlockId) => boolean =
+    enableAllowSetStateFromRefsInEffects
+      ? createControlDominators(fn, place => isDerivedFromRef(place))
+      : (): boolean => false;
+
   for (const [, block] of fn.body.blocks) {
+    if (enableAllowSetStateFromRefsInEffects) {
+      for (const phi of block.phis) {
+        if (isDerivedFromRef(phi.place)) {
+          continue;
+        }
+        let isPhiDerivedFromRef = false;
+        for (const [, operand] of phi.operands) {
+          if (isDerivedFromRef(operand)) {
+            isPhiDerivedFromRef = true;
+            break;
+          }
+        }
+        if (isPhiDerivedFromRef) {
+          refDerivedValues.add(phi.place.identifier.id);
+        } else {
+          for (const [pred] of phi.operands) {
+            if (isRefControlledBlock(pred)) {
+              refDerivedValues.add(phi.place.identifier.id);
+              break;
+            }
+          }
+        }
+      }
+    }
     for (const instr of block.instructions) {
-      if (env.config.enableAllowSetStateFromRefsInEffects) {
+      if (enableAllowSetStateFromRefsInEffects) {
         const hasRefOperand = Iterable_some(
           eachInstructionValueOperand(instr.value),
           isDerivedFromRef,
@@ -162,6 +197,46 @@ function getSetStateCall(
           for (const lvalue of eachInstructionLValue(instr)) {
             refDerivedValues.add(lvalue.identifier.id);
           }
+          // Ref-derived values can also propagate through mutation
+          for (const operand of eachInstructionValueOperand(instr.value)) {
+            switch (operand.effect) {
+              case Effect.Capture:
+              case Effect.Store:
+              case Effect.ConditionallyMutate:
+              case Effect.ConditionallyMutateIterator:
+              case Effect.Mutate: {
+                if (isMutable(instr, operand)) {
+                  refDerivedValues.add(operand.identifier.id);
+                }
+                break;
+              }
+              case Effect.Freeze:
+              case Effect.Read: {
+                // no-op
+                break;
+              }
+              case Effect.Unknown: {
+                CompilerError.invariant(false, {
+                  reason: 'Unexpected unknown effect',
+                  description: null,
+                  details: [
+                    {
+                      kind: 'error',
+                      loc: operand.loc,
+                      message: null,
+                    },
+                  ],
+                  suggestions: null,
+                });
+              }
+              default: {
+                assertExhaustive(
+                  operand.effect,
+                  `Unexpected effect kind \`${operand.effect}\``,
+                );
+              }
+            }
+          }
         }
 
         if (
@@ -203,7 +278,7 @@ function getSetStateCall(
             isSetStateType(callee.identifier) ||
             setStateFunctions.has(callee.identifier.id)
           ) {
-            if (env.config.enableAllowSetStateFromRefsInEffects) {
+            if (enableAllowSetStateFromRefsInEffects) {
               const arg = instr.value.args.at(0);
               if (
                 arg !== undefined &&
@@ -216,6 +291,8 @@ function getSetStateCall(
                  * be needed when initial layout measurements from refs need to be stored in state.
                  */
                 return null;
+              } else if (isRefControlledBlock(block.id)) {
+                continue;
               }
             }
             /*

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-invalid-destructuring-reassignment-undefined-variable.expect.md

@@ -0,0 +1,162 @@
+
+## Input
+
+```javascript
+// @flow @compilationMode:"infer"
+'use strict';
+
+function getWeekendDays(user) {
+  return [0, 6];
+}
+
+function getConfig(weekendDays) {
+  return [1, 5];
+}
+
+component Calendar(user, defaultFirstDay, currentDate, view) {
+  const weekendDays = getWeekendDays(user);
+  let firstDay = defaultFirstDay;
+  let daysToDisplay = 7;
+  if (view === 'week') {
+    let lastDay;
+    // this assignment produces invalid code
+    [firstDay, lastDay] = getConfig(weekendDays);
+    daysToDisplay = ((7 + lastDay - firstDay) % 7) + 1;
+  } else if (view === 'day') {
+    firstDay = currentDate.getDayOfWeek();
+    daysToDisplay = 1;
+  }
+
+  return [currentDate, firstDay, daysToDisplay];
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Calendar,
+  params: [
+    {
+      user: {},
+      defaultFirstDay: 1,
+      currentDate: {getDayOfWeek: () => 3},
+      view: 'week',
+    },
+  ],
+  sequentialRenders: [
+    {
+      user: {},
+      defaultFirstDay: 1,
+      currentDate: {getDayOfWeek: () => 3},
+      view: 'week',
+    },
+    {
+      user: {},
+      defaultFirstDay: 1,
+      currentDate: {getDayOfWeek: () => 3},
+      view: 'day',
+    },
+  ],
+};
+
+```
+
+## Code
+
+```javascript
+"use strict";
+import { c as _c } from "react/compiler-runtime";
+
+function getWeekendDays(user) {
+  return [0, 6];
+}
+
+function getConfig(weekendDays) {
+  return [1, 5];
+}
+
+function Calendar(t0) {
+  const $ = _c(12);
+  const { user, defaultFirstDay, currentDate, view } = t0;
+  let daysToDisplay;
+  let firstDay;
+  if (
+    $[0] !== currentDate ||
+    $[1] !== defaultFirstDay ||
+    $[2] !== user ||
+    $[3] !== view
+  ) {
+    const weekendDays = getWeekendDays(user);
+    firstDay = defaultFirstDay;
+    daysToDisplay = 7;
+    if (view === "week") {
+      let lastDay;
+
+      [firstDay, lastDay] = getConfig(weekendDays);
+      daysToDisplay = ((7 + lastDay - firstDay) % 7) + 1;
+    } else {
+      if (view === "day") {
+        let t1;
+        if ($[6] !== currentDate) {
+          t1 = currentDate.getDayOfWeek();
+          $[6] = currentDate;
+          $[7] = t1;
+        } else {
+          t1 = $[7];
+        }
+        firstDay = t1;
+        daysToDisplay = 1;
+      }
+    }
+    $[0] = currentDate;
+    $[1] = defaultFirstDay;
+    $[2] = user;
+    $[3] = view;
+    $[4] = daysToDisplay;
+    $[5] = firstDay;
+  } else {
+    daysToDisplay = $[4];
+    firstDay = $[5];
+  }
+  let t1;
+  if ($[8] !== currentDate || $[9] !== daysToDisplay || $[10] !== firstDay) {
+    t1 = [currentDate, firstDay, daysToDisplay];
+    $[8] = currentDate;
+    $[9] = daysToDisplay;
+    $[10] = firstDay;
+    $[11] = t1;
+  } else {
+    t1 = $[11];
+  }
+  return t1;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Calendar,
+  params: [
+    {
+      user: {},
+      defaultFirstDay: 1,
+      currentDate: { getDayOfWeek: () => 3 },
+      view: "week",
+    },
+  ],
+
+  sequentialRenders: [
+    {
+      user: {},
+      defaultFirstDay: 1,
+      currentDate: { getDayOfWeek: () => 3 },
+      view: "week",
+    },
+    {
+      user: {},
+      defaultFirstDay: 1,
+      currentDate: { getDayOfWeek: () => 3 },
+      view: "day",
+    },
+  ],
+};
+
+```
+      
+### Eval output
+(kind: ok) [{"getDayOfWeek":"[[ function params=0 ]]"},1,5]
+[{"getDayOfWeek":"[[ function params=0 ]]"},3,1]

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-invalid-destructuring-reassignment-undefined-variable.js

@@ -0,0 +1,53 @@
+// @flow @compilationMode:"infer"
+'use strict';
+
+function getWeekendDays(user) {
+  return [0, 6];
+}
+
+function getConfig(weekendDays) {
+  return [1, 5];
+}
+
+component Calendar(user, defaultFirstDay, currentDate, view) {
+  const weekendDays = getWeekendDays(user);
+  let firstDay = defaultFirstDay;
+  let daysToDisplay = 7;
+  if (view === 'week') {
+    let lastDay;
+    // this assignment produces invalid code
+    [firstDay, lastDay] = getConfig(weekendDays);
+    daysToDisplay = ((7 + lastDay - firstDay) % 7) + 1;
+  } else if (view === 'day') {
+    firstDay = currentDate.getDayOfWeek();
+    daysToDisplay = 1;
+  }
+
+  return [currentDate, firstDay, daysToDisplay];
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Calendar,
+  params: [
+    {
+      user: {},
+      defaultFirstDay: 1,
+      currentDate: {getDayOfWeek: () => 3},
+      view: 'week',
+    },
+  ],
+  sequentialRenders: [
+    {
+      user: {},
+      defaultFirstDay: 1,
+      currentDate: {getDayOfWeek: () => 3},
+      view: 'week',
+    },
+    {
+      user: {},
+      defaultFirstDay: 1,
+      currentDate: {getDayOfWeek: () => 3},
+      view: 'day',
+    },
+  ],
+};

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/valid-setState-in-useEffect-controlled-by-ref-value.expect.md

@@ -0,0 +1,117 @@
+
+## Input
+
+```javascript
+// @validateNoSetStateInEffects @enableAllowSetStateFromRefsInEffects @loggerTestOnly @compilationMode:"infer"
+import {useState, useRef, useEffect} from 'react';
+
+function Component({x, y}) {
+  const previousXRef = useRef(null);
+  const previousYRef = useRef(null);
+
+  const [data, setData] = useState(null);
+
+  useEffect(() => {
+    const previousX = previousXRef.current;
+    previousXRef.current = x;
+    const previousY = previousYRef.current;
+    previousYRef.current = y;
+    if (!areEqual(x, previousX) || !areEqual(y, previousY)) {
+      const data = load({x, y});
+      setData(data);
+    }
+  }, [x, y]);
+
+  return data;
+}
+
+function areEqual(a, b) {
+  return a === b;
+}
+
+function load({x, y}) {
+  return x * y;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{x: 0, y: 0}],
+  sequentialRenders: [
+    {x: 0, y: 0},
+    {x: 1, y: 0},
+    {x: 1, y: 1},
+  ],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime"; // @validateNoSetStateInEffects @enableAllowSetStateFromRefsInEffects @loggerTestOnly @compilationMode:"infer"
+import { useState, useRef, useEffect } from "react";
+
+function Component(t0) {
+  const $ = _c(4);
+  const { x, y } = t0;
+  const previousXRef = useRef(null);
+  const previousYRef = useRef(null);
+
+  const [data, setData] = useState(null);
+  let t1;
+  let t2;
+  if ($[0] !== x || $[1] !== y) {
+    t1 = () => {
+      const previousX = previousXRef.current;
+      previousXRef.current = x;
+      const previousY = previousYRef.current;
+      previousYRef.current = y;
+      if (!areEqual(x, previousX) || !areEqual(y, previousY)) {
+        const data_0 = load({ x, y });
+        setData(data_0);
+      }
+    };
+
+    t2 = [x, y];
+    $[0] = x;
+    $[1] = y;
+    $[2] = t1;
+    $[3] = t2;
+  } else {
+    t1 = $[2];
+    t2 = $[3];
+  }
+  useEffect(t1, t2);
+  return data;
+}
+
+function areEqual(a, b) {
+  return a === b;
+}
+
+function load({ x, y }) {
+  return x * y;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ x: 0, y: 0 }],
+  sequentialRenders: [
+    { x: 0, y: 0 },
+    { x: 1, y: 0 },
+    { x: 1, y: 1 },
+  ],
+};
+
+```
+
+## Logs
+
+```
+{"kind":"CompileSuccess","fnLoc":{"start":{"line":4,"column":0,"index":163},"end":{"line":22,"column":1,"index":631},"filename":"valid-setState-in-useEffect-controlled-by-ref-value.ts"},"fnName":"Component","memoSlots":4,"memoBlocks":1,"memoValues":2,"prunedMemoBlocks":0,"prunedMemoValues":0}
+```
+      
+### Eval output
+(kind: ok) 0
+0
+1

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/valid-setState-in-useEffect-controlled-by-ref-value.js

@@ -0,0 +1,40 @@
+// @validateNoSetStateInEffects @enableAllowSetStateFromRefsInEffects @loggerTestOnly @compilationMode:"infer"
+import {useState, useRef, useEffect} from 'react';
+
+function Component({x, y}) {
+  const previousXRef = useRef(null);
+  const previousYRef = useRef(null);
+
+  const [data, setData] = useState(null);
+
+  useEffect(() => {
+    const previousX = previousXRef.current;
+    previousXRef.current = x;
+    const previousY = previousYRef.current;
+    previousYRef.current = y;
+    if (!areEqual(x, previousX) || !areEqual(y, previousY)) {
+      const data = load({x, y});
+      setData(data);
+    }
+  }, [x, y]);
+
+  return data;
+}
+
+function areEqual(a, b) {
+  return a === b;
+}
+
+function load({x, y}) {
+  return x * y;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{x: 0, y: 0}],
+  sequentialRenders: [
+    {x: 0, y: 0},
+    {x: 1, y: 0},
+    {x: 1, y: 1},
+  ],
+};

packages/react-client/src/ReactFlightClient.js

@@ -813,6 +813,12 @@ function createInitializedStreamChunk<
   value: T,
   controller: FlightStreamController,
 ): InitializedChunk<T> {
+  if (__DEV__) {
+    // Retain a strong reference to the Response while we wait for chunks.
+    if (response._pendingChunks++ === 0) {
+      response._weakResponse.response = response;
+    }
+  }
   // We use the reason field to stash the controller since we already have that
   // field. It's a bit of a hack but efficient.
   // $FlowFixMe[invalid-constructor] Flow doesn't support functions as constructors
@@ -3075,7 +3081,6 @@ function resolveStream<T: ReadableStream | $AsyncIterable<any, any, void>>(
     // We already resolved. We didn't expect to see this.
     return;
   }
-  releasePendingChunk(response, chunk);
 
   const resolveListeners = chunk.value;
 
@@ -3375,6 +3380,14 @@ function stopStream(
     // We didn't expect not to have an existing stream;
     return;
   }
+  if (__DEV__) {
+    if (--response._pendingChunks === 0) {
+      // We're no longer waiting for any more chunks. We can release the strong
+      // reference to the response. We'll regain it if we ask for any more data
+      // later on.
+      response._weakResponse.response = null;
+    }
+  }
   const streamChunk: InitializedStreamChunk<any> = (chunk: any);
   const controller = streamChunk.reason;
   controller.close(row === '' ? '"$undefined"' : row);

packages/react-server/src/ReactFlightServer.js

@@ -339,8 +339,10 @@ function isPromiseAwaitInternal(url: string, functionName: string): boolean {
     case 'Function.resolve':
     case 'Function.all':
     case 'Function.allSettled':
+    case 'Function.any':
     case 'Function.race':
     case 'Function.try':
+    case 'Function.withResolvers':
       return true;
     default:
       return false;
@@ -2347,7 +2349,8 @@ function visitAsyncNodeImpl(
     // The technique for debugging the effects of uncached data on the render is to simply uncache it.
     return null;
   }
-  let previousIONode = null;
+
+  let previousIONode: void | null | PromiseNode | IONode = null;
   // First visit anything that blocked this sequence to start in the first place.
   if (node.previous !== null) {
     previousIONode = visitAsyncNode(
@@ -2363,12 +2366,20 @@ function visitAsyncNodeImpl(
       return undefined;
     }
   }
+
+  // `found` represents the return value of the following switch statement.
+  // We can't use multiple `return` statements in the switch statement
+  // since that prevents Closure compiler from inlining `visitAsyncImpl`
+  // thus doubling the call stack size.
+  let found: void | null | PromiseNode | IONode;
   switch (node.tag) {
     case IO_NODE: {
-      return node;
+      found = node;
+      break;
     }
     case UNRESOLVED_PROMISE_NODE: {
-      return previousIONode;
+      found = previousIONode;
+      break;
     }
     case PROMISE_NODE: {
       const awaited = node.awaited;
@@ -2379,7 +2390,8 @@ function visitAsyncNodeImpl(
         if (ioNode === undefined) {
           // Undefined is used as a signal that we found a suitable aborted node and we don't have to find
           // further aborted nodes.
-          return undefined;
+          found = undefined;
+          break;
         } else if (ioNode !== null) {
           // This Promise was blocked on I/O. That's a signal that this Promise is interesting to log.
           // We don't log it yet though. We return it to be logged by the point where it's awaited.
@@ -2436,10 +2448,12 @@ function visitAsyncNodeImpl(
           forwardDebugInfo(request, task, debugInfo);
         }
       }
-      return match;
+      found = match;
+      break;
     }
     case UNRESOLVED_AWAIT_NODE: {
-      return previousIONode;
+      found = previousIONode;
+      break;
     }
     case AWAIT_NODE: {
       const awaited = node.awaited;
@@ -2449,7 +2463,8 @@ function visitAsyncNodeImpl(
         if (ioNode === undefined) {
           // Undefined is used as a signal that we found a suitable aborted node and we don't have to find
           // further aborted nodes.
-          return undefined;
+          found = undefined;
+          break;
         } else if (ioNode !== null) {
           const startTime: number = node.start;
           const endTime: number = node.end;
@@ -2545,13 +2560,15 @@ function visitAsyncNodeImpl(
           forwardDebugInfo(request, task, debugInfo);
         }
       }
-      return match;
+      found = match;
+      break;
     }
     default: {
       // eslint-disable-next-line react-internal/prod-error-codes
       throw new Error('Unknown AsyncSequence tag. This is a bug in React.');
     }
   }
+  return found;
 }
 
 function emitAsyncSequence(