Refactor pending mutable source expiration time to be per root, not p…

…er root + per source

packages/react-reconciler/src/ReactFiberHooks.js

@@ -879,7 +879,7 @@ function useMutableSourceImpl<Source, Snapshot>(
   // Is it safe to read from this source during the current render?
   let isSafeToReadFromSource = false;
 
-  if (pendingExpirationTime !== null) {
+  if (pendingExpirationTime !== NoWork) {
     // If the source has pending updates, we can use the current render's expiration
     // time to determine if it's safe to read again from the source.
     isSafeToReadFromSource =
@@ -958,12 +958,11 @@ function useMutableSourceImpl<Source, Snapshot>(
 
     const create = () => {
       const scheduleUpdate = () => {
-        const alreadyScheduledExpirationTime = root.mutableSourcePendingUpdateMap.get(
-          source,
-        );
+        const alreadyScheduledExpirationTime =
+          root.mutableSourcePendingUpdateTime;
 
         // If an update is already scheduled for this source, re-use the same priority.
-        if (alreadyScheduledExpirationTime !== undefined) {
+        if (alreadyScheduledExpirationTime !== NoWork) {
           scheduleUpdateOnFiber(fiber, alreadyScheduledExpirationTime);
         } else {
           const currentTime = requestCurrentTimeForUpdate();
@@ -978,7 +977,7 @@ function useMutableSourceImpl<Source, Snapshot>(
           // Make sure reads during future renders will know there's a pending update.
           // This will prevent a higher priority update from reading a newer version of the source,
           // and causing a tear between that render and previous renders.
-          root.mutableSourcePendingUpdateMap.set(source, expirationTime);
+          root.mutableSourcePendingUpdateTime = expirationTime;
         }
       };
 

packages/react-reconciler/src/ReactFiberRoot.js

@@ -15,7 +15,6 @@ import type {Thenable} from './ReactFiberWorkLoop';
 import type {Interaction} from 'scheduler/src/Tracing';
 import type {SuspenseHydrationCallbacks} from './ReactFiberSuspenseComponent';
 import type {ReactPriorityLevel} from './SchedulerWithReactIntegration';
-import type {MutableSourcePendingUpdateMap} from 'shared/ReactTypes';
 
 import {noTimeout} from './ReactFiberHostConfig';
 import {createHostRootFiber} from './ReactFiber';
@@ -78,7 +77,7 @@ type BaseFiberRootProperties = {|
   lastExpiredTime: ExpirationTime,
   // Used by useMutableSource hook to avoid tearing within this root
   // when external, mutable sources are read from during render.
-  mutableSourcePendingUpdateMap: MutableSourcePendingUpdateMap,
+  mutableSourcePendingUpdateTime: ExpirationTime,
 |};
 
 // The following attributes are only used by interaction tracing builds.
@@ -128,7 +127,7 @@ function FiberRootNode(containerInfo, tag, hydrate) {
   this.nextKnownPendingLevel = NoWork;
   this.lastPingedTime = NoWork;
   this.lastExpiredTime = NoWork;
-  this.mutableSourcePendingUpdateMap = new Map();
+  this.mutableSourcePendingUpdateTime = NoWork;
 
   if (enableSchedulerTracing) {
     this.interactionThreadID = unstable_getThreadID();

packages/react-reconciler/src/ReactMutableSource.js

@@ -12,6 +12,7 @@ import type {FiberRoot} from 'react-reconciler/src/ReactFiberRoot';
 import type {MutableSource, MutableSourceVersion} from 'shared/ReactTypes';
 
 import {isPrimaryRenderer} from './ReactFiberHostConfig';
+import {NoWork} from './ReactFiberExpirationTime';
 
 // Work in progress version numbers only apply to a single render,
 // and should be reset before starting a new render.
@@ -23,14 +24,9 @@ export function clearPendingUpdates(
   root: FiberRoot,
   expirationTime: ExpirationTime,
 ): void {
-  // Remove pending mutable source entries that we've completed processing.
-  root.mutableSourcePendingUpdateMap.forEach(
-    (pendingExpirationTime, mutableSource) => {
-      if (pendingExpirationTime <= expirationTime) {
-        root.mutableSourcePendingUpdateMap.delete(mutableSource);
-      }
-    },
-  );
+  if (root.mutableSourcePendingUpdateTime <= expirationTime) {
+    root.mutableSourcePendingUpdateTime = NoWork;
+  }
 }
 
 export function resetWorkInProgressVersions(): void {
@@ -52,9 +48,8 @@ export function resetWorkInProgressVersions(): void {
 export function getPendingExpirationTime(
   root: FiberRoot,
   source: MutableSource<any>,
-): ExpirationTime | null {
-  const expirationTime = root.mutableSourcePendingUpdateMap.get(source);
-  return expirationTime !== undefined ? expirationTime : null;
+): ExpirationTime {
+  return root.mutableSourcePendingUpdateTime;
 }
 
 export function getWorkInProgressVersion(

packages/react-reconciler/src/__tests__/ReactHooksWithNoopRenderer-test.internal.js

@@ -2084,11 +2084,7 @@ describe('ReactHooksWithNoopRenderer', () => {
       source.value = 'two';
 
       // Render work should restart and the updated value should be used
-      expect(Scheduler).toFlushAndYieldThrough([
-        'a:two',
-        'b:two',
-        'Sync effect',
-      ]);
+      expect(Scheduler).toFlushAndYield(['a:two', 'b:two', 'Sync effect']);
     });
 
     it('should schedule an update if a new source is mutated between render and commit (subscription)', () => {

packages/shared/ReactTypes.js

@@ -7,8 +7,6 @@
  * @flow
  */
 
-import type {ExpirationTime} from 'react-reconciler/src/ReactFiberExpirationTime';
-
 export type ReactNode =
   | React$Element<any>
   | ReactPortal
@@ -194,16 +192,6 @@ export type ReactScopeInstance = {|
 // so long as it changes every time any part of the source changes.
 export type MutableSourceVersion = $NonMaybeType<mixed>;
 
-// Tracks expiration time for all mutable sources with pending updates.
-// Used to determine if a source is safe to read during updates.
-// If there are no entries in this map for a given source,
-// or if the current render’s expiration time is ≤ this value,
-// it is safe to read from the source without tearing.
-export type MutableSourcePendingUpdateMap = Map<
-  MutableSource<any>,
-  ExpirationTime,
->;
-
 export type MutableSource<Source: $NonMaybeType<mixed>> = {|
   _source: Source,