Add back useMutableSource temporarily (#22396)

Recoil uses useMutableSource behind a flag. I thought this was fine
because Recoil isn't used in any concurrent roots, so the behavior
would be the same, but it turns out that it is used by concurrent
roots in a few places.

I'm not expecting it to be hard to migrate to useSyncExternalStore, but
to de-risk the change I'm going to roll it out gradually with a flag. In
the meantime, I've added back the useMutableSource API.

packages/react-debug-tools/src/ReactDebugHooks.js

@@ -7,7 +7,13 @@
  * @flow
  */
 
-import type {ReactContext, ReactProviderType} from 'shared/ReactTypes';
+import type {
+  MutableSource,
+  MutableSourceGetSnapshotFn,
+  MutableSourceSubscribeFn,
+  ReactContext,
+  ReactProviderType,
+} from 'shared/ReactTypes';
 import type {
   Fiber,
   Dispatcher as DispatcherType,
@@ -255,6 +261,23 @@ function useMemo<T>(
   return value;
 }
 
+function useMutableSource<Source, Snapshot>(
+  source: MutableSource<Source>,
+  getSnapshot: MutableSourceGetSnapshotFn<Source, Snapshot>,
+  subscribe: MutableSourceSubscribeFn<Source, Snapshot>,
+): Snapshot {
+  // useMutableSource() composes multiple hooks internally.
+  // Advance the current hook index the same number of times
+  // so that subsequent hooks have the right memoized state.
+  nextHook(); // MutableSource
+  nextHook(); // State
+  nextHook(); // Effect
+  nextHook(); // Effect
+  const value = getSnapshot(source._source);
+  hookLog.push({primitive: 'MutableSource', stackError: new Error(), value});
+  return value;
+}
+
 function useSyncExternalStore<T>(
   subscribe: (() => void) => () => void,
   getSnapshot: () => T,
@@ -335,6 +358,7 @@ const Dispatcher: DispatcherType = {
   useRef,
   useState,
   useTransition,
+  useMutableSource,
   useSyncExternalStore,
   useDeferredValue,
   useOpaqueIdentifier,

packages/react-debug-tools/src/__tests__/ReactHooksInspectionIntegration-test.js

@@ -1019,6 +1019,43 @@ describe('ReactHooksInspectionIntegration', () => {
     ]);
   });
 
+  it('should support composite useMutableSource hook', () => {
+    const createMutableSource =
+      React.createMutableSource || React.unstable_createMutableSource;
+    const useMutableSource =
+      React.useMutableSource || React.unstable_useMutableSource;
+
+    const mutableSource = createMutableSource({}, () => 1);
+    function Foo(props) {
+      useMutableSource(
+        mutableSource,
+        () => 'snapshot',
+        () => {},
+      );
+      React.useMemo(() => 'memo', []);
+      return <div />;
+    }
+    const renderer = ReactTestRenderer.create(<Foo />);
+    const childFiber = renderer.root.findByType(Foo)._currentFiber();
+    const tree = ReactDebugTools.inspectHooksOfFiber(childFiber);
+    expect(tree).toEqual([
+      {
+        id: 0,
+        isStateEditable: false,
+        name: 'MutableSource',
+        value: 'snapshot',
+        subHooks: [],
+      },
+      {
+        id: 1,
+        isStateEditable: false,
+        name: 'Memo',
+        value: 'memo',
+        subHooks: [],
+      },
+    ]);
+  });
+
   // @gate experimental || www
   it('should support composite useSyncExternalStore hook', () => {
     const useSyncExternalStore = React.unstable_useSyncExternalStore;

packages/react-dom/src/client/ReactDOMRoot.js

@@ -8,7 +8,7 @@
  */
 
 import type {Container} from './ReactDOMHostConfig';
-import type {ReactNodeList} from 'shared/ReactTypes';
+import type {MutableSource, ReactNodeList} from 'shared/ReactTypes';
 import type {FiberRoot} from 'react-reconciler/src/ReactInternalTypes';
 
 export type RootType = {
@@ -24,6 +24,7 @@ export type CreateRootOptions = {
   hydrationOptions?: {
     onHydrated?: (suspenseNode: Comment) => void,
     onDeleted?: (suspenseNode: Comment) => void,
+    mutableSources?: Array<MutableSource<any>>,
     ...
   },
   // END OF TODO
@@ -34,6 +35,7 @@ export type CreateRootOptions = {
 
 export type HydrateRootOptions = {
   // Hydration options
+  hydratedSources?: Array<MutableSource<any>>,
   onHydrated?: (suspenseNode: Comment) => void,
   onDeleted?: (suspenseNode: Comment) => void,
   // Options for all roots
@@ -59,6 +61,7 @@ import {
   createContainer,
   updateContainer,
   findHostInstanceWithNoPortals,
+  registerMutableSourceForHydration,
 } from 'react-reconciler/src/ReactFiberReconciler';
 import invariant from 'shared/invariant';
 import {ConcurrentRoot} from 'react-reconciler/src/ReactRootTags';
@@ -126,6 +129,11 @@ export function createRoot(
   const hydrate = options != null && options.hydrate === true;
   const hydrationCallbacks =
     (options != null && options.hydrationOptions) || null;
+  const mutableSources =
+    (options != null &&
+      options.hydrationOptions != null &&
+      options.hydrationOptions.mutableSources) ||
+    null;
   // END TODO
 
   const isStrictMode = options != null && options.unstable_strictMode === true;
@@ -151,6 +159,15 @@ export function createRoot(
     container.nodeType === COMMENT_NODE ? container.parentNode : container;
   listenToAllSupportedEvents(rootContainerElement);
 
+  // TODO: Delete this path
+  if (mutableSources) {
+    for (let i = 0; i < mutableSources.length; i++) {
+      const mutableSource = mutableSources[i];
+      registerMutableSourceForHydration(root, mutableSource);
+    }
+  }
+  // END TODO
+
   return new ReactDOMRoot(root);
 }
 
@@ -168,6 +185,7 @@ export function hydrateRoot(
   // For now we reuse the whole bag of options since they contain
   // the hydration callbacks.
   const hydrationCallbacks = options != null ? options : null;
+  const mutableSources = (options != null && options.hydratedSources) || null;
   const isStrictMode = options != null && options.unstable_strictMode === true;
 
   let concurrentUpdatesByDefaultOverride = null;
@@ -190,6 +208,13 @@ export function hydrateRoot(
   // This can't be a comment node since hydration doesn't work on comment nodes anyway.
   listenToAllSupportedEvents(container);
 
+  if (mutableSources) {
+    for (let i = 0; i < mutableSources.length; i++) {
+      const mutableSource = mutableSources[i];
+      registerMutableSourceForHydration(root, mutableSource);
+    }
+  }
+
   // Render the initial children
   updateContainer(initialChildren, root, null, null);
 

packages/react-dom/src/server/ReactPartialRendererHooks.js

@@ -9,7 +9,12 @@
 
 import type {Dispatcher as DispatcherType} from 'react-reconciler/src/ReactInternalTypes';
 
-import type {ReactContext} from 'shared/ReactTypes';
+import type {
+  MutableSource,
+  MutableSourceGetSnapshotFn,
+  MutableSourceSubscribeFn,
+  ReactContext,
+} from 'shared/ReactTypes';
 import type PartialRenderer from './ReactPartialRenderer';
 
 import {validateContextBounds} from './ReactPartialRendererContext';
@@ -461,6 +466,18 @@ export function useCallback<T>(
   return useMemo(() => callback, deps);
 }
 
+// TODO Decide on how to implement this hook for server rendering.
+// If a mutation occurs during render, consider triggering a Suspense boundary
+// and falling back to client rendering.
+function useMutableSource<Source, Snapshot>(
+  source: MutableSource<Source>,
+  getSnapshot: MutableSourceGetSnapshotFn<Source, Snapshot>,
+  subscribe: MutableSourceSubscribeFn<Source, Snapshot>,
+): Snapshot {
+  resolveCurrentlyRenderingComponent();
+  return getSnapshot(source._source);
+}
+
 function useSyncExternalStore<T>(
   subscribe: (() => void) => () => void,
   getSnapshot: () => T,
@@ -527,6 +544,8 @@ export const Dispatcher: DispatcherType = {
   useDeferredValue,
   useTransition,
   useOpaqueIdentifier,
+  // Subscriptions are not setup in a server environment.
+  useMutableSource,
   useSyncExternalStore,
 };
 

packages/react-reconciler/src/ReactFiberBeginWork.new.js

@@ -12,6 +12,7 @@ import type {LazyComponent as LazyComponentType} from 'react/src/ReactLazy';
 import type {Fiber, FiberRoot} from './ReactInternalTypes';
 import type {TypeOfMode} from './ReactTypeOfMode';
 import type {Lanes, Lane} from './ReactFiberLane.new';
+import type {MutableSource} from 'shared/ReactTypes';
 import type {
   SuspenseState,
   SuspenseListRenderState,
@@ -144,6 +145,7 @@ import {
   isSuspenseInstancePending,
   isSuspenseInstanceFallback,
   registerSuspenseInstanceRetry,
+  supportsHydration,
   isPrimaryRenderer,
   supportsPersistence,
   getOffscreenContainerProps,
@@ -218,6 +220,7 @@ import {
   RetryAfterError,
   NoContext,
 } from './ReactFiberWorkLoop.new';
+import {setWorkInProgressVersion} from './ReactMutableSource.new';
 import {
   requestCacheFromPool,
   pushCacheProvider,
@@ -1297,6 +1300,21 @@ function updateHostRoot(current, workInProgress, renderLanes) {
     // We always try to hydrate. If this isn't a hydration pass there won't
     // be any children to hydrate which is effectively the same thing as
     // not hydrating.
+
+    if (supportsHydration) {
+      const mutableSourceEagerHydrationData =
+        root.mutableSourceEagerHydrationData;
+      if (mutableSourceEagerHydrationData != null) {
+        for (let i = 0; i < mutableSourceEagerHydrationData.length; i += 2) {
+          const mutableSource = ((mutableSourceEagerHydrationData[
+            i
+          ]: any): MutableSource<any>);
+          const version = mutableSourceEagerHydrationData[i + 1];
+          setWorkInProgressVersion(mutableSource, version);
+        }
+      }
+    }
+
     const child = mountChildFibers(
       workInProgress,
       null,

packages/react-reconciler/src/ReactFiberBeginWork.old.js

@@ -12,6 +12,7 @@ import type {LazyComponent as LazyComponentType} from 'react/src/ReactLazy';
 import type {Fiber, FiberRoot} from './ReactInternalTypes';
 import type {TypeOfMode} from './ReactTypeOfMode';
 import type {Lanes, Lane} from './ReactFiberLane.old';
+import type {MutableSource} from 'shared/ReactTypes';
 import type {
   SuspenseState,
   SuspenseListRenderState,
@@ -144,6 +145,7 @@ import {
   isSuspenseInstancePending,
   isSuspenseInstanceFallback,
   registerSuspenseInstanceRetry,
+  supportsHydration,
   isPrimaryRenderer,
   supportsPersistence,
   getOffscreenContainerProps,
@@ -218,6 +220,7 @@ import {
   RetryAfterError,
   NoContext,
 } from './ReactFiberWorkLoop.old';
+import {setWorkInProgressVersion} from './ReactMutableSource.old';
 import {
   requestCacheFromPool,
   pushCacheProvider,
@@ -1297,6 +1300,21 @@ function updateHostRoot(current, workInProgress, renderLanes) {
     // We always try to hydrate. If this isn't a hydration pass there won't
     // be any children to hydrate which is effectively the same thing as
     // not hydrating.
+
+    if (supportsHydration) {
+      const mutableSourceEagerHydrationData =
+        root.mutableSourceEagerHydrationData;
+      if (mutableSourceEagerHydrationData != null) {
+        for (let i = 0; i < mutableSourceEagerHydrationData.length; i += 2) {
+          const mutableSource = ((mutableSourceEagerHydrationData[
+            i
+          ]: any): MutableSource<any>);
+          const version = mutableSourceEagerHydrationData[i + 1];
+          setWorkInProgressVersion(mutableSource, version);
+        }
+      }
+    }
+
     const child = mountChildFibers(
       workInProgress,
       null,

packages/react-reconciler/src/ReactFiberCompleteWork.new.js

@@ -30,6 +30,8 @@ import type {SuspenseContext} from './ReactFiberSuspenseContext.new';
 import type {OffscreenState} from './ReactFiberOffscreenComponent';
 import type {Cache, SpawnedCachePool} from './ReactFiberCacheComponent.new';
 
+import {resetWorkInProgressVersions as resetMutableSourceWorkInProgressVersions} from './ReactMutableSource.new';
+
 import {now} from './Scheduler';
 
 import {
@@ -852,6 +854,7 @@ function completeWork(
       }
       popHostContainer(workInProgress);
       popTopLevelLegacyContextObject(workInProgress);
+      resetMutableSourceWorkInProgressVersions();
       if (fiberRoot.pendingContext) {
         fiberRoot.context = fiberRoot.pendingContext;
         fiberRoot.pendingContext = null;

packages/react-reconciler/src/ReactFiberCompleteWork.old.js

@@ -30,6 +30,8 @@ import type {SuspenseContext} from './ReactFiberSuspenseContext.old';
 import type {OffscreenState} from './ReactFiberOffscreenComponent';
 import type {Cache, SpawnedCachePool} from './ReactFiberCacheComponent.old';
 
+import {resetWorkInProgressVersions as resetMutableSourceWorkInProgressVersions} from './ReactMutableSource.old';
+
 import {now} from './Scheduler';
 
 import {
@@ -852,6 +854,7 @@ function completeWork(
       }
       popHostContainer(workInProgress);
       popTopLevelLegacyContextObject(workInProgress);
+      resetMutableSourceWorkInProgressVersions();
       if (fiberRoot.pendingContext) {
         fiberRoot.context = fiberRoot.pendingContext;
         fiberRoot.pendingContext = null;