Add cache() API (#25506)

Like memo() but longer lived.

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

@@ -2,7 +2,6 @@ let React;
 let ReactNoop;
 let Cache;
 let getCacheSignal;
-let getCacheForType;
 let Scheduler;
 let act;
 let Suspense;
@@ -10,8 +9,10 @@ let Offscreen;
 let useCacheRefresh;
 let startTransition;
 let useState;
+let cache;
 
-let caches;
+let getTextCache;
+let textCaches;
 let seededCache;
 
 describe('ReactCache', () => {
@@ -24,66 +25,68 @@ describe('ReactCache', () => {
     Scheduler = require('scheduler');
     act = require('jest-react').act;
     Suspense = React.Suspense;
+    cache = React.experimental_cache;
     Offscreen = React.unstable_Offscreen;
     getCacheSignal = React.unstable_getCacheSignal;
-    getCacheForType = React.unstable_getCacheForType;
     useCacheRefresh = React.unstable_useCacheRefresh;
     startTransition = React.startTransition;
     useState = React.useState;
 
-    caches = [];
+    textCaches = [];
     seededCache = null;
-  });
 
-  function createTextCache() {
+    if (gate(flags => flags.enableCache)) {
-    if (seededCache !== null) {
+      getTextCache = cache(() => {
-      // Trick to seed a cache before it exists.
+        if (seededCache !== null) {
-      // TODO: Need a built-in API to seed data before the initial render (i.e.
+          // Trick to seed a cache before it exists.
-      // not a refresh because nothing has mounted yet).
+          // TODO: Need a built-in API to seed data before the initial render (i.e.
-      const cache = seededCache;
+          // not a refresh because nothing has mounted yet).
-      seededCache = null;
+          const textCache = seededCache;
-      return cache;
+          seededCache = null;
+          return textCache;
+        }
+
+        const data = new Map();
+        const version = textCaches.length + 1;
+        const textCache = {
+          version,
+          data,
+          resolve(text) {
+            const record = data.get(text);
+            if (record === undefined) {
+              const newRecord = {
+                status: 'resolved',
+                value: text,
+                cleanupScheduled: false,
+              };
+              data.set(text, newRecord);
+            } else if (record.status === 'pending') {
+              record.value.resolve();
+            }
+          },
+          reject(text, error) {
+            const record = data.get(text);
+            if (record === undefined) {
+              const newRecord = {
+                status: 'rejected',
+                value: error,
+                cleanupScheduled: false,
+              };
+              data.set(text, newRecord);
+            } else if (record.status === 'pending') {
+              record.value.reject();
+            }
+          },
+        };
+        textCaches.push(textCache);
+        return textCache;
+      });
     }
-
+  });
-    const data = new Map();
-    const version = caches.length + 1;
-    const cache = {
-      version,
-      data,
-      resolve(text) {
-        const record = data.get(text);
-        if (record === undefined) {
-          const newRecord = {
-            status: 'resolved',
-            value: text,
-            cleanupScheduled: false,
-          };
-          data.set(text, newRecord);
-        } else if (record.status === 'pending') {
-          record.value.resolve();
-        }
-      },
-      reject(text, error) {
-        const record = data.get(text);
-        if (record === undefined) {
-          const newRecord = {
-            status: 'rejected',
-            value: error,
-            cleanupScheduled: false,
-          };
-          data.set(text, newRecord);
-        } else if (record.status === 'pending') {
-          record.value.reject();
-        }
-      },
-    };
-    caches.push(cache);
-    return cache;
-  }
 
   function readText(text) {
     const signal = getCacheSignal();
-    const textCache = getCacheForType(createTextCache);
+    const textCache = getTextCache();
     const record = textCache.data.get(text);
     if (record !== undefined) {
       if (!record.cleanupScheduled) {
@@ -160,18 +163,18 @@ describe('ReactCache', () => {
 
   function seedNextTextCache(text) {
     if (seededCache === null) {
-      seededCache = createTextCache();
+      seededCache = getTextCache();
     }
     seededCache.resolve(text);
   }
 
   function resolveMostRecentTextCache(text) {
-    if (caches.length === 0) {
+    if (textCaches.length === 0) {
       throw Error('Cache does not exist.');
     } else {
       // Resolve the most recently created cache. An older cache can by
-      // resolved with `caches[index].resolve(text)`.
+      // resolved with `textCaches[index].resolve(text)`.
-      caches[caches.length - 1].resolve(text);
+      textCaches[textCaches.length - 1].resolve(text);
     }
   }
 
@@ -815,9 +818,18 @@ describe('ReactCache', () => {
 
   // @gate experimental || www
   test('refresh a cache with seed data', async () => {
-    let refresh;
+    let refreshWithSeed;
     function App() {
-      refresh = useCacheRefresh();
+      const refresh = useCacheRefresh();
+      const [seed, setSeed] = useState({fn: null});
+      if (seed.fn) {
+        seed.fn();
+        seed.fn = null;
+      }
+      refreshWithSeed = fn => {
+        setSeed({fn});
+        refresh();
+      };
       return <AsyncText showVersion={true} text="A" />;
     }
 
@@ -845,11 +857,14 @@ describe('ReactCache', () => {
     await act(async () => {
       // Refresh the cache with seeded data, like you would receive from a
       // server mutation.
-      // TODO: Seeding multiple typed caches. Should work by calling `refresh`
+      // TODO: Seeding multiple typed textCaches. Should work by calling `refresh`
       // multiple times with different key/value pairs
-      const cache = createTextCache();
+      startTransition(() =>
-      cache.resolve('A');
+        refreshWithSeed(() => {
-      startTransition(() => refresh(createTextCache, cache));
+          const textCache = getTextCache();
+          textCache.resolve('A');
+        }),
+      );
     });
     // The root should re-render without a cache miss.
     // The cache is not cleared up yet, since it's still reference by the root
@@ -1624,4 +1639,152 @@ describe('ReactCache', () => {
     expect(Scheduler).toHaveYielded(['More']);
     expect(root).toMatchRenderedOutput(<div hidden={true}>More</div>);
   });
+
+  // @gate enableCache
+  it('cache objects and primitive arguments and a mix of them', async () => {
+    const root = ReactNoop.createRoot();
+    const types = cache((a, b) => ({a: typeof a, b: typeof b}));
+    function Print({a, b}) {
+      return types(a, b).a + ' ' + types(a, b).b + ' ';
+    }
+    function Same({a, b}) {
+      const x = types(a, b);
+      const y = types(a, b);
+      return (x === y).toString() + ' ';
+    }
+    function FlippedOrder({a, b}) {
+      return (types(a, b) === types(b, a)).toString() + ' ';
+    }
+    function FewerArgs({a, b}) {
+      return (types(a, b) === types(a)).toString() + ' ';
+    }
+    function MoreArgs({a, b}) {
+      return (types(a) === types(a, b)).toString() + ' ';
+    }
+    await act(async () => {
+      root.render(
+        <>
+          <Print a="e" b="f" />
+          <Same a="a" b="b" />
+          <FlippedOrder a="c" b="d" />
+          <FewerArgs a="e" b="f" />
+          <MoreArgs a="g" b="h" />
+        </>,
+      );
+    });
+    expect(root).toMatchRenderedOutput('string string true false false false ');
+    await act(async () => {
+      root.render(
+        <>
+          <Print a="e" b={null} />
+          <Same a="a" b={null} />
+          <FlippedOrder a="c" b={null} />
+          <FewerArgs a="e" b={null} />
+          <MoreArgs a="g" b={null} />
+        </>,
+      );
+    });
+    expect(root).toMatchRenderedOutput('string object true false false false ');
+    const obj = {};
+    await act(async () => {
+      root.render(
+        <>
+          <Print a="e" b={obj} />
+          <Same a="a" b={obj} />
+          <FlippedOrder a="c" b={obj} />
+          <FewerArgs a="e" b={obj} />
+          <MoreArgs a="g" b={obj} />
+        </>,
+      );
+    });
+    expect(root).toMatchRenderedOutput('string object true false false false ');
+    const sameObj = {};
+    await act(async () => {
+      root.render(
+        <>
+          <Print a={sameObj} b={sameObj} />
+          <Same a={sameObj} b={sameObj} />
+          <FlippedOrder a={sameObj} b={sameObj} />
+          <FewerArgs a={sameObj} b={sameObj} />
+          <MoreArgs a={sameObj} b={sameObj} />
+        </>,
+      );
+    });
+    expect(root).toMatchRenderedOutput('object object true true false false ');
+    const objA = {};
+    const objB = {};
+    await act(async () => {
+      root.render(
+        <>
+          <Print a={objA} b={objB} />
+          <Same a={objA} b={objB} />
+          <FlippedOrder a={objA} b={objB} />
+          <FewerArgs a={objA} b={objB} />
+          <MoreArgs a={objA} b={objB} />
+        </>,
+      );
+    });
+    expect(root).toMatchRenderedOutput('object object true false false false ');
+    const sameSymbol = Symbol();
+    await act(async () => {
+      root.render(
+        <>
+          <Print a={sameSymbol} b={sameSymbol} />
+          <Same a={sameSymbol} b={sameSymbol} />
+          <FlippedOrder a={sameSymbol} b={sameSymbol} />
+          <FewerArgs a={sameSymbol} b={sameSymbol} />
+          <MoreArgs a={sameSymbol} b={sameSymbol} />
+        </>,
+      );
+    });
+    expect(root).toMatchRenderedOutput('symbol symbol true true false false ');
+    const notANumber = +'nan';
+    await act(async () => {
+      root.render(
+        <>
+          <Print a={1} b={notANumber} />
+          <Same a={1} b={notANumber} />
+          <FlippedOrder a={1} b={notANumber} />
+          <FewerArgs a={1} b={notANumber} />
+          <MoreArgs a={1} b={notANumber} />
+        </>,
+      );
+    });
+    expect(root).toMatchRenderedOutput('number number true false false false ');
+  });
+
+  // @gate enableCache
+  it('cached functions that throw should cache the error', async () => {
+    const root = ReactNoop.createRoot();
+    const throws = cache(v => {
+      throw new Error(v);
+    });
+    let x;
+    let y;
+    let z;
+    function Test() {
+      try {
+        throws(1);
+      } catch (e) {
+        x = e;
+      }
+      try {
+        throws(1);
+      } catch (e) {
+        y = e;
+      }
+      try {
+        throws(2);
+      } catch (e) {
+        z = e;
+      }
+
+      return 'Blank';
+    }
+    await act(async () => {
+      root.render(<Test />);
+    });
+    expect(x).toBe(y);
+    expect(z).not.toBe(x);
+  });
 });

packages/react/index.classic.fb.js

@@ -32,6 +32,7 @@ export {
   isValidElement,
   lazy,
   memo,
+  experimental_cache,
   startTransition,
   startTransition as unstable_startTransition, // TODO: Remove once call sights updated to startTransition
   unstable_Cache,

packages/react/index.experimental.js

@@ -29,6 +29,7 @@ export {
   isValidElement,
   lazy,
   memo,
+  experimental_cache,
   startTransition,
   unstable_Cache,
   unstable_DebugTracingMode,

packages/react/index.js

@@ -54,6 +54,7 @@ export {
   isValidElement,
   lazy,
   memo,
+  experimental_cache,
   startTransition,
   unstable_Cache,
   unstable_DebugTracingMode,

packages/react/index.modern.fb.js

@@ -31,6 +31,7 @@ export {
   isValidElement,
   lazy,
   memo,
+  experimental_cache,
   startTransition,
   startTransition as unstable_startTransition, // TODO: Remove once call sights updated to startTransition
   unstable_Cache,

packages/react/src/React.js

@@ -35,6 +35,7 @@ import {createContext} from './ReactContext';
 import {lazy} from './ReactLazy';
 import {forwardRef} from './ReactForwardRef';
 import {memo} from './ReactMemo';
+import {cache} from './ReactCache';
 import {
   getCacheSignal,
   getCacheForType,
@@ -100,6 +101,7 @@ export {
   forwardRef,
   lazy,
   memo,
+  cache as experimental_cache,
   useCallback,
   useContext,
   useEffect,

packages/react/src/ReactCache.js

@@ -0,0 +1,126 @@
+/**
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ *
+ * @flow
+ */
+
+import ReactCurrentCache from './ReactCurrentCache';
+
+const UNTERMINATED = 0;
+const TERMINATED = 1;
+const ERRORED = 2;
+
+type UnterminatedCacheNode<T> = {
+  s: 0,
+  v: void,
+  o: null | WeakMap<Function | Object, CacheNode<T>>,
+  p: null | Map<string | number | null | void | symbol | boolean, CacheNode<T>>,
+};
+
+type TerminatedCacheNode<T> = {
+  s: 1,
+  v: T,
+  o: null | WeakMap<Function | Object, CacheNode<T>>,
+  p: null | Map<string | number | null | void | symbol | boolean, CacheNode<T>>,
+};
+
+type ErroredCacheNode<T> = {
+  s: 2,
+  v: mixed,
+  o: null | WeakMap<Function | Object, CacheNode<T>>,
+  p: null | Map<string | number | null | void | symbol | boolean, CacheNode<T>>,
+};
+
+type CacheNode<T> =
+  | TerminatedCacheNode<T>
+  | UnterminatedCacheNode<T>
+  | ErroredCacheNode<T>;
+
+function createCacheRoot<T>(): WeakMap<Function | Object, CacheNode<T>> {
+  return new WeakMap();
+}
+
+function createCacheNode<T>(): CacheNode<T> {
+  return {
+    s: UNTERMINATED, // status, represents whether the cached computation returned a value or threw an error
+    v: undefined, // value, either the cached result or an error, depending on s
+    o: null, // object cache, a WeakMap where non-primitive arguments are stored
+    p: null, // primitive cache, a regular Map where primitive arguments are stored.
+  };
+}
+
+export function cache<A: Iterable<mixed>, T>(fn: (...A) => T): (...A) => T {
+  return function() {
+    const dispatcher = ReactCurrentCache.current;
+    if (!dispatcher) {
+      // If there is no dispatcher, then we treat this as not being cached.
+      // $FlowFixMe: We don't want to use rest arguments since we transpile the code.
+      return fn.apply(null, arguments);
+    }
+    const fnMap = dispatcher.getCacheForType(createCacheRoot);
+    const fnNode = fnMap.get(fn);
+    let cacheNode: CacheNode<T>;
+    if (fnNode === undefined) {
+      cacheNode = createCacheNode();
+      fnMap.set(fn, cacheNode);
+    } else {
+      cacheNode = fnNode;
+    }
+    for (let i = 0, l = arguments.length; i < l; i++) {
+      const arg = arguments[i];
+      if (
+        typeof arg === 'function' ||
+        (typeof arg === 'object' && arg !== null)
+      ) {
+        // Objects go into a WeakMap
+        let objectCache = cacheNode.o;
+        if (objectCache === null) {
+          cacheNode.o = objectCache = new WeakMap();
+        }
+        const objectNode = objectCache.get(arg);
+        if (objectNode === undefined) {
+          cacheNode = createCacheNode();
+          objectCache.set(arg, cacheNode);
+        } else {
+          cacheNode = objectNode;
+        }
+      } else {
+        // Primitives go into a regular Map
+        let primitiveCache = cacheNode.p;
+        if (primitiveCache === null) {
+          cacheNode.p = primitiveCache = new Map();
+        }
+        const primitiveNode = primitiveCache.get(arg);
+        if (primitiveNode === undefined) {
+          cacheNode = createCacheNode();
+          primitiveCache.set(arg, cacheNode);
+        } else {
+          cacheNode = primitiveNode;
+        }
+      }
+    }
+    if (cacheNode.s === TERMINATED) {
+      return cacheNode.v;
+    }
+    if (cacheNode.s === ERRORED) {
+      throw cacheNode.v;
+    }
+    try {
+      // $FlowFixMe: We don't want to use rest arguments since we transpile the code.
+      const result = fn.apply(null, arguments);
+      const terminatedNode: TerminatedCacheNode<T> = (cacheNode: any);
+      terminatedNode.s = TERMINATED;
+      terminatedNode.v = result;
+      return result;
+    } catch (error) {
+      // We store the first error that's thrown and rethrow it.
+      const erroredNode: ErroredCacheNode<T> = (cacheNode: any);
+      erroredNode.s = ERRORED;
+      erroredNode.v = error;
+      throw error;
+    }
+  };
+}

packages/react/src/ReactSharedSubset.experimental.js

@@ -24,6 +24,7 @@ export {
   isValidElement,
   lazy,
   memo,
+  experimental_cache,
   startTransition,
   unstable_DebugTracingMode,
   unstable_getCacheSignal,