[Experiment] Reuse memo cache after interruption

packages/react-reconciler/src/ReactFiberHooks.js

@@ -46,6 +46,7 @@ import {
   enableAsyncActions,
   enableUseDeferredValueInitialArg,
   disableLegacyMode,
+  enableNoCloningMemoCache,
 } from 'shared/ReactFeatureFlags';
 import {
   REACT_CONTEXT_TYPE,
@@ -1130,7 +1131,32 @@ function useMemoCache(size: number): Array<any> {
         const currentMemoCache: ?MemoCache = currentUpdateQueue.memoCache;
         if (currentMemoCache != null) {
           memoCache = {
-            data: currentMemoCache.data.map(array => array.slice()),
+            // When enableNoCloningMemoCache is enabled, instead of treating the
+            // cache as copy-on-write, like we do with fibers, we share the same
+            // cache instance across all render attempts, even if the component
+            // is interrupted before it commits.
+            //
+            // If an update is interrupted, either because it suspended or
+            // because of another update, we can reuse the memoized computations
+            // from the previous attempt. We can do this because the React
+            // Compiler performs atomic writes to the memo cache, i.e. it will
+            // not record the inputs to a memoization without also recording its
+            // output.
+            //
+            // This gives us a form of "resuming" within components and hooks.
+            //
+            // This only works when updating a component that already mounted.
+            // It has no impact during initial render, because the memo cache is
+            // stored on the fiber, and since we have not implemented resuming
+            // for fibers, it's always a fresh memo cache, anyway.
+            //
+            // However, this alone is pretty useful — it happens whenever you
+            // update the UI with fresh data after a mutation/action, which is
+            // extremely common in a Suspense-driven (e.g. RSC or Relay) app.
+            data: enableNoCloningMemoCache
+              ? currentMemoCache.data
+              : // Clone the memo cache before each render (copy-on-write)
+                currentMemoCache.data.map(array => array.slice()),
             index: 0,
           };
         }

packages/react-reconciler/src/__tests__/useMemoCache-test.js

@@ -10,7 +10,9 @@
 
 let React;
 let ReactNoop;
+let Scheduler;
 let act;
+let assertLog;
 let useState;
 let useMemoCache;
 let MemoCacheSentinel;
@@ -22,7 +24,9 @@ describe('useMemoCache()', () => {
 
     React = require('react');
     ReactNoop = require('react-noop-renderer');
+    Scheduler = require('scheduler');
     act = require('internal-test-utils').act;
+    assertLog = require('internal-test-utils').assertLog;
     useState = React.useState;
     useMemoCache = React.unstable_useMemoCache;
     MemoCacheSentinel = Symbol.for('react.memo_cache_sentinel');
@@ -363,4 +367,258 @@ describe('useMemoCache()', () => {
     expect(Text).toBeCalledTimes(3);
     expect(data).toBe(data1); // confirm that the cache persisted across renders
   });
+
+  // @gate enableUseMemoCacheHook
+  test('reuses computations from suspended/interrupted render attempts during an update', async () => {
+    // This test demonstrates the benefit of a shared memo cache. By "shared" I
+    // mean multiple concurrent render attempts of the same component/hook use
+    // the same cache. (When the feature flag is off, we don't do this — the
+    // cache is copy-on-write.)
+    //
+    // If an update is interrupted, either because it suspended or because of
+    // another update, we can reuse the memoized computations from the previous
+    // attempt. We can do this because the React Compiler performs atomic writes
+    // to the memo cache, i.e. it will not record the inputs to a memoization
+    // without also recording its output.
+    //
+    // This gives us a form of "resuming" within components and hooks.
+    //
+    // This only works when updating a component that already mounted. It has no
+    // impact during initial render, because the memo cache is stored on the
+    // fiber, and since we have not implemented resuming for fibers, it's always
+    // a fresh memo cache, anyway.
+    //
+    // However, this alone is pretty useful — it happens whenever you update the
+    // UI with fresh data after a mutation/action, which is extremely common in
+    // a Suspense-driven (e.g. RSC or Relay) app. That's the scenario that this
+    // test simulates.
+    //
+    // So the impact of this feature is faster data mutations/actions.
+
+    function someExpensiveProcessing(t) {
+      Scheduler.log(`Some expensive processing... [${t}]`);
+      return t;
+    }
+
+    function useWithLog(t, msg) {
+      try {
+        return React.use(t);
+      } catch (x) {
+        Scheduler.log(`Suspend! [${msg}]`);
+        throw x;
+      }
+    }
+
+    // Original code:
+    //
+    //   function Data({chunkA, chunkB}) {
+    //     const a = someExpensiveProcessing(useWithLog(chunkA, 'chunkA'));
+    //     const b = useWithLog(chunkB, 'chunkB');
+    //     return (
+    //       <>
+    //         {a}
+    //         {b}
+    //       </>
+    //     );
+    //   }
+    //
+    //   function Input() {
+    //     const [input, _setText] = useState('');
+    //     return input;
+    //   }
+    //
+    //   function App({chunkA, chunkB}) {
+    //     return (
+    //       <>
+    //         <div>
+    //           Input: <Input />
+    //         </div>
+    //         <div>
+    //           Data: <Data chunkA={chunkA} chunkB={chunkB} />
+    //         </div>
+    //       </>
+    //     );
+    //   }
+    function Data(t0) {
+      const $ = useMemoCache(5);
+      const {chunkA, chunkB} = t0;
+      const t1 = useWithLog(chunkA, 'chunkA');
+      let t2;
+
+      if ($[0] !== t1) {
+        t2 = someExpensiveProcessing(t1);
+        $[0] = t1;
+        $[1] = t2;
+      } else {
+        t2 = $[1];
+      }
+
+      const a = t2;
+      const b = useWithLog(chunkB, 'chunkB');
+      let t3;
+
+      if ($[2] !== a || $[3] !== b) {
+        t3 = (
+          <>
+            {a}
+            {b}
+          </>
+        );
+        $[2] = a;
+        $[3] = b;
+        $[4] = t3;
+      } else {
+        t3 = $[4];
+      }
+
+      return t3;
+    }
+
+    let setInput;
+    function Input() {
+      const [input, _set] = useState('');
+      setInput = _set;
+      return input;
+    }
+
+    function App(t0) {
+      const $ = useMemoCache(4);
+      const {chunkA, chunkB} = t0;
+      let t1;
+
+      if ($[0] === Symbol.for('react.memo_cache_sentinel')) {
+        t1 = (
+          <div>
+            Input: <Input />
+          </div>
+        );
+        $[0] = t1;
+      } else {
+        t1 = $[0];
+      }
+
+      let t2;
+
+      if ($[1] !== chunkA || $[2] !== chunkB) {
+        t2 = (
+          <>
+            {t1}
+            <div>
+              Data: <Data chunkA={chunkA} chunkB={chunkB} />
+            </div>
+          </>
+        );
+        $[1] = chunkA;
+        $[2] = chunkB;
+        $[3] = t2;
+      } else {
+        t2 = $[3];
+      }
+
+      return t2;
+    }
+
+    function createInstrumentedResolvedPromise(value) {
+      return {
+        then() {},
+        status: 'fulfilled',
+        value,
+      };
+    }
+
+    function createDeferred() {
+      let resolve;
+      const p = new Promise(res => {
+        resolve = res;
+      });
+      p.resolve = resolve;
+      return p;
+    }
+
+    // Initial render. We pass the data in as two separate "chunks" to simulate
+    // a stream (e.g. RSC).
+    const root = ReactNoop.createRoot();
+    const initialChunkA = createInstrumentedResolvedPromise('A1');
+    const initialChunkB = createInstrumentedResolvedPromise('B1');
+    await act(() =>
+      root.render(<App chunkA={initialChunkA} chunkB={initialChunkB} />),
+    );
+    assertLog(['Some expensive processing... [A1]']);
+    expect(root).toMatchRenderedOutput(
+      <>
+        <div>Input: </div>
+        <div>Data: A1B1</div>
+      </>,
+    );
+
+    // Update the UI in a transition. This would happen after a data mutation.
+    const updatedChunkA = createDeferred();
+    const updatedChunkB = createDeferred();
+    await act(() => {
+      React.startTransition(() => {
+        root.render(<App chunkA={updatedChunkA} chunkB={updatedChunkB} />);
+      });
+    });
+    assertLog(['Suspend! [chunkA]']);
+
+    // The data starts to stream in. Loading the data in the first chunk
+    // triggers an expensive computation in the UI. Later, we'll test whether
+    // this computation is reused.
+    await act(() => updatedChunkA.resolve('A2'));
+    assertLog(['Some expensive processing... [A2]', 'Suspend! [chunkB]']);
+
+    // The second chunk hasn't loaded yet, so we're still showing the
+    // initial UI.
+    expect(root).toMatchRenderedOutput(
+      <>
+        <div>Input: </div>
+        <div>Data: A1B1</div>
+      </>,
+    );
+
+    // While waiting for the data to finish loading, update a different part of
+    // the screen. This interrupts the refresh transition.
+    //
+    // In a real app, this might be an input or hover event.
+    await act(() => setInput('hi!'));
+
+    // Once the input has updated, we go back to rendering the transition.
+    if (gate(flags => flags.enableNoCloningMemoCache)) {
+      // We did not have process the first chunk again. We reused the
+      // computation from the earlier attempt.
+      assertLog(['Suspend! [chunkB]']);
+    } else {
+      // Because we clone/reset the memo cache after every aborted attempt, we
+      // must process the first chunk again.
+      assertLog(['Some expensive processing... [A2]', 'Suspend! [chunkB]']);
+    }
+
+    expect(root).toMatchRenderedOutput(
+      <>
+        <div>Input: hi!</div>
+        <div>Data: A1B1</div>
+      </>,
+    );
+
+    // Finish loading the data.
+    await act(() => updatedChunkB.resolve('B2'));
+    if (gate(flags => flags.enableNoCloningMemoCache)) {
+      // We did not have process the first chunk again. We reused the
+      // computation from the earlier attempt.
+      assertLog([]);
+    } else {
+      // Because we clone/reset the memo cache after every aborted attempt, we
+      // must process the first chunk again.
+      //
+      // That's three total times we've processed the first chunk, compared to
+      // just once when enableNoCloningMemoCache is on.
+      assertLog(['Some expensive processing... [A2]']);
+    }
+    expect(root).toMatchRenderedOutput(
+      <>
+        <div>Input: hi!</div>
+        <div>Data: A2B2</div>
+      </>,
+    );
+  });
 });

packages/shared/ReactFeatureFlags.js

@@ -105,6 +105,8 @@ export const enableCPUSuspense = __EXPERIMENTAL__;
 // Enables unstable_useMemoCache hook, intended as a compilation target for
 // auto-memoization.
 export const enableUseMemoCacheHook = __EXPERIMENTAL__;
+// Test this at Meta before enabling.
+export const enableNoCloningMemoCache = false;
 
 export const enableUseEffectEventHook = __EXPERIMENTAL__;
 

packages/shared/forks/ReactFeatureFlags.native-fb.js

@@ -62,6 +62,7 @@ export const enableSuspenseAvoidThisFallback = false;
 export const enableSuspenseAvoidThisFallbackFizz = false;
 export const enableCPUSuspense = true;
 export const enableUseMemoCacheHook = true;
+export const enableNoCloningMemoCache = false;
 export const enableUseEffectEventHook = false;
 export const favorSafetyOverHydrationPerf = true;
 export const enableLegacyFBSupport = false;

packages/shared/forks/ReactFeatureFlags.native-oss.js

@@ -62,6 +62,7 @@ export const debugRenderPhaseSideEffectsForStrictMode = __DEV__;
 
 // TODO: decide on React 19
 export const enableUseMemoCacheHook = false;
+export const enableNoCloningMemoCache = false;
 export const enableUseDeferredValueInitialArg = __EXPERIMENTAL__;
 
 // -----------------------------------------------------------------------------

packages/shared/forks/ReactFeatureFlags.test-renderer.js

@@ -37,6 +37,7 @@ export const enableSuspenseAvoidThisFallback = false;
 export const enableSuspenseAvoidThisFallbackFizz = false;
 export const enableCPUSuspense = false;
 export const enableUseMemoCacheHook = true;
+export const enableNoCloningMemoCache = false;
 export const enableUseEffectEventHook = false;
 export const favorSafetyOverHydrationPerf = true;
 export const enableComponentStackLocations = true;

packages/shared/forks/ReactFeatureFlags.test-renderer.native-fb.js

@@ -42,6 +42,7 @@ export const enableSuspenseAvoidThisFallback = false;
 export const enableSuspenseAvoidThisFallbackFizz = false;
 export const enableCPUSuspense = true;
 export const enableUseMemoCacheHook = true;
+export const enableNoCloningMemoCache = false;
 export const enableUseEffectEventHook = false;
 export const favorSafetyOverHydrationPerf = true;
 export const enableInfiniteRenderLoopDetection = false;

packages/shared/forks/ReactFeatureFlags.test-renderer.www.js

@@ -38,6 +38,7 @@ export const enableSuspenseAvoidThisFallback = true;
 export const enableSuspenseAvoidThisFallbackFizz = false;
 export const enableCPUSuspense = false;
 export const enableUseMemoCacheHook = true;
+export const enableNoCloningMemoCache = false;
 export const enableUseEffectEventHook = false;
 export const favorSafetyOverHydrationPerf = true;
 export const enableComponentStackLocations = true;

packages/shared/forks/ReactFeatureFlags.www-dynamic.js

@@ -27,6 +27,7 @@ export const enableRefAsProp = __VARIANT__;
 export const enableRetryLaneExpiration = __VARIANT__;
 export const favorSafetyOverHydrationPerf = __VARIANT__;
 export const disableDefaultPropsExceptForClasses = __VARIANT__;
+export const enableNoCloningMemoCache = __VARIANT__;
 export const retryLaneExpirationMs = 5000;
 export const syncLaneExpirationMs = 250;
 export const transitionLaneExpirationMs = 5000;

packages/shared/forks/ReactFeatureFlags.www.js

@@ -34,6 +34,7 @@ export const {
   enableRefAsProp,
   favorSafetyOverHydrationPerf,
   disableDefaultPropsExceptForClasses,
+  enableNoCloningMemoCache,
 } = dynamicFeatureFlags;
 
 // On WWW, __EXPERIMENTAL__ is used for a new modern build.