[fresh] Reset useMemoCache in response to a FastRefresh run

packages/react-reconciler/src/ReactFiber.js

@@ -210,6 +210,7 @@ function FiberNode(
       this._debugTask = null;
     }
     this._debugNeedsRemount = false;
+    this._debugNeedsMemoCacheReset = false;
     this._debugHookTypes = null;
     if (!hasBadMapPolyfill && typeof Object.preventExtensions === 'function') {
       Object.preventExtensions(this);
@@ -301,6 +302,7 @@ function createFiberImplObject(
       fiber._debugTask = null;
     }
     fiber._debugNeedsRemount = false;
+    fiber._debugNeedsMemoCacheReset = false;
     fiber._debugHookTypes = null;
     if (!hasBadMapPolyfill && typeof Object.preventExtensions === 'function') {
       Object.preventExtensions(fiber);
@@ -423,6 +425,8 @@ export function createWorkInProgress(current: Fiber, pendingProps: any): Fiber {
   if (__DEV__) {
     workInProgress._debugInfo = current._debugInfo;
     workInProgress._debugNeedsRemount = current._debugNeedsRemount;
+    workInProgress._debugNeedsMemoCacheReset =
+      current._debugNeedsMemoCacheReset;
     switch (workInProgress.tag) {
       case FunctionComponent:
       case SimpleMemoComponent:

packages/react-reconciler/src/ReactFiberHooks.js

@@ -1232,7 +1232,15 @@ function useMemoCache(size: number): Array<any> {
   // values from being reused. In prod environments this is never expected to happen. However, in
   // the unlikely case that it does vary between renders, we reset the cache anyway so behavior is
   // consistent in both environments.
-  if (data === undefined || data.length !== size) {
+  //
+  // The cache is also reset if the fiber is being rerendered as a result of Fast Refresh run, even
+  // if the cache size incidentally happens to be the same. This is to ensure that we don't  see
+  // incorrect values after a refresh.
+  if (
+    data === undefined ||
+    data.length !== size ||
+    currentlyRenderingFiber._debugNeedsMemoCacheReset === true
+  ) {
     data = memoCache.data[memoCache.index] = new Array(size);
     for (let i = 0; i < size; i++) {
       data[i] = REACT_MEMO_CACHE_SENTINEL;

packages/react-reconciler/src/ReactFiberHotReloading.js

@@ -307,6 +307,10 @@ function scheduleFibersWithFamiliesRecursively(
       }
     }
 
+    if (fiber.updateQueue != null && fiber.updateQueue.memoCache != null) {
+      fiber._debugNeedsMemoCacheReset = true;
+    }
+
     if (needsRemount) {
       fiber._debugNeedsRemount = true;
     }

packages/react-reconciler/src/ReactInternalTypes.js

@@ -210,6 +210,7 @@ export type Fiber = {
   _debugTask?: ConsoleTask | null,
   _debugIsCurrentlyTiming?: boolean,
   _debugNeedsRemount?: boolean,
+  _debugNeedsMemoCacheReset?: boolean,
 
   // Used to verify that the order of hooks does not change between renders.
   _debugHookTypes?: Array<HookType> | null,

packages/react-refresh/src/__tests__/ReactFreshIntegration-test.js

@@ -1637,7 +1637,69 @@ describe('ReactFreshIntegration', () => {
       }
     });
 
-    it('resets useMemoCache cache slots', async () => {
+    it('resets useMemoCache cache slots in a refresh', async () => {
+      if (__DEV__) {
+        await render(`
+          const useMemoCache = require('react/compiler-runtime').c;
+          let cacheMisses = 0;
+          const cacheMiss = (id) => {
+            cacheMisses++;
+            return id;
+          };
+          export default function App(t0) {
+            const $ = useMemoCache(2);
+            const {reset1, reset2} = t0;
+            let t1;
+            if ($[0] !== reset1) {
+              $[0] = t1 = cacheMiss({reset1});
+            } else {
+              t1 = $[1];
+            }
+            let t2;
+            if ($[1] !== reset2) {
+              $[1] = t2 = cacheMiss({reset2});
+            } else {
+              t2 = $[1];
+            }
+            return <h1>{cacheMisses}</h1>;
+          }
+        `);
+        const el = container.firstChild;
+        expect(el.textContent).toBe('2');
+        await patch(`
+          const useMemoCache = require('react/compiler-runtime').c;
+          let cacheMisses = 0;
+          const cacheMiss = (id) => {
+            cacheMisses++;
+            return id;
+          };
+          export default function App(t0) {
+            const $ = useMemoCache(2);
+            const {foo, bar} = t0;
+            let t1;
+            if ($[0] !== foo) {
+              $[0] = t1 = cacheMiss({foo});
+            } else {
+              t1 = $[1];
+            }
+            let t2;
+            if ($[1] !== bar) {
+              $[1] = t2 = cacheMiss({bar});
+            } else {
+              t2 = $[1];
+            }
+            return <h1>{cacheMisses}</h1>;
+          }
+        `);
+        expect(container.firstChild).toBe(el);
+        // cache size is constant but the cache is cleared anyway as the component in question was
+        // fast refreshed. this can occur in the case where the cache size incidentally happens to
+        // stay constant even though the source code was changed.
+        expect(el.textContent).toBe('2');
+      }
+    });
+
+    it('resets useMemoCache cache slots when cache size changes', async () => {
       if (__DEV__) {
         await render(`
           const useMemoCache = require('react/compiler-runtime').c;