diff --git a/packages/core/src/hooks/useControlledState/useControlledState.spec.tsx b/packages/core/src/hooks/useControlledState/useControlledState.spec.tsx index afdf9138..f9867c91 100644 --- a/packages/core/src/hooks/useControlledState/useControlledState.spec.tsx +++ b/packages/core/src/hooks/useControlledState/useControlledState.spec.tsx @@ -115,4 +115,19 @@ describe('useControlledState', () => { const [nextValue] = result.current; expect(nextValue).toBe(8); }); + + it('should handle multiple function setState actions', async () => { + const { result } = renderHookSSR(() => useControlledState({ defaultValue: 5 })); + const [value] = result.current; + expect(value).toBe(5); + + await act(async () => { + const [, setValue] = result.current; + setValue(prev => prev + 3); + setValue(prev => prev + 3); + }); + + const [nextValue] = result.current; + expect(nextValue).toBe(11); + }); }); diff --git a/packages/core/src/hooks/useControlledState/useControlledState.ts b/packages/core/src/hooks/useControlledState/useControlledState.ts index 43c5ce92..b6eee4ad 100644 --- a/packages/core/src/hooks/useControlledState/useControlledState.ts +++ b/packages/core/src/hooks/useControlledState/useControlledState.ts @@ -1,4 +1,8 @@ -import { type Dispatch, type SetStateAction, useCallback, useState } from 'react'; +import React, { type Dispatch, type SetStateAction, useCallback, useReducer, useRef, useState } from 'react'; + +const useSafeInsertionEffect: typeof React.useLayoutEffect = + /* c8 ignore next */ + typeof document !== 'undefined' ? (React['useInsertionEffect'] ?? React.useLayoutEffect) : () => {}; type ControlledState = { value: T; defaultValue?: never } | { defaultValue: T; value?: T }; @@ -49,19 +53,35 @@ export function useControlledState({ equalityFn = Object.is, }: UseControlledStateProps): [T, Dispatch>] { const [uncontrolledState, setUncontrolledState] = useState(defaultValue as T); + const valueRef = useRef(uncontrolledState); + const controlled = valueProp !== undefined; const value = controlled ? valueProp : uncontrolledState; + // After each render, reset valueRef to the actual value. + // This keeps the ref in sync with the actual value after optimistic updates. + const [, triggerRerender] = useReducer(() => ({}), {}); + useSafeInsertionEffect(() => { + valueRef.current = value; + }); + const setValue = useCallback( (next: SetStateAction) => { - const nextValue = isSetStateAction(next) ? next(value) : next; + const nextValue = isSetStateAction(next) ? next(valueRef.current) : next; + + if (equalityFn(valueRef.current, nextValue) === true) return; + + valueRef.current = nextValue; + // In controlled mode the parent may decide not to update the value + // (e.g. conditional logic inside onChange). + // Force a re-render so the effect above can resync the ref with the actual value. + triggerRerender(); - if (equalityFn(value, nextValue) === true) return; if (controlled === false) setUncontrolledState(nextValue); if (controlled === true && nextValue === undefined) setUncontrolledState(nextValue); onChange?.(nextValue); }, - [controlled, onChange, equalityFn, value] + [controlled, onChange, equalityFn] ); return [value, setValue];