fix: apply batched functional updates in useControllableState (#80)

Author: TheGreatAxios

src/hooks/use-controllable-state.test.ts

@@ -1,11 +1,12 @@
 import { afterEach, beforeEach, describe, expect, mock, test } from "bun:test";
-import { act, createElement } from "react";
+import { StrictMode, act, createElement } from "react";
 import { createRoot } from "react-dom/client";
 
 import {
   useControllableState,
   type UseControllableStateOptions,
 } from "./use-controllable-state.js";
+import { renderHook } from "../test/render-hook.js";
 
 function mount<T>(options: UseControllableStateOptions<T>) {
   const container = document.createElement("div");
@@ -52,6 +53,45 @@ describe("useControllableState", () => {
     handle.unmount();
   });
 
+  test("uncontrolled: batched functional updates each see the previous result", () => {
+    const changes: number[] = [];
+    const { result, unmount } = renderHook(() =>
+      useControllableState({
+        defaultValue: 0,
+        name: "test",
+        onChange: (v) => changes.push(v),
+      }),
+    );
+    act(() => {
+      result.current[1]((previous) => previous + 1);
+      result.current[1]((previous) => previous + 1);
+    });
+    expect(result.current[0]).toBe(2);
+    expect(changes).toEqual([1, 2]);
+    unmount();
+  });
+
+  test("uncontrolled: onChange fires once per update under StrictMode", () => {
+    const changes: number[] = [];
+    const container = document.createElement("div");
+    const root = createRoot(container);
+    let set: ((action: (previous: number) => number) => void) | undefined;
+    function Host() {
+      set = useControllableState({
+        defaultValue: 0,
+        name: "test",
+        onChange: (v) => changes.push(v),
+      })[1];
+      return null;
+    }
+    act(() =>
+      root.render(createElement(StrictMode, null, createElement(Host))),
+    );
+    act(() => set?.((previous) => previous + 1));
+    expect(changes).toEqual([1]);
+    act(() => root.unmount());
+  });
+
   test("controlled: the value tracks the prop, not internal state", () => {
     let changed: boolean | undefined;
     const handle = mount({

src/hooks/use-controllable-state.ts

@@ -52,16 +52,27 @@ export function useControllableState<T>(
   }
   wasControlledRef.current = isControlled;
 
+  // Mirrors the uncontrolled state synchronously so batched functional updates
+  // chain off each other. Resolving inside a React updater instead would call
+  // `onChange` from that updater, which StrictMode invokes twice.
+  const latestUncontrolledRef = useRef(uncontrolledValue);
+
   const setValue = useCallback(
     (action: ControllableStateAction<T>) => {
-      const resolved =
+      const previous = isControlled
+        ? controlledValue
+        : latestUncontrolledRef.current;
+      const next =
         typeof action === "function"
-          ? (action as (previous: T) => T)(value)
+          ? (action as (previous: T) => T)(previous)
           : action;
-      if (!isControlled) setUncontrolledValue(resolved);
-      onChange?.(resolved);
+      if (!isControlled) {
+        latestUncontrolledRef.current = next;
+        setUncontrolledValue(next);
+      }
+      onChange?.(next);
     },
-    [isControlled, value, onChange],
+    [isControlled, controlledValue, onChange],
   );
 
   return [value, setValue] as const;

src/hooks/use-dismissable-popover.test.ts

@@ -73,6 +73,16 @@ describe("useDismissablePopover", () => {
     handle.unmount();
   });
 
+  test("uncontrolled: two batched toggles return to closed", () => {
+    const handle = mount();
+    act(() => {
+      handle.get().setOpen((previous) => !previous);
+      handle.get().setOpen((previous) => !previous);
+    });
+    expect(handle.get().open).toBe(false);
+    handle.unmount();
+  });
+
   test("uncontrolled: honors defaultOpen", () => {
     const handle = mount({ defaultOpen: true });
     expect(handle.get().open).toBe(true);