Remove getCacheForType from dispatcher

Author: eps1lon

packages/react-reconciler/src/ReactFiberAsyncDispatcher.js

@@ -7,7 +7,7 @@
  * @flow
  */
 
-import type {AsyncDispatcher, Fiber} from './ReactInternalTypes';
+import type {AsyncCache, AsyncDispatcher, Fiber} from './ReactInternalTypes';
 import type {Cache} from './ReactFiberCacheComponent';
 
 import {enableCache} from 'shared/ReactFeatureFlags';
@@ -18,28 +18,18 @@ import {disableStringRefs} from 'shared/ReactFeatureFlags';
 
 import {current as currentOwner} from './ReactCurrentFiber';
 
-function getActiveCache(): Map<Function, mixed> {
-  const cache: Cache = readContext(CacheContext);
-
-  return cache.data;
-}
-
-function getCacheForType<T>(resourceType: () => T): T {
+function getActiveCache(): AsyncCache {
   if (!enableCache) {
     throw new Error('Not implemented.');
   }
-  const cache = getActiveCache();
-  let cacheForType: T | void = (cache.get(resourceType): any);
-  if (cacheForType === undefined) {
-    cacheForType = resourceType();
-    cache.set(resourceType, cacheForType);
-  }
-  return cacheForType;
+
+  const cache: Cache = readContext(CacheContext);
+
+  return cache.data;
 }
 
 export const DefaultAsyncDispatcher: AsyncDispatcher = ({
   getActiveCache,
-  getCacheForType,
 }: any);
 
 if (__DEV__ || !disableStringRefs) {

packages/react-reconciler/src/ReactFiberCacheComponent.js

@@ -8,7 +8,7 @@
  */
 
 import type {ReactContext} from 'shared/ReactTypes';
-import type {Fiber} from 'react-reconciler/src/ReactInternalTypes';
+import type {AsyncCache, Fiber} from 'react-reconciler/src/ReactInternalTypes';
 
 import {enableCache} from 'shared/ReactFeatureFlags';
 import {REACT_CONTEXT_TYPE} from 'shared/ReactSymbols';
@@ -42,7 +42,7 @@ const AbortControllerLocal: typeof AbortController = enableCache
 
 export type Cache = {
   controller: AbortController,
-  data: Map<() => mixed, mixed>,
+  data: AsyncCache,
   refCount: number,
 };
 

packages/react-reconciler/src/ReactInternalTypes.js

@@ -451,9 +451,13 @@ export type Dispatcher = {
   ) => [Awaited<S>, (P) => void, boolean],
 };
 
+export interface AsyncCache {
+  get(resourceType: Function): mixed;
+  set(resourceType: Function, value: mixed): AsyncCache;
+}
+
 export type AsyncDispatcher = {
-  getActiveCache: () => Map<Function, mixed> | null,
-  getCacheForType: <T>(resourceType: () => T) => T,
+  getActiveCache: () => AsyncCache,
   // DEV-only (or !disableStringRefs)
   getOwner: () => null | Fiber | ReactComponentInfo | ComponentStackNode,
 };

packages/react-server/src/ReactFizzAsyncDispatcher.js

@@ -7,19 +7,22 @@
  * @flow
  */
 
-import type {AsyncDispatcher} from 'react-reconciler/src/ReactInternalTypes';
+import type {
+  AsyncCache,
+  AsyncDispatcher,
+} from 'react-reconciler/src/ReactInternalTypes';
 import type {ComponentStackNode} from './ReactFizzComponentStack';
 
 import {disableStringRefs} from 'shared/ReactFeatureFlags';
 
 import {currentTaskInDEV} from './ReactFizzCurrentTask';
 
-function getCacheForType<T>(resourceType: () => T): T {
+function getActiveCache(): AsyncCache {
   throw new Error('Not implemented.');
 }
 
 export const DefaultAsyncDispatcher: AsyncDispatcher = ({
-  getCacheForType,
+  getActiveCache,
 }: any);
 
 if (__DEV__) {

packages/react-server/src/flight/ReactFlightAsyncDispatcher.js

@@ -9,7 +9,10 @@
 
 import type {ReactComponentInfo} from 'shared/ReactTypes';
 
-import type {AsyncDispatcher} from 'react-reconciler/src/ReactInternalTypes';
+import type {
+  AsyncCache,
+  AsyncDispatcher,
+} from 'react-reconciler/src/ReactInternalTypes';
 
 import {resolveRequest, getCache} from '../ReactFlightServer';
 import ReactSharedInternals from '../ReactSharedInternalsServer';
@@ -28,38 +31,42 @@ function resolveCache(): Map<Function, mixed> {
   return new Map();
 }
 
-export const DefaultAsyncDispatcher: AsyncDispatcher = ({
-  getActiveCache(): Map<Function, mixed> | null {
-    const request = resolveRequest();
-    if (request) {
-      return getCache(request);
-    }
-    return null;
-  },
-  getCacheForType<T>(resourceType: () => T): T {
-    const outerCache: Map<Function, mixed> | null =
-      previousAsyncDispatcher !== null
-        ? previousAsyncDispatcher.getActiveCache()
-        : null;
-    if (outerCache !== null) {
-      const entry: T | void = (outerCache.get(resourceType): any);
-      if (entry !== undefined) {
-        return entry;
+function getActiveCache(): AsyncCache {
+  const outerCache: AsyncCache | null =
+    previousAsyncDispatcher !== null
+      ? previousAsyncDispatcher.getActiveCache()
+      : null;
+
+  const innerCache = resolveCache();
+
+  if (outerCache === null) {
+    return innerCache;
+  }
+
+  // If both caches are active, reads will prefer the outer cache
+  // Writes will go into both caches.
+  const chainedCache: AsyncCache = {
+    get(resourceType: Function) {
+      const outerEntry = outerCache.get(resourceType);
+      if (outerEntry !== undefined) {
+        return outerEntry;
       }
-    }
+      return innerCache.get(resourceType);
+    },
+    set(resourceType: Function, value: mixed) {
+      if (outerCache !== null) {
+        outerCache.set(resourceType, value);
+      }
+      innerCache.set(resourceType, value);
+      return chainedCache;
+    },
+  };
 
-    const cache = resolveCache();
-    let entry: T | void = (cache.get(resourceType): any);
-    if (entry === undefined) {
-      entry = resourceType();
-      // TODO: Warn if undefined?
-      cache.set(resourceType, entry);
-    }
-    if (outerCache !== null) {
-      outerCache.set(resourceType, entry);
-    }
-    return entry;
-  },
+  return chainedCache;
+}
+
+export const DefaultAsyncDispatcher: AsyncDispatcher = ({
+  getActiveCache,
 }: any);
 
 if (__DEV__) {

packages/react/src/ReactCacheImpl.js

@@ -60,9 +60,15 @@ export function cache<A: Iterable<mixed>, T>(fn: (...A) => T): (...A) => T {
       // $FlowFixMe[incompatible-call]: We don't want to use rest arguments since we transpile the code.
       return fn.apply(null, arguments);
     }
-    const fnMap: WeakMap<any, CacheNode<T>> = dispatcher.getCacheForType(
+    const activeCache = dispatcher.getActiveCache();
+    let fnMap: WeakMap<any, CacheNode<T>> | void = (activeCache.get(
       createCacheRoot,
-    );
+    ): any);
+    if (fnMap === undefined) {
+      fnMap = createCacheRoot();
+      // TODO: Warn if undefined?
+      activeCache.set(createCacheRoot, fnMap);
+    }
     const fnNode = fnMap.get(fn);
     let cacheNode: CacheNode<T>;
     if (fnNode === undefined) {

packages/react/src/ReactHooks.js

@@ -53,7 +53,14 @@ export function getCacheForType<T>(resourceType: () => T): T {
     // If there is no dispatcher, then we treat this as not being cached.
     return resourceType();
   }
-  return dispatcher.getCacheForType(resourceType);
+  const activeCache = dispatcher.getActiveCache();
+  let entry: T | void = (activeCache.get(resourceType): any);
+  if (entry === undefined) {
+    entry = resourceType();
+    // TODO: Warn if undefined?
+    activeCache.set(resourceType, entry);
+  }
+  return entry;
 }
 
 export function useContext<T>(Context: ReactContext<T>): T {