diff --git a/packages/react/core/src/__tests__/useWeakRef-test.tsx b/packages/react/core/src/__tests__/useWeakRef-test.tsx new file mode 100644 index 00000000..4f55725f --- /dev/null +++ b/packages/react/core/src/__tests__/useWeakRef-test.tsx @@ -0,0 +1,73 @@ +import React from 'react'; +import { act, create } from 'react-test-renderer'; + +import { renderHook } from '../test-renderer.js'; +import { useWeakRef } from '../useWeakRef.js'; + +function TestComponent({ + keyObject, + thiefRef, +}: { + keyObject: WeakKey; + thiefRef?: React.MutableRefObject | undefined>; +}) { + const weakRef = useWeakRef(keyObject); + if (thiefRef) { + thiefRef.current = weakRef; + } + return <>{weakRef.current}; +} + +describe('useWeakRef', () => { + it('rerenders every component in which it is used when the value of `current` is mutated', async () => { + expect.assertions(2); + const keyObject = {}; + const thiefRef: React.MutableRefObject | undefined> = { current: undefined }; // We use this to reach in and grab the ref produced. + let testRenderer; + await act(() => { + testRenderer = create( + <> + + {'/'} + + + ); + }); + expect(testRenderer).toMatchInlineSnapshot(`"/"`); + await act(() => { + const weakRefFromRendererA = thiefRef.current!; + // Merely mutating this ref should cause both values to update. + weakRefFromRendererA.current = 'updated value'; + }); + expect(testRenderer).toMatchInlineSnapshot(` + [ + "updated value", + "/", + "updated value", + ] + `); + }); + it('vends the latest value of the ref to new components that mount', async () => { + expect.assertions(1); + // STEP 1: Set a value on the ref. + const keyObject = {}; + let result: ReturnType['result']; + await act(() => { + ({ result } = renderHook(() => useWeakRef(keyObject))); + if (result.__type === 'error') { + fail('`useWeakRef()` threw an unexpected exception'); + } + }); + await act(() => { + const weakRefFromRenderer = result.current! as React.MutableRefObject; + // Mutating this value should cause a re-render of the hook from above. + weakRefFromRenderer.current = 'initial value'; + }); + // STEP 2: Ensure that a brand new component receives that value on mount + let testRenderer; + await act(() => { + testRenderer = create(); + }); + expect(testRenderer).toMatchInlineSnapshot(`"initial value"`); + }); +}); diff --git a/packages/react/core/src/useWeakRef.ts b/packages/react/core/src/useWeakRef.ts new file mode 100644 index 00000000..3e3de8cc --- /dev/null +++ b/packages/react/core/src/useWeakRef.ts @@ -0,0 +1,71 @@ +import type React from 'react'; +import { useCallback, useRef, useSyncExternalStore } from 'react'; + +const values = new WeakMap(); +const subscribers = new Map void>>(); + +function getServerSnapshot() { + return { current: undefined }; +} + +function createMutableRef(keyObj: WeakKey): React.MutableRefObject { + return { + get current(): TValue | undefined { + return values.get(keyObj) as TValue | undefined; + }, + set current(newValue: TValue | undefined) { + if (newValue === values.get(keyObj)) { + return; + } + if (newValue === undefined) { + values.delete(keyObj); + } else { + values.set(keyObj, newValue); + } + subscribers.get(keyObj)?.forEach((cb) => cb()); + }, + }; +} + +/** + * Given an object as a key, this hook will vend you a mutable ref that causes all of its consumers + * to rerender whenever the `current` property is mutated. At all times, the value of `current` + * points to the latest value. + * + * This is particularly useful when you need to share a `Promise` across your entire application and + * you need every consumer to have access to the latest value of it synchronously, without having to + * wait for a rerender. Use this hook to create a ref related to some stable object, then store your + * promise in that ref. + * + * Note that the value related to the key object will persist even when every component that uses + * this hook unmounts. The next component to mount and call this hook with the same `keyObj` will + * recieve the same value. The value will only be released when `keyObj` is garbage collected. + */ +export function useWeakRef(keyObj: WeakKey): React.MutableRefObject { + const subscribe = useCallback( + (onStoreChange: () => void) => { + let callbacks = subscribers.get(keyObj); + if (callbacks == null) { + subscribers.set(keyObj, (callbacks = new Set())); + } + function handleChange() { + // This is super subtle; When there's a change, we want to create a new ref object + // so that `useSyncExternalStore()` perceives it as changed and triggers a rerender. + ref.current = createMutableRef(keyObj); + onStoreChange(); + } + callbacks.add(handleChange); + return () => { + ref.current = undefined; + callbacks.delete(handleChange); + if (callbacks.size === 0) { + subscribers.delete(keyObj); + } + }; + }, + [keyObj] + ); + const ref = useRef>(); + const getSnapshot = useCallback(() => (ref.current ||= createMutableRef(keyObj)), [keyObj]); + return useSyncExternalStore(subscribe, getSnapshot, getServerSnapshot); +}