Move unstable_scheduleHydration to ReactDOMHydrationRoot

packages/react-dom/index.classic.fb.js

@@ -34,6 +34,5 @@ export {
   unstable_isNewReconciler,
   unstable_renderSubtreeIntoContainer,
   unstable_runWithPriority, // DO NOT USE: Temporarily exposed to migrate off of Scheduler.runWithPriority.
-  unstable_scheduleHydration,
   version,
 } from './src/client/ReactDOM';

packages/react-dom/index.experimental.js

@@ -21,6 +21,5 @@ export {
   unstable_flushControlled,
   unstable_renderSubtreeIntoContainer,
   unstable_runWithPriority, // DO NOT USE: Temporarily exposed to migrate off of Scheduler.runWithPriority.
-  unstable_scheduleHydration,
   version,
 } from './src/client/ReactDOM';

packages/react-dom/index.js

@@ -25,6 +25,5 @@ export {
   unstable_isNewReconciler,
   unstable_renderSubtreeIntoContainer,
   unstable_runWithPriority, // DO NOT USE: Temporarily exposed to migrate off of Scheduler.runWithPriority.
-  unstable_scheduleHydration,
   version,
 } from './src/client/ReactDOM';

packages/react-dom/index.modern.fb.js

@@ -18,6 +18,5 @@ export {
   unstable_flushControlled,
   unstable_isNewReconciler,
   unstable_runWithPriority, // DO NOT USE: Temporarily exposed to migrate off of Scheduler.runWithPriority.
-  unstable_scheduleHydration,
   version,
 } from './src/client/ReactDOM';

packages/react-dom/src/__tests__/ReactDOMServerPartialHydration-test.internal.js

@@ -1561,10 +1561,9 @@ describe('ReactDOMServerPartialHydration', () => {
     const b = container.getElementsByTagName('span')[1];
     expect(b.textContent).toBe('B');
 
-    const root = ReactDOM.createRoot(container, {hydrate: true});
-
+    const root = ReactDOM.hydrateRoot(container);
     // Increase hydration priority to higher than "offscreen".
-    ReactDOM.unstable_scheduleHydration(b);
+    root.unstable_scheduleHydration(b);
 
     suspend = true;
 

packages/react-dom/src/__tests__/ReactDOMServerSelectiveHydration-test.internal.js

@@ -844,15 +844,14 @@ describe('ReactDOMServerSelectiveHydration', () => {
     const spanB = container.getElementsByTagName('span')[1];
     const spanC = container.getElementsByTagName('span')[2];
 
-    const root = ReactDOM.createRoot(container, {hydrate: true});
-    root.render(<App />);
+    const root = ReactDOM.hydrateRoot(container, <App />);
 
     // Nothing has been hydrated so far.
     expect(Scheduler).toHaveYielded([]);
 
     // Increase priority of B and then C.
-    ReactDOM.unstable_scheduleHydration(spanB);
-    ReactDOM.unstable_scheduleHydration(spanC);
+    root.unstable_scheduleHydration(spanB);
+    root.unstable_scheduleHydration(spanC);
 
     // We should prioritize hydrating C first because the last added
     // gets highest priority followed by the next added.

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

@@ -57,7 +57,6 @@ import {
   setAttemptDiscreteHydration,
   setAttemptContinuousHydration,
   setAttemptHydrationAtCurrentPriority,
-  queueExplicitHydrationTarget,
   setGetCurrentUpdatePriority,
   setAttemptHydrationAtPriority,
 } from '../events/ReactDOMEventReplaying';
@@ -117,12 +116,6 @@ function createPortal(
   return createPortalImpl(children, container, null, key);
 }
 
-function scheduleHydration(target: Node) {
-  if (target) {
-    queueExplicitHydrationTarget(target);
-  }
-}
-
 function renderSubtreeIntoContainer(
   parentComponent: React$Component<any, any>,
   element: React$Element<any>,
@@ -197,7 +190,6 @@ export {
   createRoot,
   hydrateRoot,
   flushControlled as unstable_flushControlled,
-  scheduleHydration as unstable_scheduleHydration,
   // Disabled behind disableUnstableRenderSubtreeIntoContainer
   renderSubtreeIntoContainer as unstable_renderSubtreeIntoContainer,
   // enableCreateEventHandleAPI

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

@@ -11,6 +11,8 @@ import type {Container} from './ReactDOMHostConfig';
 import type {MutableSource, ReactNodeList} from 'shared/ReactTypes';
 import type {FiberRoot} from 'react-reconciler/src/ReactInternalTypes';
 
+import {queueExplicitHydrationTarget} from '../events/ReactDOMEventReplaying';
+
 export type RootType = {
   render(children: ReactNodeList): void,
   unmount(): void,
@@ -190,6 +192,18 @@ export function createRoot(
   return new ReactDOMRoot(root);
 }
 
+function ReactDOMHydrationRoot(internalRoot: FiberRoot) {
+  ReactDOMRoot.call(this, internalRoot);
+}
+function scheduleHydration(target: Node) {
+  if (target) {
+    queueExplicitHydrationTarget(target);
+  }
+}
+Object.assign(ReactDOMHydrationRoot.prototype, ReactDOMRoot.prototype, {
+  unstable_scheduleHydration: scheduleHydration,
+});
+
 export function hydrateRoot(
   container: Container,
   initialChildren: ReactNodeList,
@@ -237,7 +251,7 @@ export function hydrateRoot(
   // Render the initial children
   updateContainer(initialChildren, root, null, null);
 
-  return new ReactDOMRoot(root);
+  return new ReactDOMHydrationRoot(root);
 }
 
 export function isValidContainer(node: any): boolean {