[Fresh] Track mounted roots via DevTools Hook

packages/react-reconciler/src/ReactFiberHotReloading.js

@@ -38,16 +38,32 @@ export type Family = {|
   current: any,
 |};
 
-export type HotUpdate = {|
-  resolveFamily: (any => Family | void) | null,
+export type RefreshUpdate = {|
   staleFamilies: Set<Family>,
   updatedFamilies: Set<Family>,
 |};
 
-let resolveFamily: (any => Family | void) | null = null;
+// Resolves type to a family.
+type RefreshHandler = any => Family | void;
+
+// Used by React Refresh runtime through DevTools Global Hook.
+export type SetRefreshHandler = (handler: RefreshHandler | null) => void;
+export type ScheduleRefresh = (root: FiberRoot, update: RefreshUpdate) => void;
+export type FindHostInstancesForRefresh = (
+  root: FiberRoot,
+  families: Array<Family>,
+) => Set<Instance>;
+
+let resolveFamily: RefreshHandler | null = null;
 // $FlowFixMe Flow gets confused by a WeakSet feature check below.
 let failedBoundaries: WeakSet<Fiber> | null = null;
 
+export let setRefreshHandler = (handler: RefreshHandler | null): void => {
+  if (__DEV__) {
+    resolveFamily = handler;
+  }
+};
+
 export function resolveFunctionForHotReloading(type: any): any {
   if (__DEV__) {
     if (resolveFamily === null) {
@@ -192,7 +208,7 @@ export function isCompatibleFamilyForHotReloading(
 export function markFailedErrorBoundaryForHotReloading(fiber: Fiber) {
   if (__DEV__) {
     if (resolveFamily === null) {
-      // Not hot reloading.
+      // Hot reloading is disabled.
       return;
     }
     if (typeof WeakSet !== 'function') {
@@ -205,12 +221,16 @@ export function markFailedErrorBoundaryForHotReloading(fiber: Fiber) {
   }
 }
 
-export function scheduleHotUpdate(root: FiberRoot, hotUpdate: HotUpdate): void {
+export let scheduleRefresh: ScheduleRefresh = (
+  root: FiberRoot,
+  update: RefreshUpdate,
+): void => {
   if (__DEV__) {
-    // TODO: warn if its identity changes over time?
-    resolveFamily = hotUpdate.resolveFamily;
-
-    const {staleFamilies, updatedFamilies} = hotUpdate;
+    if (resolveFamily === null) {
+      // Hot reloading is disabled.
+      return;
+    }
+    const {staleFamilies, updatedFamilies} = update;
     flushPassiveEffects();
     flushSync(() => {
       scheduleFibersWithFamiliesRecursively(
@@ -220,7 +240,7 @@ export function scheduleHotUpdate(root: FiberRoot, hotUpdate: HotUpdate): void {
       );
     });
   }
-}
+};
 
 function scheduleFibersWithFamiliesRecursively(
   fiber: Fiber,
@@ -292,10 +312,10 @@ function scheduleFibersWithFamiliesRecursively(
   }
 }
 
-export function findHostInstancesForHotUpdate(
+export let findHostInstancesForRefresh: FindHostInstancesForRefresh = (
   root: FiberRoot,
   families: Array<Family>,
-): Set<Instance> {
+): Set<Instance> => {
   if (__DEV__) {
     const hostInstances = new Set();
     const types = new Set(families.map(family => family.current));
@@ -307,10 +327,10 @@ export function findHostInstancesForHotUpdate(
     return hostInstances;
   } else {
     throw new Error(
-      'Did not expect findHostInstancesForHotUpdate to be called in production.',
+      'Did not expect findHostInstancesForRefresh to be called in production.',
     );
   }
-}
+};
 
 function findHostInstancesForMatchingFibersRecursively(
   fiber: Fiber,

packages/react-reconciler/src/ReactFiberReconciler.js

@@ -71,8 +71,9 @@ import {Sync} from './ReactFiberExpirationTime';
 import {revertPassiveEffectsChange} from 'shared/ReactFeatureFlags';
 import {requestCurrentSuspenseConfig} from './ReactFiberSuspenseConfig';
 import {
-  scheduleHotUpdate,
-  findHostInstancesForHotUpdate,
+  scheduleRefresh,
+  setRefreshHandler,
+  findHostInstancesForRefresh,
 } from './ReactFiberHotReloading';
 
 type OpaqueRoot = FiberRoot;
@@ -475,10 +476,6 @@ export function injectIntoDevTools(devToolsConfig: DevToolsConfig): boolean {
 
   return injectInternals({
     ...devToolsConfig,
-    findHostInstancesForHotUpdate: __DEV__
-      ? findHostInstancesForHotUpdate
-      : null,
-    scheduleHotUpdate: __DEV__ ? scheduleHotUpdate : null,
     overrideHookState,
     overrideProps,
     setSuspenseHandler,
@@ -498,5 +495,9 @@ export function injectIntoDevTools(devToolsConfig: DevToolsConfig): boolean {
       }
       return findFiberByHostInstance(instance);
     },
+    // React Refresh
+    findHostInstancesForRefresh: __DEV__ ? findHostInstancesForRefresh : null,
+    scheduleRefresh: __DEV__ ? scheduleRefresh : null,
+    setRefreshHandler: __DEV__ ? setRefreshHandler : null,
   });
 }