react-hmr-hack.tsx

import { detectReactBuildType, getRDTHook, type ReactRenderer } from "bippy";

const devRenderers: ReactRenderer[] = [];
const originalScheduleRefreshFns: ReactRenderer["scheduleRefresh"][] = [];
const originalScheduleUpdateFns: ReactRenderer["scheduleUpdate"][] = [];
const originalScheduleRootFns: ReactRenderer["scheduleRoot"][] = [];
let isHMRDisabled = false;

interface QueuedUpdate {
  args: any[];
  functionName: "scheduleRefresh" | "scheduleUpdate" | "scheduleRoot";
  rendererIndex: number;
}

const queuedUpdates: QueuedUpdate[] = [];

export const disableHMR = () => {
  const rdtHook = getRDTHook();

  for (const renderer of rdtHook.renderers.values()) {
    if (detectReactBuildType(renderer) === "development") {
      devRenderers.push(renderer);
    }
  }

  if (devRenderers.length === 0) return;

  if (isHMRDisabled) return;

  for (const renderer of devRenderers) {
    originalScheduleRefreshFns.push(renderer.scheduleRefresh);
    originalScheduleUpdateFns.push(renderer.scheduleUpdate);
    originalScheduleRootFns.push(renderer.scheduleRoot);
  }

  for (let i = 0; i < devRenderers.length; i++) {
    const renderer = devRenderers[i];
    const createQueuedUpdateFn =
      (functionName: "scheduleRefresh" | "scheduleUpdate" | "scheduleRoot") =>
      (...args: any[]) => {
        queuedUpdates.push({ args, functionName, rendererIndex: i });
      };
    renderer.scheduleRefresh = createQueuedUpdateFn("scheduleRefresh");
    renderer.scheduleUpdate = createQueuedUpdateFn("scheduleUpdate");
    renderer.scheduleRoot = createQueuedUpdateFn("scheduleRoot");
  }

  isHMRDisabled = true;
};

export const enableHMR = () => {
  if (devRenderers.length === 0) return;
  if (!isHMRDisabled) return;
  for (let i = 0; i < devRenderers.length; i++) {
    const renderer = devRenderers[i];
    renderer.scheduleRefresh = originalScheduleRefreshFns[i];
    renderer.scheduleUpdate = originalScheduleUpdateFns[i];
    renderer.scheduleRoot = originalScheduleRootFns[i];
  }

  isHMRDisabled = false;

  flushQueuedUpdates();
};

export const flushQueuedUpdates = () => {
  for (const { args, functionName, rendererIndex } of queuedUpdates) {
    if (functionName === "scheduleRefresh") {
      originalScheduleRefreshFns[rendererIndex]?.apply(
        devRenderers[rendererIndex],
        args as any,
      );
    }
    if (functionName === "scheduleUpdate") {
      originalScheduleUpdateFns[rendererIndex]?.apply(
        devRenderers[rendererIndex],
        args as any,
      );
    }
    if (functionName === "scheduleRoot") {
      originalScheduleRootFns[rendererIndex]?.apply(
        devRenderers[rendererIndex],
        args as any,
      );
    }
  }
  queuedUpdates.length = 0;
};

export const getQueuedUpdates = () => {
  return queuedUpdates;
};