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

@@ -1771,10 +1771,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, <App />);
     // 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

@@ -886,7 +886,6 @@ describe('ReactDOMServerSelectiveHydration', () => {
     document.body.removeChild(container);
   });
 
-  // @gate experimental || www
   it('hydrates the last explicitly hydrated target at higher priority', async () => {
     function Child({text}) {
       Scheduler.unstable_yieldValue(text);
@@ -920,15 +919,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

@@ -56,7 +56,6 @@ import {
   setAttemptDiscreteHydration,
   setAttemptContinuousHydration,
   setAttemptHydrationAtCurrentPriority,
-  queueExplicitHydrationTarget,
   setGetCurrentUpdatePriority,
   setAttemptHydrationAtPriority,
 } from '../events/ReactDOMEventReplaying';
@@ -116,12 +115,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>,
@@ -196,7 +189,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,
@@ -72,7 +74,9 @@ function ReactDOMRoot(internalRoot: FiberRoot) {
   this._internalRoot = internalRoot;
 }
 
-ReactDOMRoot.prototype.render = function(children: ReactNodeList): void {
+ReactDOMHydrationRoot.prototype.render = ReactDOMRoot.prototype.render = function(
+  children: ReactNodeList,
+): void {
   const root = this._internalRoot;
   if (root === null) {
     throw new Error('Cannot update an unmounted root.');
@@ -104,7 +108,7 @@ ReactDOMRoot.prototype.render = function(children: ReactNodeList): void {
   updateContainer(children, root, null, null);
 };
 
-ReactDOMRoot.prototype.unmount = function(): void {
+ReactDOMHydrationRoot.prototype.unmount = ReactDOMRoot.prototype.unmount = function(): void {
   if (__DEV__) {
     if (typeof arguments[0] === 'function') {
       console.error(
@@ -189,6 +193,16 @@ export function createRoot(
   return new ReactDOMRoot(root);
 }
 
+function ReactDOMHydrationRoot(internalRoot: FiberRoot) {
+  this._internalRoot = internalRoot;
+}
+function scheduleHydration(target: Node) {
+  if (target) {
+    queueExplicitHydrationTarget(target);
+  }
+}
+ReactDOMHydrationRoot.prototype.unstable_scheduleHydration = scheduleHydration;
+
 export function hydrateRoot(
   container: Container,
   initialChildren: ReactNodeList,
@@ -236,7 +250,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 {