feat: Support setState callback in useControlledState (#9041)

* feat: Support function setState callback in useControlledState

* fixes

* fix docs build

---------

Co-authored-by: Robert Snow <[email protected]>

Author: devongovett

packages/@react-stately/calendar/docs/useRangeCalendarState.mdx

@@ -33,7 +33,7 @@ keywords: [date, calendar, state]
 
 ## Interface
 
-<ClassAPI links={docs.links} class={docs.links[docs.exports.useRangeCalendarState.return.id]} />
+<ClassAPI links={docs.links} class={docs.links[docs.exports.useRangeCalendarState.return.base.id]} />
 
 ## Example
 

packages/@react-stately/calendar/src/types.ts

@@ -106,11 +106,11 @@ export interface CalendarState extends CalendarStateBase {
   setValue(value: CalendarDate | null): void
 }
 
-export interface RangeCalendarState extends CalendarStateBase {
+export interface RangeCalendarState<T extends DateValue = DateValue> extends CalendarStateBase {
   /** The currently selected date range. */
-  readonly value: RangeValue<DateValue> | null,
+  readonly value: RangeValue<T> | null,
   /** Sets the currently selected date range. */
-  setValue(value: RangeValue<DateValue> | null): void,
+  setValue(value: RangeValue<T> | null): void,
   /** Highlights the given date during selection, e.g. by hovering or dragging. */
   highlightDate(date: CalendarDate): void,
   /** The current anchor date that the user clicked on to begin range selection. */

packages/@react-stately/calendar/src/useRangeCalendarState.ts

@@ -45,7 +45,7 @@ export interface RangeCalendarStateOptions<T extends DateValue = DateValue> exte
  * Provides state management for a range calendar component.
  * A range calendar displays one or more date grids and allows users to select a contiguous range of dates.
  */
-export function useRangeCalendarState<T extends DateValue = DateValue>(props: RangeCalendarStateOptions<T>): RangeCalendarState {
+export function useRangeCalendarState<T extends DateValue = DateValue>(props: RangeCalendarStateOptions<T>): RangeCalendarState<T> {
   let {
     value: valueProp,
     defaultValue,

packages/@react-stately/checkbox/src/useCheckboxGroupState.ts

@@ -98,10 +98,12 @@ export function useCheckboxGroupState(props: CheckboxGroupProps = {}): CheckboxG
       if (props.isReadOnly || props.isDisabled) {
         return;
       }
-      if (!selectedValues.includes(value)) {
-        selectedValues = selectedValues.concat(value);
-        setValue(selectedValues);
-      }
+      setValue(selectedValues => {
+        if (!selectedValues.includes(value)) {
+          return selectedValues.concat(value);
+        }
+        return selectedValues;
+      });
     },
     removeValue(value) {
       if (props.isReadOnly || props.isDisabled) {

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

@@ -10,12 +10,20 @@
  * governing permissions and limitations under the License.
  */
 
-import {useCallback, useEffect, useRef, useState} from 'react';
+import React, {SetStateAction, useCallback, useEffect, useRef, useState} from 'react';
 
-export function useControlledState<T, C = T>(value: Exclude<T, undefined>, defaultValue: Exclude<T, undefined> | undefined, onChange?: (v: C, ...args: any[]) => void): [T, (value: T, ...args: any[]) => void];
-export function useControlledState<T, C = T>(value: Exclude<T, undefined> | undefined, defaultValue: Exclude<T, undefined>, onChange?: (v: C, ...args: any[]) => void): [T, (value: T, ...args: any[]) => void];
-export function useControlledState<T, C = T>(value: T, defaultValue: T, onChange?: (v: C, ...args: any[]) => void): [T, (value: T, ...args: any[]) => void] {
+// Use the earliest effect possible to reset the ref below.
+const useEarlyEffect: typeof React.useLayoutEffect = typeof document !== 'undefined'
+  ? React['useInsertionEffect'] ?? React.useLayoutEffect
+  : () => {};
+
+export function useControlledState<T, C = T>(value: Exclude<T, undefined>, defaultValue: Exclude<T, undefined> | undefined, onChange?: (v: C, ...args: any[]) => void): [T, (value: SetStateAction<T>, ...args: any[]) => void];
+export function useControlledState<T, C = T>(value: Exclude<T, undefined> | undefined, defaultValue: Exclude<T, undefined>, onChange?: (v: C, ...args: any[]) => void): [T, (value: SetStateAction<T>, ...args: any[]) => void];
+export function useControlledState<T, C = T>(value: T, defaultValue: T, onChange?: (v: C, ...args: any[]) => void): [T, (value: SetStateAction<T>, ...args: any[]) => void] {
+  // Store the value in both state and a ref. The state value will only be used when uncontrolled.
+  // The ref is used to track the most current value, which is passed to the function setState callback.
   let [stateValue, setStateValue] = useState(value || defaultValue);
+  let valueRef = useRef(stateValue);
 
   let isControlledRef = useRef(value !== undefined);
   let isControlled = value !== undefined;
@@ -27,49 +35,29 @@ export function useControlledState<T, C = T>(value: T, defaultValue: T, onChange
     isControlledRef.current = isControlled;
   }, [isControlled]);
 
+  // After each render, update the ref to the current value.
+  // This ensures that the setState callback argument is reset.
+  // Note: the effect should not have any dependencies so that controlled values always reset.
   let currentValue = isControlled ? value : stateValue;
-  let setValue = useCallback((value, ...args) => {
-    let onChangeCaller = (value, ...onChangeArgs) => {
-      if (onChange) {
-        if (!Object.is(currentValue, value)) {
-          onChange(value, ...onChangeArgs);
-        }
-      }
-      if (!isControlled) {
-        // If uncontrolled, mutate the currentValue local variable so that
-        // calling setState multiple times with the same value only emits onChange once.
-        // We do not use a ref for this because we specifically _do_ want the value to
-        // reset every render, and assigning to a ref in render breaks aborted suspended renders.
-        // eslint-disable-next-line react-hooks/exhaustive-deps
-        currentValue = value;
-      }
-    };
+  useEarlyEffect(() => {
+    valueRef.current = currentValue;
+  });
+
+  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);
 
-    if (typeof value === 'function') {
-      if (process.env.NODE_ENV !== 'production') {
-        console.warn('We can not support a function callback. See Github Issues for details https://github.com/adobe/react-spectrum/issues/2320');
-      }
-      // this supports functional updates https://reactjs.org/docs/hooks-reference.html#functional-updates
-      // when someone using useControlledState calls setControlledState(myFunc)
-      // this will call our useState setState with a function as well which invokes myFunc and calls onChange with the value from myFunc
-      // if we're in an uncontrolled state, then we also return the value of myFunc which to setState looks as though it was just called with myFunc from the beginning
-      // otherwise we just return the controlled value, which won't cause a rerender because React knows to bail out when the value is the same
-      let updateFunction = (oldValue, ...functionArgs) => {
-        let interceptedValue = value(isControlled ? currentValue : oldValue, ...functionArgs);
-        onChangeCaller(interceptedValue, ...args);
-        if (!isControlled) {
-          return interceptedValue;
-        }
-        return oldValue;
-      };
-      setStateValue(updateFunction);
-    } else {
-      if (!isControlled) {
-        setStateValue(value);
-      }
-      onChangeCaller(value, ...args);
+      // 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);
     }
-  }, [isControlled, currentValue, onChange]);
+  }, [onChange]);
 
   return [currentValue, setValue];
 }