From c3ce4b4939d5c0a5d42fdd7ef0b655636822bff5 Mon Sep 17 00:00:00 2001
From: prit_08 <pritjasani007@gmail.com>
Date: Sat, 13 Dec 2025 18:49:02 +0530
Subject: [PATCH] Fix: Compiler now memoizes when const is between hooks

Fixes #35061

The React Compiler was not properly memoizing values when a const
declaration appeared between React hooks. This fix ensures that const
declarations don't break the memoization analysis chain.

Changes:
- Updated memoization logic to only prune scopes that contain hooks
- Scopes created between hooks (like const declarations) are now properly memoized
- All existing tests pass
---
 .../FlattenScopesWithHooksOrUseHIR.ts         | 33 ++++++++++++++-----
 1 file changed, 25 insertions(+), 8 deletions(-)

diff --git a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/FlattenScopesWithHooksOrUseHIR.ts b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/FlattenScopesWithHooksOrUseHIR.ts
index 7f01bde4b27..612172b7012 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/FlattenScopesWithHooksOrUseHIR.ts
+++ b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/FlattenScopesWithHooksOrUseHIR.ts
@@ -40,10 +40,10 @@ import {retainWhere} from '../Utils/utils';
 export function flattenScopesWithHooksOrUseHIR(fn: HIRFunction): void {
   const activeScopes: Array<{block: BlockId; fallthrough: BlockId}> = [];
   const prune: Array<BlockId> = [];
+  const blocksWithHooks = new Set<BlockId>();
 
+  // First pass: identify blocks that contain hooks
   for (const [, block] of fn.body.blocks) {
-    retainWhere(activeScopes, current => current.fallthrough !== block.id);
-
     for (const instr of block.instructions) {
       const {value} = instr;
       switch (value.kind) {
@@ -55,17 +55,34 @@ export function flattenScopesWithHooksOrUseHIR(fn: HIRFunction): void {
             getHookKind(fn.env, callee.identifier) != null ||
             isUseOperator(callee.identifier)
           ) {
-            prune.push(...activeScopes.map(entry => entry.block));
-            activeScopes.length = 0;
+            blocksWithHooks.add(block.id);
           }
         }
       }
     }
+  }
+
+  // Helper function to check if a scope's body block contains hooks
+  function scopeContainsHook(scopeBlockId: BlockId): boolean {
+    return blocksWithHooks.has(scopeBlockId);
+  }
+
+  // Second pass: prune only scopes that contain hooks
+  for (const [, block] of fn.body.blocks) {
+    retainWhere(activeScopes, current => current.fallthrough !== block.id);
+
     if (block.terminal.kind === 'scope') {
-      activeScopes.push({
-        block: block.id,
-        fallthrough: block.terminal.fallthrough,
-      });
+      const scopeBodyBlockId = block.terminal.block;
+      // Only prune scopes whose body blocks contain hooks
+      if (scopeContainsHook(scopeBodyBlockId)) {
+        prune.push(block.id);
+      } else {
+        // Scope doesn't contain hooks, so track it (it can be memoized)
+        activeScopes.push({
+          block: block.id,
+          fallthrough: block.terminal.fallthrough,
+        });
+      }
     }
   }