diff --git a/packages/bippy/src/core.ts b/packages/bippy/src/core.ts index d4465fc5..6670fa9c 100644 --- a/packages/bippy/src/core.ts +++ b/packages/bippy/src/core.ts @@ -1105,6 +1105,418 @@ export const overrideContext = ( } }; +interface DispatcherRef { + H?: unknown; + current?: unknown; +} + +interface Dispatcher { + useState: (initialState: S | (() => S)) => [S, (action: S | ((prev: S) => S)) => void]; + useReducer: ( + reducer: (state: S, action: A) => S, + initialArg: S, + init?: (arg: S) => S, + ) => [S, (action: A) => void]; + useSyncExternalStore: ( + subscribe: (onStoreChange: () => void) => () => void, + getSnapshot: () => T, + getServerSnapshot?: () => T, + ) => T; + [key: string]: unknown; +} + +interface HookQueue { + pending?: unknown; + dispatch?: (...args: unknown[]) => void; + lastRenderedReducer?: unknown; + lanes?: number; + value?: unknown; + getSnapshot?: () => unknown; +} + +interface HookState { + memoizedState: unknown; + baseState: unknown; + baseQueue: unknown; + queue: HookQueue | null; + next: HookState | null; +} + +interface PausedQueueState { + originalPendingDescriptor?: PropertyDescriptor; + originalGetSnapshot?: () => unknown; + snapshotValueAtPause?: unknown; +} + +interface PausedDispatcherState { + dispatcherKey: 'H' | 'current'; + originalDescriptor: PropertyDescriptor | undefined; +} + +interface PausedContextState { + originalDescriptor?: PropertyDescriptor; + frozenValue: unknown; +} + +let isUpdatesPaused = false; +const pausedDispatcherStates = new Map(); +const pausedQueueStates = new WeakMap(); +const pausedContextStates = new WeakMap, PausedContextState>(); + +const pauseHookQueue = (queue: HookQueue): void => { + if (!queue || pausedQueueStates.has(queue)) return; + + const queuePauseState: PausedQueueState = {}; + + if ('pending' in queue) { + queuePauseState.originalPendingDescriptor = Object.getOwnPropertyDescriptor(queue, 'pending'); + let currentPendingValue = queue.pending; + + Object.defineProperty(queue, 'pending', { + configurable: true, + enumerable: true, + get() { + return currentPendingValue; + }, + set(newValue) { + if (isUpdatesPaused) return; + currentPendingValue = newValue; + }, + }); + } + + if ('getSnapshot' in queue && typeof queue.getSnapshot === 'function') { + queuePauseState.originalGetSnapshot = queue.getSnapshot; + queuePauseState.snapshotValueAtPause = queue.getSnapshot(); + + queue.getSnapshot = () => { + if (isUpdatesPaused) { + return queuePauseState.snapshotValueAtPause; + } + return queuePauseState.originalGetSnapshot!(); + }; + } + + pausedQueueStates.set(queue, queuePauseState); +}; + +const resumeHookQueue = (queue: HookQueue): void => { + const queuePauseState = pausedQueueStates.get(queue); + if (!queuePauseState) return; + + if (queuePauseState.originalPendingDescriptor) { + const currentPendingValue = queue.pending; + Object.defineProperty(queue, 'pending', queuePauseState.originalPendingDescriptor); + if (!queuePauseState.originalPendingDescriptor.get && !queuePauseState.originalPendingDescriptor.set) { + queue.pending = currentPendingValue; + } + } else if ('pending' in queue) { + const currentPendingValue = queue.pending; + delete (queue as Record).pending; + queue.pending = currentPendingValue; + } + + if (queuePauseState.originalGetSnapshot) { + queue.getSnapshot = queuePauseState.originalGetSnapshot; + } + + pausedQueueStates.delete(queue); +}; + +const pauseContextDependency = (contextDependency: ContextDependency): void => { + if (pausedContextStates.has(contextDependency)) return; + + const frozenValue = contextDependency.memoizedValue; + const originalDescriptor = Object.getOwnPropertyDescriptor(contextDependency, 'memoizedValue'); + + const pausedState: PausedContextState = { + originalDescriptor, + frozenValue, + }; + + Object.defineProperty(contextDependency, 'memoizedValue', { + configurable: true, + enumerable: true, + get() { + if (isUpdatesPaused) { + return pausedState.frozenValue; + } + if (originalDescriptor?.get) { + return originalDescriptor.get.call(this); + } + return (this as unknown as { _memoizedValue?: unknown })._memoizedValue; + }, + set(newValue) { + if (isUpdatesPaused) return; + if (originalDescriptor?.set) { + originalDescriptor.set.call(this, newValue); + } else { + (this as unknown as { _memoizedValue: unknown })._memoizedValue = newValue; + } + }, + }); + + pausedContextStates.set(contextDependency, pausedState); +}; + +const resumeContextDependency = (contextDependency: ContextDependency): void => { + const pausedState = pausedContextStates.get(contextDependency); + if (!pausedState) return; + + if (pausedState.originalDescriptor) { + Object.defineProperty(contextDependency, 'memoizedValue', pausedState.originalDescriptor); + } else { + delete (contextDependency as unknown as Record).memoizedValue; + } + + pausedContextStates.delete(contextDependency); +}; + +const pauseFiberContextDependencies = (fiber: Fiber): void => { + if (!fiber.dependencies) return; + + try { + let contextDependency = fiber.dependencies.firstContext as ContextDependency | null; + + while (contextDependency && typeof contextDependency === 'object' && 'memoizedValue' in contextDependency) { + pauseContextDependency(contextDependency); + contextDependency = contextDependency.next; + } + } catch {} +}; + +const resumeFiberContextDependencies = (fiber: Fiber): void => { + if (!fiber.dependencies) return; + + try { + let contextDependency = fiber.dependencies.firstContext as ContextDependency | null; + + while (contextDependency && typeof contextDependency === 'object' && 'memoizedValue' in contextDependency) { + resumeContextDependency(contextDependency); + contextDependency = contextDependency.next; + } + } catch {} +}; + +const pauseFiberHookQueues = (fiber: Fiber): void => { + let currentHookState = fiber.memoizedState as HookState | null; + while (currentHookState) { + if (currentHookState.queue && typeof currentHookState.queue === 'object') { + pauseHookQueue(currentHookState.queue); + } + currentHookState = currentHookState?.next ?? null; + } +}; + +const resumeFiberHookQueues = (fiber: Fiber): void => { + let currentHookState = fiber.memoizedState as HookState | null; + while (currentHookState) { + if (currentHookState.queue && typeof currentHookState.queue === 'object') { + resumeHookQueue(currentHookState.queue); + } + currentHookState = currentHookState?.next ?? null; + } +}; + +const traverseAndPauseHookQueues = (fiber: Fiber | null): void => { + if (!fiber) return; + + if (isCompositeFiber(fiber)) { + pauseFiberHookQueues(fiber); + pauseFiberContextDependencies(fiber); + } + + traverseAndPauseHookQueues(fiber.child); + traverseAndPauseHookQueues(fiber.sibling); +}; + +const traverseAndResumeHookQueues = (fiber: Fiber | null): void => { + if (!fiber) return; + + if (isCompositeFiber(fiber)) { + resumeFiberHookQueues(fiber); + resumeFiberContextDependencies(fiber); + } + + traverseAndResumeHookQueues(fiber.child); + traverseAndResumeHookQueues(fiber.sibling); +}; + +const createPausedDispatcher = (originalDispatcher: Dispatcher): Dispatcher => { + return new Proxy(originalDispatcher, { + get(target, prop, receiver) { + const originalMethod = Reflect.get(target, prop, receiver); + + if (prop === 'useState') { + return (initialState: S | (() => S)): [S, (action: S | ((prev: S) => S)) => void] => { + return (originalMethod as Dispatcher['useState'])(initialState); + }; + } + + if (prop === 'useReducer') { + return ( + reducer: (state: S, action: A) => S, + initialArg: S, + init?: (arg: S) => S, + ): [S, (action: A) => void] => { + return (originalMethod as Dispatcher['useReducer'])(reducer, initialArg, init); + }; + } + + if (prop === 'useSyncExternalStore') { + return ( + subscribe: (onStoreChange: () => void) => () => void, + getSnapshot: () => T, + getServerSnapshot?: () => T, + ): T => { + const pauseAwareSubscribe = (onStoreChange: () => void) => { + const pauseAwareCallback = () => { + if (isUpdatesPaused) return; + onStoreChange(); + }; + return subscribe(pauseAwareCallback); + }; + return (originalMethod as Dispatcher['useSyncExternalStore'])(pauseAwareSubscribe, getSnapshot, getServerSnapshot); + }; + } + + if (prop === 'useTransition' && typeof originalMethod === 'function') { + return (...args: unknown[]): unknown => { + const result = originalMethod(...args); + + if (!Array.isArray(result) || result.length !== 2) { + return result; + } + + const [isPending, startTransition] = result as [boolean, (callback: () => void) => void]; + + if (typeof startTransition !== 'function') { + return result; + } + + const pausedStartTransition = (callback: () => void) => { + if (isUpdatesPaused) return; + startTransition(callback); + }; + + return [isPending, pausedStartTransition]; + }; + } + + if (prop === 'useDeferredValue' && typeof originalMethod === 'function') { + return (...args: unknown[]): T => { + return (originalMethod as (...args: unknown[]) => T)(...args); + }; + } + + return originalMethod; + }, + }); +}; + +const installDispatcherProxy = (renderer: ReactRenderer): void => { + const dispatcherRef = renderer.currentDispatcherRef as DispatcherRef | null; + if (!dispatcherRef || typeof dispatcherRef !== 'object') return; + if (pausedDispatcherStates.has(renderer)) return; + + const dispatcherKey: 'H' | 'current' = 'H' in dispatcherRef ? 'H' : 'current'; + const originalDescriptor = Object.getOwnPropertyDescriptor(dispatcherRef, dispatcherKey); + pausedDispatcherStates.set(renderer, { dispatcherKey, originalDescriptor }); + + let currentDispatcherValue = dispatcherRef[dispatcherKey]; + + Object.defineProperty(dispatcherRef, dispatcherKey, { + configurable: true, + enumerable: true, + get() { + if (isUpdatesPaused && currentDispatcherValue) { + return createPausedDispatcher(currentDispatcherValue as Dispatcher); + } + return currentDispatcherValue; + }, + set(newDispatcher) { + currentDispatcherValue = newDispatcher; + }, + }); +}; + +const uninstallDispatcherProxy = (renderer: ReactRenderer): void => { + const dispatcherPauseState = pausedDispatcherStates.get(renderer); + if (!dispatcherPauseState) return; + + const dispatcherRef = renderer.currentDispatcherRef as DispatcherRef | null; + if (!dispatcherRef) return; + + if (dispatcherPauseState.originalDescriptor) { + Object.defineProperty(dispatcherRef, dispatcherPauseState.dispatcherKey, dispatcherPauseState.originalDescriptor); + } else { + delete (dispatcherRef as Record)[dispatcherPauseState.dispatcherKey]; + } + + pausedDispatcherStates.delete(renderer); +}; + +/** + * Pauses all React state updates by patching hook queues, dispatchers, and context dependencies. + * When paused, all `useState`, `useReducer`, `useSyncExternalStore`, `useTransition`, `useDeferredValue`, + * and context updates become no-ops. + * + * This works by: + * 1. Intercepting the update queue's `pending` property to prevent updates from being enqueued + * 2. Wrapping `getSnapshot` for external stores to return the value captured at pause time + * 3. Freezing context dependencies to prevent context changes from triggering updates + * 4. Blocking `startTransition` calls when paused + * 5. Patching dispatchers so new components also get paused behavior + * + * @returns A function to resume normal React behavior. + * + * @example + * ```ts + * const resumeUpdates = pauseUpdates(); + * // All setState/dispatch calls are now no-ops + * resumeUpdates(); + * // React updates work normally again + * ``` + */ +export const pauseUpdates = (): (() => void) => { + if (isUpdatesPaused) { + return () => {}; + } + + const rdtHook = getRDTHook(); + + for (const renderer of rdtHook.renderers.values()) { + installDispatcherProxy(renderer); + } + + for (const fiberRoot of _fiberRoots) { + traverseAndPauseHookQueues(fiberRoot.current); + } + + isUpdatesPaused = true; + + return () => { + if (!isUpdatesPaused) return; + + isUpdatesPaused = false; + + for (const fiberRoot of _fiberRoots) { + traverseAndResumeHookQueues(fiberRoot.current); + } + + for (const renderer of rdtHook.renderers.values()) { + uninstallDispatcherProxy(renderer); + } + }; +}; + +/** + * Returns whether React updates are currently paused. + */ +export const areUpdatesPaused = (): boolean => { + return isUpdatesPaused; +}; + export interface InstrumentationOptions { name?: string; onActive?: () => unknown; diff --git a/packages/bippy/src/test/pause-updates.test.tsx b/packages/bippy/src/test/pause-updates.test.tsx new file mode 100644 index 00000000..222c3f11 --- /dev/null +++ b/packages/bippy/src/test/pause-updates.test.tsx @@ -0,0 +1,350 @@ +import '../index.js'; + +import * as React from 'react'; +import { describe, it, expect, beforeEach, afterEach } from 'vitest'; +import { render, screen, fireEvent, act, cleanup } from '@testing-library/react'; +import { pauseUpdates, areUpdatesPaused, instrument, secure, getRDTHook, _fiberRoots } from '../core.js'; + +const Counter = () => { + const [count, setCount] = React.useState(0); + return ( +
+ {count} + +
+ ); +}; + +const ReducerCounter = () => { + const [state, dispatch] = React.useReducer( + (currentState: { count: number }, action: { type: 'increment' | 'decrement' }) => { + switch (action.type) { + case 'increment': + return { count: currentState.count + 1 }; + case 'decrement': + return { count: currentState.count - 1 }; + default: + return currentState; + } + }, + { count: 0 }, + ); + + return ( +
+ {state.count} + +
+ ); +}; + +const createExternalStore = (initialValue: number) => { + let value = initialValue; + const listeners = new Set<() => void>(); + + return { + getSnapshot: () => value, + subscribe: (listener: () => void) => { + listeners.add(listener); + return () => listeners.delete(listener); + }, + increment: () => { + value++; + listeners.forEach((listener) => listener()); + }, + }; +}; + +const ExternalStoreCounter = ({ store }: { store: ReturnType }) => { + const count = React.useSyncExternalStore(store.subscribe, store.getSnapshot); + return ( +
+ {count} +
+ ); +}; + +describe('pauseUpdates', () => { + beforeEach(() => { + cleanup(); + instrument( + secure( + { + onCommitFiberRoot: () => {}, + }, + { dangerouslyRunInProduction: true }, + ), + ); + }); + + afterEach(() => { + cleanup(); + }); + + it('should return resume function', () => { + const resumeUpdates = pauseUpdates(); + expect(typeof resumeUpdates).toBe('function'); + resumeUpdates(); + }); + + it('should report areUpdatesPaused correctly', () => { + expect(areUpdatesPaused()).toBe(false); + const resumeUpdates = pauseUpdates(); + expect(areUpdatesPaused()).toBe(true); + resumeUpdates(); + expect(areUpdatesPaused()).toBe(false); + }); + + it('should handle multiple pauseUpdates calls', () => { + const resume1 = pauseUpdates(); + expect(areUpdatesPaused()).toBe(true); + const resume2 = pauseUpdates(); + expect(areUpdatesPaused()).toBe(true); + resume1(); + expect(areUpdatesPaused()).toBe(false); + resume2(); + expect(areUpdatesPaused()).toBe(false); + }); + + it('should pause useState updates', async () => { + render(); + + expect(_fiberRoots.size).toBeGreaterThan(0); + + const countElement = screen.getByTestId('count'); + const incrementButton = screen.getByTestId('increment'); + + expect(countElement.textContent).toBe('0'); + + await act(async () => { + fireEvent.click(incrementButton); + }); + expect(countElement.textContent).toBe('1'); + + const resumeUpdates = pauseUpdates(); + + await act(async () => { + fireEvent.click(incrementButton); + }); + expect(countElement.textContent).toBe('1'); + + await act(async () => { + fireEvent.click(incrementButton); + }); + expect(countElement.textContent).toBe('1'); + + resumeUpdates(); + + await act(async () => { + fireEvent.click(incrementButton); + }); + expect(countElement.textContent).toBe('2'); + }); + + it('should pause useReducer updates', async () => { + render(); + + expect(_fiberRoots.size).toBeGreaterThan(0); + + const countElement = screen.getByTestId('reducer-count'); + const incrementButton = screen.getByTestId('reducer-increment'); + + expect(countElement.textContent).toBe('0'); + + await act(async () => { + fireEvent.click(incrementButton); + }); + expect(countElement.textContent).toBe('1'); + + const resumeUpdates = pauseUpdates(); + + await act(async () => { + fireEvent.click(incrementButton); + }); + expect(countElement.textContent).toBe('1'); + + resumeUpdates(); + + await act(async () => { + fireEvent.click(incrementButton); + }); + expect(countElement.textContent).toBe('2'); + }); + + it('should pause useSyncExternalStore updates', async () => { + const store = createExternalStore(0); + render(); + + expect(_fiberRoots.size).toBeGreaterThan(0); + + const countElement = screen.getByTestId('external-count'); + + expect(countElement.textContent).toBe('0'); + + await act(async () => { + store.increment(); + }); + expect(countElement.textContent).toBe('1'); + + const resumeUpdates = pauseUpdates(); + + await act(async () => { + store.increment(); + }); + expect(countElement.textContent).toBe('1'); + + await act(async () => { + store.increment(); + }); + expect(countElement.textContent).toBe('1'); + + resumeUpdates(); + + await act(async () => { + store.increment(); + }); + expect(countElement.textContent).toBe('4'); + }); + + it('should work with renderers accessed via getRDTHook', () => { + const rdtHook = getRDTHook(); + expect(rdtHook.renderers).toBeDefined(); + }); + + it('should pause context updates', async () => { + const TestContext = React.createContext(0); + + const ContextConsumer = () => { + const value = React.useContext(TestContext); + return {value}; + }; + + const ContextProvider = () => { + const [contextValue, setContextValue] = React.useState(0); + return ( + + + + + ); + }; + + render(); + + const valueElement = screen.getByTestId('context-value'); + const updateButton = screen.getByTestId('update-context'); + + expect(valueElement.textContent).toBe('0'); + + await act(async () => { + fireEvent.click(updateButton); + }); + expect(valueElement.textContent).toBe('1'); + + const resumeUpdates = pauseUpdates(); + + await act(async () => { + fireEvent.click(updateButton); + }); + expect(valueElement.textContent).toBe('1'); + + resumeUpdates(); + + await act(async () => { + fireEvent.click(updateButton); + }); + expect(valueElement.textContent).toBe('2'); + }); + + it('should pause useTransition updates', async () => { + const TransitionComponent = () => { + const [count, setCount] = React.useState(0); + const [isPending, startTransition] = React.useTransition(); + + return ( +
+ {count} + {isPending ? 'pending' : 'idle'} + +
+ ); + }; + + render(); + + const countElement = screen.getByTestId('transition-count'); + const updateButton = screen.getByTestId('transition-update'); + + expect(countElement.textContent).toBe('0'); + + await act(async () => { + fireEvent.click(updateButton); + }); + expect(countElement.textContent).toBe('1'); + + const resumeUpdates = pauseUpdates(); + + await act(async () => { + fireEvent.click(updateButton); + }); + expect(countElement.textContent).toBe('1'); + + resumeUpdates(); + + await act(async () => { + fireEvent.click(updateButton); + }); + expect(countElement.textContent).toBe('2'); + }); + + it('should handle useDeferredValue gracefully', async () => { + const DeferredComponent = () => { + const [input, setInput] = React.useState('initial'); + const deferredInput = React.useDeferredValue(input); + + return ( +
+ {deferredInput} + +
+ ); + }; + + render(); + + const valueElement = screen.getByTestId('deferred-value'); + const updateButton = screen.getByTestId('update-input'); + + expect(valueElement.textContent).toBe('initial'); + + await act(async () => { + fireEvent.click(updateButton); + }); + expect(valueElement.textContent).toBe('updated'); + + const resumeUpdates = pauseUpdates(); + + await act(async () => { + fireEvent.click(screen.getByTestId('update-input')); + }); + + resumeUpdates(); + }); +});