[LG-5607] feat: useControlled - accept useState Dispatch (#3272)

Author: adamrasheed

.changeset/proud-clouds-hide.md

@@ -0,0 +1,5 @@
+---
+'@leafygreen-ui/hooks': minor
+---
+
+Updated the useControlled hook argument to utilize the useState Dispatch while ensuring backwards compatibility

packages/hooks/src/useControlled/useControlled.spec.tsx

@@ -120,6 +120,21 @@ describe('packages/hooks/useControlled', () => {
       });
       expect(result.current.value).toBe('apple');
     });
+
+    test('updateValue accepts a function for functional updates', () => {
+      const handler = jest.fn();
+      const { result } = renderUseControlledHook<number>(5, handler);
+      result.current.updateValue(prev => prev + 1);
+      expect(handler).toHaveBeenCalledWith(6);
+    });
+
+    test('functional update uses current controlled value', () => {
+      const handler = jest.fn();
+      const { result, rerender } = renderUseControlledHook<number>(10, handler);
+      rerender(20);
+      result.current.updateValue(prev => prev * 2);
+      expect(handler).toHaveBeenCalledWith(40);
+    });
   });
 
   describe('Uncontrolled', () => {
@@ -170,6 +185,49 @@ describe('packages/hooks/useControlled', () => {
       expect(result.current.value).toBe('apple');
       expect(handler).not.toHaveBeenCalled();
     });
+
+    test('updateValue accepts a function for functional updates', () => {
+      const handler = jest.fn();
+      const { result } = renderUseControlledHook<number>(undefined, handler, 5);
+      act(() => {
+        result.current.updateValue(prev => prev + 1);
+      });
+      expect(handler).toHaveBeenCalledWith(6);
+      expect(result.current.value).toBe(6);
+    });
+
+    test('functional update uses current uncontrolled value', () => {
+      const handler = jest.fn();
+      const { result } = renderUseControlledHook<number>(
+        undefined,
+        handler,
+        10,
+      );
+      act(() => {
+        result.current.updateValue(prev => prev * 2);
+      });
+      expect(handler).toHaveBeenCalledWith(20);
+      expect(result.current.value).toBe(20);
+    });
+
+    test('functional update avoids stale closures', () => {
+      const handler = jest.fn();
+      const { result } = renderUseControlledHook<number>(undefined, handler, 0);
+
+      // Simulate multiple updates that might capture stale values
+      act(() => {
+        result.current.updateValue(prev => prev + 1);
+        result.current.updateValue(prev => prev + 1);
+        result.current.updateValue(prev => prev + 1);
+      });
+
+      // With functional updates, each update should use the latest value
+      expect(handler).toHaveBeenCalledTimes(3);
+      expect(handler).toHaveBeenNthCalledWith(1, 1);
+      expect(handler).toHaveBeenNthCalledWith(2, 2);
+      expect(handler).toHaveBeenNthCalledWith(3, 3);
+      expect(result.current.value).toBe(3);
+    });
   });
 
   describe('Within test component', () => {

packages/hooks/src/useControlled/useControlled.ts

@@ -48,12 +48,30 @@ export const useControlled = <T extends any = undefined>(
    * If the component is controlled, it will not update the controlled value.
    *
    * onChange callback is called if provided.
+   * Accepts either a direct value or a function that receives the previous value.
    */
-  const updateValue = (newVal: T) => {
+  const updateValue = (newVal: React.SetStateAction<T>) => {
     if (!isControlled) {
-      setUncontrolledValue(newVal);
+      // In uncontrolled mode, use the state setter's functional update
+      // to ensure we always get the latest value
+      if (typeof newVal === 'function') {
+        setUncontrolledValue(prev => {
+          const resolvedValue = (newVal as (prev: T) => T)(prev);
+          onChange?.(resolvedValue);
+          return resolvedValue;
+        });
+      } else {
+        setUncontrolledValue(newVal);
+        onChange?.(newVal);
+      }
+    } else {
+      // In controlled mode, resolve the value using the current controlled value
+      const resolvedValue =
+        typeof newVal === 'function'
+          ? (newVal as (prev: T) => T)(value)
+          : newVal;
+      onChange?.(resolvedValue);
     }
-    onChange?.(newVal);
   };
 
   /**

packages/hooks/src/useControlled/useControlled.types.ts

@@ -7,9 +7,10 @@ export interface ControlledReturnObject<T extends any> {
 
   /**
    * Either updates the uncontrolled value,
-   * or calls the provided `onChange` callback
+   * or calls the provided `onChange` callback.
+   * Accepts either a direct value or a function that receives the previous value.
    */
-  updateValue: (newVal: T) => void;
+  updateValue: (newVal: React.SetStateAction<T>) => void;
 
   /**
    * A setter for the internal value.