refactor: introduce store selection bindings

src/core/__tests__/createShallowStore.test.ts

@@ -1,162 +1,48 @@
 import { act, renderHook } from "@testing-library/react";
-import { describe, expect, it } from "vitest";
+import { describe, expect, expectTypeOf, it, vi } from "vitest";
+import { subscribeWithSelector } from "zustand/middleware";
 import { createShallowStore } from "../createShallowStore";
 
 interface TestStore {
   count: number;
-  name: string;
   increment: () => void;
-  setName: (name: string) => void;
 }
 
 describe("createShallowStore", () => {
-  it("should create a store with shallow comparison", () => {
-    const { useStore } = createShallowStore<TestStore>((set) => ({
-      count: 0,
-      name: "test",
-      increment: () => set((state) => ({ count: state.count + 1 })),
-      setName: (name: string) => set({ name }),
-    }));
-
-    const { result } = renderHook(() => useStore());
-
-    expect(result.current.count).toBe(0);
-    expect(result.current.name).toBe("test");
-  });
-
-  it("should work with selector", () => {
-    const { useStore } = createShallowStore<TestStore>((set) => ({
-      count: 0,
-      name: "test",
-      increment: () => set((state) => ({ count: state.count + 1 })),
-      setName: (name: string) => set({ name }),
-    }));
-
-    const { result } = renderHook(() => useStore((state) => state.count));
-
-    expect(result.current).toBe(0);
-  });
-
-  it("should expose plain selector access", () => {
-    const { useStorePlain } = createShallowStore<TestStore>((set) => ({
-      count: 0,
-      name: "test",
-      increment: () => set((state) => ({ count: state.count + 1 })),
-      setName: (name: string) => set({ name }),
-    }));
-
-    const { result } = renderHook(() => useStorePlain((state) => state.name));
-
-    expect(result.current).toBe("test");
-  });
-
-  it("should update state correctly", () => {
+  it("should select and update the global store", () => {
     const { useStore, useStoreApi } = createShallowStore<TestStore>((set) => ({
       count: 0,
-      name: "test",
       increment: () => set((state) => ({ count: state.count + 1 })),
-      setName: (name: string) => set({ name }),
     }));
-
     const { result } = renderHook(() => useStore((state) => state.count));
 
-    act(() => {
-      useStoreApi.getState().increment();
-    });
+    act(() => useStoreApi.getState().increment());
 
     expect(result.current).toBe(1);
-
-    act(() => {
-      useStoreApi.getState().increment();
-    });
-
-    expect(result.current).toBe(2);
-  });
-
-  it("should provide access to store API", () => {
-    const { useStoreApi } = createShallowStore<TestStore>((set) => ({
-      count: 0,
-      name: "test",
-      increment: () => set((state) => ({ count: state.count + 1 })),
-      setName: (name: string) => set({ name }),
-    }));
-
-    const state = useStoreApi.getState();
-    expect(state.count).toBe(0);
-
-    act(() => {
-      state.increment();
-    });
-
     expect(useStoreApi.getState().count).toBe(1);
   });
 
-  it("should handle multiple selectors independently", () => {
-    const { useStore } = createShallowStore<TestStore>((set) => ({
-      count: 0,
-      name: "test",
-      increment: () => set((state) => ({ count: state.count + 1 })),
-      setName: (name: string) => set({ name }),
-    }));
-
-    const { result: countResult } = renderHook(() => useStore((state) => state.count));
-    const { result: nameResult } = renderHook(() => useStore((state) => state.name));
-
-    expect(countResult.current).toBe(0);
-    expect(nameResult.current).toBe("test");
-  });
-
-  it("should support custom equality function", () => {
-    const { useStore, useStoreApi } = createShallowStore<TestStore>((set) => ({
-      count: 0,
-      name: "test",
-      increment: () => set((state) => ({ count: state.count + 1 })),
-      setName: (name: string) => set({ name }),
-    }));
-
-    const { result } = renderHook(() =>
-      useStore(
-        (state) => state.count,
-        () => true
-      )
+  it("should preserve middleware-enhanced store capabilities", () => {
+    const listener = vi.fn();
+    const { useStoreApi } = createShallowStore<
+      TestStore,
+      [["zustand/subscribeWithSelector", never]]
+    >(
+      subscribeWithSelector((set) => ({
+        count: 0,
+        increment: () => set((state) => ({ count: state.count + 1 })),
+      }))
     );
 
-    expect(result.current).toBe(0);
-
-    act(() => {
-      useStoreApi.getState().increment();
-    });
-
-    // custom equality always returns true, so selection remains stable
-    expect(result.current).toBe(0);
-  });
-
-  it("should keep shallow mode and plain mode behavior distinct", () => {
-    const { useStore, useStorePlain, useStoreApi } = createShallowStore<TestStore>((set) => ({
-      count: 0,
-      name: "test",
-      increment: () => set((state) => ({ count: state.count + 1 })),
-      setName: (name: string) => set({ name }),
-    }));
-
-    let shallowRenders = 0;
-    let plainRenders = 0;
-
-    renderHook(() => {
-      shallowRenders += 1;
-      return useStore((state) => ({ count: state.count }));
-    });
-
-    renderHook(() => {
-      plainRenders += 1;
-      return useStorePlain();
-    });
+    expectTypeOf(useStoreApi.subscribe).toBeCallableWith(
+      (state: TestStore) => state.count,
+      listener
+    );
 
-    act(() => {
-      useStoreApi.getState().setName("updated");
-    });
+    const unsubscribe = useStoreApi.subscribe((state) => state.count, listener);
+    useStoreApi.getState().increment();
 
-    expect(shallowRenders).toBe(1);
-    expect(plainRenders).toBe(2);
+    expect(listener).toHaveBeenCalledWith(1, 0);
+    unsubscribe();
   });
 });

src/core/createShallowStore.ts

@@ -1,5 +1,5 @@
 import { createStore } from "zustand";
-import { useStoreSelection, useStoreSelectionPlain } from "../hooks";
+import { createStoreSelectionBindings } from "../hooks";
 import type { MutatorsStateCreator, StoreApiWithMutators, StoreMutatorTuple } from "../types";
 import type { ShallowStoreBindings } from "./createShallowStore.types";
 
@@ -139,27 +139,10 @@ export function createShallowStore<TState, TMutators extends Array<StoreMutatorT
     storeCreator
   );
 
-  function useShallowStore(): TState;
-  function useShallowStore<T>(selector: (state: TState) => T): T;
-  function useShallowStore<T>(
-    selector: (state: TState) => T,
-    equalityFn: (a: T, b: T) => boolean
-  ): T;
-  function useShallowStore<T>(
-    selector?: (state: TState) => T,
-    equalityFn?: (a: T | TState, b: T | TState) => boolean
-  ): T | TState {
-    return useStoreSelection(storeApi, selector, equalityFn);
-  }
-
-  function useStorePlain(): TState;
-  function useStorePlain<T>(selector: (state: TState) => T): T;
-  function useStorePlain<T>(selector?: (state: TState) => T): T | TState {
-    return useStoreSelectionPlain(storeApi, selector);
-  }
+  const { useStoreValue: useStore, useStorePlain } = createStoreSelectionBindings(() => storeApi);
 
   return {
-    useStore: useShallowStore,
+    useStore,
     useStorePlain,
     useStoreApi: storeApi,
   };

src/hooks/index.ts

@@ -7,5 +7,11 @@ export type {
   UseSelectorWithEqualityReturn,
   UseSelectorWithEqualitySelector,
 } from "./useSelectorWithEquality";
+export { createStoreSelectionBindings } from "./storeSelection";
 export { useStoreSelection, useStoreSelectionPlain } from "./storeSelection";
-export type { StorePlainHook, StoreValueHook } from "./storeSelection";
+export type {
+  StorePlainHook,
+  StoreResolver,
+  StoreSelectionBindings,
+  StoreValueHook,
+} from "./storeSelection";

src/hooks/storeSelection/__tests__/createStoreSelectionBindings.test.ts

@@ -0,0 +1,135 @@
+import { act, renderHook } from "@testing-library/react";
+import { describe, expect, it } from "vitest";
+import { createStore } from "zustand";
+import { createStoreSelectionBindings } from "../createStoreSelectionBindings";
+
+describe("createStoreSelectionBindings", () => {
+  it("should expose shallow and plain selection through one interface", () => {
+    const store = createStore(() => ({ count: 1, items: [1, 2] }));
+    const { useStoreValue, useStorePlain } = createStoreSelectionBindings(() => store);
+    const { result } = renderHook(() => ({
+      shallow: useStoreValue((state) => state.items),
+      plain: useStorePlain((state) => state.items),
+    }));
+    const initialItems = store.getState().items;
+    const replacementItems = [1, 2];
+
+    act(() => store.setState({ items: replacementItems }));
+
+    expect(result.current.shallow).toBe(initialItems);
+    expect(result.current.plain).toBe(replacementItems);
+  });
+
+  it("should retain shallow references across rerenders and unrelated updates", () => {
+    const store = createStore(() => ({ count: 1, label: "initial" }));
+    const { useStoreValue } = createStoreSelectionBindings(() => store);
+    const { result, rerender } = renderHook(() =>
+      useStoreValue((state) => ({ count: state.count }))
+    );
+    const initialSelection = result.current;
+
+    rerender();
+    act(() => store.setState({ label: "updated" }));
+
+    expect(result.current).toBe(initialSelection);
+
+    act(() => store.setState({ count: 2 }));
+
+    expect(result.current).toEqual({ count: 2 });
+    expect(result.current).not.toBe(initialSelection);
+  });
+
+  it("should use the latest selector and equality after rerender", () => {
+    const store = createStore(() => ({ count: 1 }));
+    const { useStoreValue } = createStoreSelectionBindings(() => store);
+    const { result, rerender } = renderHook(
+      ({ multiplier, equalityFn }) =>
+        useStoreValue((state) => state.count * multiplier, equalityFn),
+      { initialProps: { multiplier: 1, equalityFn: Object.is } }
+    );
+
+    rerender({ multiplier: 2, equalityFn: Object.is });
+    expect(result.current).toBe(2);
+
+    rerender({ multiplier: 2, equalityFn: () => true });
+    act(() => store.setState({ count: 3 }));
+    expect(result.current).toBe(2);
+
+    rerender({ multiplier: 2, equalityFn: Object.is });
+    expect(result.current).toBe(6);
+  });
+
+  it("should reset retained values and subscriptions when the store changes", () => {
+    const first = createStore(() => ({ items: [1, 2] }));
+    const second = createStore(() => ({ items: [1, 2] }));
+    let currentStore = first;
+    const { useStoreValue } = createStoreSelectionBindings(() => currentStore);
+    const { result, rerender } = renderHook(() =>
+      useStoreValue(
+        (state) => state.items,
+        () => true
+      )
+    );
+
+    expect(result.current).toBe(first.getState().items);
+
+    currentStore = second;
+    rerender();
+
+    expect(result.current).toBe(second.getState().items);
+
+    act(() => first.setState({ items: [9] }));
+    expect(result.current).toBe(second.getState().items);
+  });
+
+  it("should subscribe both selection modes to the current store", () => {
+    const first = createStore(() => ({ count: 1 }));
+    const second = createStore(() => ({ count: 1 }));
+    let currentStore = first;
+    const { useStoreValue, useStorePlain } = createStoreSelectionBindings(() => currentStore);
+    const { result, rerender } = renderHook(() => ({
+      shallow: useStoreValue((state) => state.count),
+      plain: useStorePlain((state) => state.count),
+    }));
+
+    currentStore = second;
+    rerender();
+    act(() => second.setState({ count: 4 }));
+
+    expect(result.current).toEqual({ shallow: 4, plain: 4 });
+
+    act(() => first.setState({ count: 9 }));
+
+    expect(result.current).toEqual({ shallow: 4, plain: 4 });
+  });
+
+  it("should return full state when selectors are omitted", () => {
+    const store = createStore(() => ({ count: 1 }));
+    const { useStoreValue, useStorePlain } = createStoreSelectionBindings(() => store);
+    const { result } = renderHook(() => ({
+      shallow: useStoreValue(),
+      plain: useStorePlain(),
+    }));
+
+    expect(result.current.shallow).toBe(store.getState());
+    expect(result.current.plain).toBe(store.getState());
+
+    act(() => store.setState({ count: 2 }));
+
+    expect(result.current.shallow).toBe(store.getState());
+    expect(result.current.plain).toBe(store.getState());
+  });
+
+  it("should retain independent values for separate consumers", () => {
+    const store = createStore(() => ({ count: 1, label: "initial" }));
+    const { useStoreValue } = createStoreSelectionBindings(() => store);
+    const { result } = renderHook(() => ({
+      count: useStoreValue((state) => state.count),
+      label: useStoreValue((state) => state.label),
+    }));
+
+    act(() => store.setState({ count: 2 }));
+
+    expect(result.current).toEqual({ count: 2, label: "initial" });
+  });
+});