React refresh component swap

packages/bippy/src/core.ts

@@ -14,10 +14,12 @@ import type {
 
 import {
   BIPPY_INSTRUMENTATION_STRING,
+  ensureReactRefreshHandler,
   getRDTHook,
   hasRDTHook,
   isReactRefresh,
   isRealReactDevtools,
+  linkReactRefreshFamily,
 } from './rdt-hook.js';
 
 // https://github.com/facebook/react/blob/main/packages/react-reconciler/src/ReactWorkTags.js
@@ -1149,6 +1151,9 @@ export const instrument = (
       priority: number | void,
     ) => {
       if (prevOnCommitFiberRoot === nextOnCommitFiberRoot) return;
+      if (!_fiberRoots.has(root)) {
+        _fiberRoots.add(root);
+      }
       // TODO: validate whether the bottom version is more correct here
       // for preventing infinite loops
       // if (rdtHook.onCommitFiberRoot !== handler) return;
@@ -1214,6 +1219,43 @@ export const INSTALL_ERROR = new Error();
 
 export const _fiberRoots = new Set<FiberRoot>();
 
+export interface HmrSwapComponentOptions {
+  shouldRemount?: boolean;
+}
+
+export const hmrSwapComponent = (
+  previousType: unknown,
+  nextType: unknown,
+  options: HmrSwapComponentOptions = {},
+): boolean => {
+  const family = linkReactRefreshFamily(previousType, nextType);
+  if (!family) return false;
+
+  family.current = nextType;
+
+  const updatedFamilies = options.shouldRemount ? new Set() : new Set([family]);
+  const staleFamilies = options.shouldRemount ? new Set([family]) : new Set();
+
+  const rootsToRefresh = new Set<FiberRoot>(_fiberRoots);
+  if (!rootsToRefresh.size) return false;
+
+  const rdtHook = getRDTHook();
+
+  let didSchedule = false;
+  for (const renderer of rdtHook.renderers.values()) {
+    try {
+      ensureReactRefreshHandler(renderer);
+      if (!renderer.scheduleRefresh) continue;
+      for (const root of rootsToRefresh) {
+        renderer.scheduleRefresh(root, { staleFamilies, updatedFamilies });
+      }
+      didSchedule = true;
+    } catch {}
+  }
+
+  return didSchedule;
+};
+
 export const secure = (
   options: InstrumentationOptions,
   secureOptions: {

packages/bippy/src/rdt-hook.ts

@@ -3,7 +3,7 @@
 // without this, we can't stub the React DevTools global hook, we don't have a way to instrument the application
 // make sure you import this file first before anything else (particularly React)
 
-import type { ReactDevToolsGlobalHook, ReactRenderer } from './types.js';
+import type { Family, ReactDevToolsGlobalHook, ReactRenderer } from './types.js';
 
 export const version = process.env.VERSION;
 export const BIPPY_INSTRUMENTATION_STRING = `bippy-${version}`;
@@ -16,6 +16,121 @@ const NO_OP = () => {
   /**/
 };
 
+const patchedRefreshRenderers = new WeakSet<object>();
+const refreshFamiliesByType = new WeakMap<object, Family>();
+const didSetRefreshHandlerByRenderer = new WeakSet<object>();
+
+const canBeRefreshKey = (value: unknown): value is object => {
+  return (
+    typeof value === 'function' || (typeof value === 'object' && value != null)
+  );
+};
+
+const hasTypeProperty = (value: unknown): value is { type: unknown } => {
+  return typeof value === 'object' && value != null && 'type' in value;
+};
+
+const hasRenderProperty = (value: unknown): value is { render: unknown } => {
+  return typeof value === 'object' && value != null && 'render' in value;
+};
+
+const getRefreshCandidates = (value: unknown): object[] => {
+  if (!canBeRefreshKey(value)) return [];
+
+  const candidates: object[] = [];
+  const visited = new Set<object>();
+
+  let current: unknown = value;
+  while (canBeRefreshKey(current) && !visited.has(current)) {
+    visited.add(current);
+    candidates.push(current);
+
+    if (hasTypeProperty(current) && canBeRefreshKey(current.type)) {
+      current = current.type;
+      continue;
+    }
+
+    if (hasRenderProperty(current) && canBeRefreshKey(current.render)) {
+      current = current.render;
+      continue;
+    }
+
+    break;
+  }
+
+  return candidates;
+};
+
+export const getReactRefreshFamily = (type: unknown): Family | null => {
+  const candidates = getRefreshCandidates(type);
+  for (const candidate of candidates) {
+    const existing = refreshFamiliesByType.get(candidate);
+    if (existing) return existing;
+  }
+  return null;
+};
+
+export const linkReactRefreshFamily = (
+  previousType: unknown,
+  nextType: unknown,
+): Family | null => {
+  const previousCandidates = getRefreshCandidates(previousType);
+  const nextCandidates = getRefreshCandidates(nextType);
+
+  const existingFamily =
+    getReactRefreshFamily(previousType) ?? getReactRefreshFamily(nextType);
+
+  const family: Family = existingFamily ?? { current: nextType };
+
+  for (const candidate of [...previousCandidates, ...nextCandidates]) {
+    refreshFamiliesByType.set(candidate, family);
+  }
+
+  return family;
+};
+
+const patchRefreshRenderer = (renderer: ReactRenderer): void => {
+  if (!canBeRefreshKey(renderer)) return;
+  if (patchedRefreshRenderers.has(renderer)) return;
+  patchedRefreshRenderers.add(renderer);
+
+  const originalSetRefreshHandler = renderer.setRefreshHandler;
+  if (typeof originalSetRefreshHandler !== 'function') return;
+
+  renderer.setRefreshHandler = (handler) => {
+    const wrappedHandler: typeof handler =
+      handler == null
+        ? null
+        : (fiberOrType) => {
+            const maybeFiberOrType: unknown = fiberOrType;
+            const maybeType =
+              hasTypeProperty(maybeFiberOrType) ? maybeFiberOrType.type : null;
+
+            return (
+              getReactRefreshFamily(maybeType ?? maybeFiberOrType) ??
+              handler(fiberOrType)
+            );
+          };
+
+    didSetRefreshHandlerByRenderer.add(renderer);
+    originalSetRefreshHandler(wrappedHandler);
+  };
+};
+
+export const ensureReactRefreshHandler = (renderer: ReactRenderer): void => {
+  if (!canBeRefreshKey(renderer)) return;
+  if (didSetRefreshHandlerByRenderer.has(renderer)) return;
+  if (typeof renderer.setRefreshHandler !== 'function') return;
+
+  renderer.setRefreshHandler((fiberOrType) => {
+    const maybeFiberOrType: unknown = fiberOrType;
+    const maybeType = hasTypeProperty(maybeFiberOrType)
+      ? maybeFiberOrType.type
+      : null;
+    return getReactRefreshFamily(maybeType ?? maybeFiberOrType);
+  });
+};
+
 const checkDCE = (fn: unknown): void => {
   try {
     const code = Function.prototype.toString.call(fn);
@@ -66,6 +181,7 @@ export const installRDTHook = (
       const nextID = ++i;
       renderers.set(nextID, renderer);
       _renderers.add(renderer);
+      patchRefreshRenderer(renderer);
       if (!rdtHook._instrumentationIsActive) {
         rdtHook._instrumentationIsActive = true;
         onActiveListeners.forEach((listener) => listener());
@@ -94,6 +210,7 @@ export const installRDTHook = (
           if (ourRenderers.size > 0) {
             ourRenderers.forEach((renderer, id) => {
               _renderers.add(renderer);
+              patchRefreshRenderer(renderer);
               // eslint-disable-next-line @typescript-eslint/no-unsafe-call
               newHook.renderers.set(id, renderer);
             });
@@ -168,6 +285,7 @@ export const patchRDTHook = (onActive?: () => unknown): void => {
       rdtHook.inject = (renderer) => {
         const id = prevInject(renderer);
         _renderers.add(renderer);
+        patchRefreshRenderer(renderer);
         if (isRefresh) {
           // react refresh doesn't inject this properly
           // https://github.com/facebook/react/blob/18eaf51bd51fed8dfed661d64c306759101d0bfd/packages/react-refresh/src/ReactFreshRuntime.js#L430

packages/bippy/src/test/hmr-swap-component.test.tsx

@@ -0,0 +1,58 @@
+import '../index.js'; // KEEP THIS LINE ON TOP
+
+/* eslint-disable @typescript-eslint/no-unsafe-call */
+/* eslint-disable @typescript-eslint/no-require-imports */
+
+import { fireEvent, render, waitFor } from '@testing-library/react';
+import React from 'react';
+import { expect, it, vi } from 'vitest';
+
+import { hmrSwapComponent, instrument } from '../index.js';
+
+declare global {
+  interface Window {
+    $RefreshReg$?: () => void;
+    $RefreshSig$?: () => <T>(type: T) => T;
+  }
+}
+
+const runtime = require('react-refresh/runtime');
+runtime.injectIntoGlobalHook(window);
+window.$RefreshReg$ = () => {};
+window.$RefreshSig$ = () => (type) => type;
+
+it('can swap components with different refs via refresh', async () => {
+  instrument({ onCommitFiberRoot: vi.fn() });
+
+  const OldCounter = () => {
+    const [count, setCount] = React.useState(0);
+    return (
+      <button type="button" onClick={() => setCount(count + 1)}>
+        count:{count}
+      </button>
+    );
+  };
+
+  const NewCounter = () => {
+    const [count, setCount] = React.useState(0);
+    return (
+      <button type="button" onClick={() => setCount(count + 1)}>
+        count:{count + 1}
+      </button>
+    );
+  };
+
+  const result = render(<OldCounter />);
+  const button = result.getByRole('button');
+
+  fireEvent.click(button);
+  expect(button.textContent).toBe('count:1');
+
+  const didSwap = hmrSwapComponent(OldCounter, NewCounter);
+  expect(didSwap).toBe(true);
+
+  await waitFor(() => {
+    expect(button.textContent).toBe('count:2');
+  });
+});
+