fix: refactor useControlledState to address state-outdated issues (#9331)

* fix: refactor useControlledState to address state-outdated issues

* Reduce fix to only required changes

---------

Co-authored-by: Devon Govett <[email protected]>

Author: lixiaoyan

packages/@react-stately/utils/src/useControlledState.ts

@@ -10,7 +10,7 @@
  * governing permissions and limitations under the License.
  */
 
-import React, {SetStateAction, useCallback, useEffect, useRef, useState} from 'react';
+import React, {SetStateAction, useCallback, useEffect, useReducer, useRef, useState} from 'react';
 
 // Use the earliest effect possible to reset the ref below.
 const useEarlyEffect: typeof React.useLayoutEffect = typeof document !== 'undefined'
@@ -43,16 +43,19 @@ export function useControlledState<T, C = T>(value: T, defaultValue: T, onChange
     valueRef.current = currentValue;
   });
 
+  let [, forceUpdate] = useReducer(() => ({}), {});
   let setValue = useCallback((value: SetStateAction<T>, ...args: any[]) => {
     // @ts-ignore - TS doesn't know that T cannot be a function.
     let newValue = typeof value === 'function' ? value(valueRef.current) : value;
     if (!Object.is(valueRef.current, newValue)) {
       // Update the ref so that the next setState callback has the most recent value.
       valueRef.current = newValue;
 
-      // Always trigger a setState, even when controlled, so that the layout effect above runs to reset the value.
       setStateValue(newValue);
 
+      // Always trigger a re-render, even when controlled, so that the layout effect above runs to reset the value.
+      forceUpdate();
+
       // Trigger onChange. Note that if setState is called multiple times in a single event,
       // onChange will be called for each one instead of only once.
       onChange?.(newValue, ...args);

packages/@react-stately/utils/test/useControlledState.test.tsx

@@ -212,6 +212,16 @@ describe('useControlledState tests', function () {
 
     onChangeSpy.mockClear();
 
+    act(() => setValue((prevValue) => {
+      expect(prevValue).toBe('updated');
+      return 'newValue';
+    }));
+    [value, setValue] = result.current;
+    expect(value).toBe('updated');
+    expect(onChangeSpy).toHaveBeenLastCalledWith('newValue');
+
+    onChangeSpy.mockClear();
+
     act(() => setValue((prevValue) => {
       expect(prevValue).toBe('updated');
       return 'updated';