refactor: centralize store selection policy

src/core/__tests__/createShallowStore.test.ts

@@ -159,20 +159,4 @@ describe("createShallowStore", () => {
     expect(shallowRenders).toBe(1);
     expect(plainRenders).toBe(2);
   });
-
-  it("should preserve shallow selection reference across parent rerenders", () => {
-    const { useStore } = createShallowStore<TestStore>((set) => ({
-      count: 0,
-      name: "test",
-      increment: () => set((state) => ({ count: state.count + 1 })),
-      setName: (name: string) => set({ name }),
-    }));
-
-    const { result, rerender } = renderHook(() => useStore((state) => ({ count: state.count })));
-    const initialSelection = result.current;
-
-    rerender();
-
-    expect(result.current).toBe(initialSelection);
-  });
 });

src/core/__tests__/createStoreToolkit.test.tsx

@@ -206,21 +206,6 @@ describe("createStoreToolkit", () => {
     expect(typeof plainResult.current.increment).toBe("function");
   });
 
-  it("should preserve resolved shallow selection reference across parent rerenders", () => {
-    const toolkit = createStoreToolkit<TestStore>((set) => ({
-      count: 0,
-      increment: () => set((state) => ({ count: state.count + 1 })),
-    }));
-    const { result, rerender } = renderHook(() =>
-      toolkit.useResolvedValue((state) => ({ count: state.count }))
-    );
-    const initialSelection = result.current;
-
-    rerender();
-
-    expect(result.current).toBe(initialSelection);
-  });
-
   it("should reset resolved selection reference when the resolved store changes", () => {
     interface CollectionStore {
       items: Array<number>;

src/core/createShallowStore.ts

@@ -1,7 +1,5 @@
-import { createStore, useStore } from "zustand";
-import { shallow } from "zustand/shallow";
-import { useSelectorWithEquality } from "../hooks";
-import { identitySelector } from "../utils";
+import { createStore } from "zustand";
+import { useStoreSelection, useStoreSelectionPlain } from "../hooks";
 import type {
   MutatorsStateCreator,
   ShallowStoreBindings,
@@ -155,20 +153,13 @@ export function createShallowStore<TState, TMutators extends Array<StoreMutatorT
     selector?: (state: TState) => T,
     equalityFn?: (a: T | TState, b: T | TState) => boolean
   ): T | TState {
-    const actualSelector = useSelectorWithEquality({
-      cacheKey: storeApi,
-      equalityFn: equalityFn ?? shallow,
-      selector: (state: TState): T | TState => (selector ? selector(state) : state),
-    });
-
-    return useStore(storeApi, actualSelector);
+    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 {
-    const actualSelector = selector ?? identitySelector<TState>;
-    return useStore<typeof storeApi, T | TState>(storeApi, actualSelector);
+    return useStoreSelectionPlain(storeApi, selector);
   }
 
   return {

src/hooks/createResolvedStoreHooks.ts

@@ -1,7 +1,4 @@
-import { useStore } from "zustand";
-import { shallow } from "zustand/shallow";
-import { identitySelector } from "../utils";
-import { useSelectorWithEquality } from "./useSelectorWithEquality";
+import { useStoreSelection, useStoreSelectionPlain } from "./storeSelection";
 import type { StoreApiWithMutators, StoreMutatorTuple } from "../types";
 
 /**
@@ -50,21 +47,14 @@ export function createResolvedStoreHooks<TState, TMutators extends Array<StoreMu
     equalityFn?: (a: T | TState, b: T | TState) => boolean
   ): T | TState {
     const store = useResolvedStoreApi();
-    const actualSelector = useSelectorWithEquality({
-      cacheKey: store,
-      equalityFn: equalityFn ?? shallow,
-      selector: (state: TState): T | TState => (selector ? selector(state) : state),
-    });
-
-    return useStore(store, actualSelector);
+    return useStoreSelection(store, selector, equalityFn);
   }
 
   function useResolvedStorePlain(): TState;
   function useResolvedStorePlain<T>(selector: (state: TState) => T): T;
   function useResolvedStorePlain<T>(selector?: (state: TState) => T): T | TState {
     const store = useResolvedStoreApi();
-    const actualSelector = selector ?? identitySelector<TState>;
-    return useStore<typeof store, T | TState>(store, actualSelector);
+    return useStoreSelectionPlain(store, selector);
   }
 
   return {

src/hooks/index.ts

@@ -6,3 +6,4 @@ export type {
   UseSelectorWithEqualityReturn,
   UseSelectorWithEqualitySelector,
 } from "./useSelectorWithEquality";
+export { useStoreSelection, useStoreSelectionPlain } from "./storeSelection";

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

@@ -0,0 +1,126 @@
+import { act, renderHook } from "@testing-library/react";
+import { describe, expect, it } from "vitest";
+import { createStore } from "zustand";
+import { useStoreSelection, useStoreSelectionPlain } from "../useStoreSelection";
+
+function createTestStore() {
+  return createStore(() => ({ count: 1, label: "initial", items: [1, 2] }));
+}
+
+describe("store selection", () => {
+  it("should retain shallow references across rerenders and unrelated updates", () => {
+    const store = createTestStore();
+    const { result, rerender } = renderHook(() =>
+      useStoreSelection(store, (state) => ({ count: state.count }))
+    );
+    const initialSelection = result.current;
+
+    rerender();
+    expect(result.current).toBe(initialSelection);
+
+    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 = createTestStore();
+    const { result, rerender } = renderHook(
+      ({ multiplier, equalityFn }) =>
+        useStoreSelection(store, (state) => state.count * multiplier, equalityFn),
+      { initialProps: { multiplier: 1, equalityFn: Object.is } }
+    );
+    expect(result.current).toBe(1);
+
+    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 = createTestStore();
+    const second = createTestStore();
+    const { result, rerender } = renderHook(
+      ({ store }) =>
+        useStoreSelection(
+          store,
+          (state) => state.items,
+          () => true
+        ),
+      { initialProps: { store: first } }
+    );
+    expect(result.current).toBe(first.getState().items);
+
+    rerender({ store: second });
+    expect(result.current).toBe(second.getState().items);
+
+    act(() => first.setState({ items: [9] }));
+    expect(result.current).toBe(second.getState().items);
+  });
+
+  it("should subscribe to the replacement store in both selection modes", () => {
+    const first = createTestStore();
+    const second = createTestStore();
+    const { result, rerender, unmount } = renderHook(
+      ({ store }) => ({
+        shallow: useStoreSelection(store, (state) => state.count),
+        plain: useStoreSelectionPlain(store, (state) => state.count),
+      }),
+      { initialProps: { store: first } }
+    );
+    rerender({ store: second });
+    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 });
+    unmount();
+  });
+
+  it("should return full state when the selector is omitted", () => {
+    const store = createTestStore();
+    const { result } = renderHook(() => ({
+      shallow: useStoreSelection(store),
+      plain: useStoreSelectionPlain(store),
+    }));
+    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 shallow-equal values while plain selection follows reference changes", () => {
+    const store = createTestStore();
+    const { result } = renderHook(() => ({
+      shallow: useStoreSelection(store, (state) => state.items),
+      plain: useStoreSelectionPlain(store, (state) => state.items),
+    }));
+    const initialItems = store.getState().items;
+    const replacement = [1, 2];
+
+    act(() => store.setState({ items: replacement }));
+    expect(result.current.shallow).toBe(initialItems);
+    expect(result.current.plain).toBe(replacement);
+  });
+
+  it("should retain separate selections for each consumer", () => {
+    const store = createTestStore();
+    const { result } = renderHook(() => ({
+      count: useStoreSelection(store, (state) => state.count),
+      label: useStoreSelection(store, (state) => state.label),
+    }));
+    act(() => store.setState({ count: 2 }));
+    expect(result.current).toEqual({ count: 2, label: "initial" });
+  });
+});

src/hooks/storeSelection/index.ts

@@ -0,0 +1 @@
+export { useStoreSelection, useStoreSelectionPlain } from "./useStoreSelection";

src/hooks/storeSelection/storeSelection.types.ts

@@ -0,0 +1,6 @@
+import type { StoreApi } from "zustand";
+
+export type SelectionStore<TState> = Pick<
+  StoreApi<TState>,
+  "getState" | "getInitialState" | "subscribe"
+>;

src/hooks/storeSelection/useStoreSelection.ts

@@ -0,0 +1,30 @@
+import { useStore } from "zustand";
+import { shallow } from "zustand/shallow";
+import { identitySelector } from "../../utils";
+import { useSelectorWithEquality } from "../useSelectorWithEquality";
+import type { SelectionStore } from "./storeSelection.types";
+
+/** Selects from a store with store-scoped reference retention. */
+export function useStoreSelection<TState, TSelected = TState>(
+  store: SelectionStore<TState>,
+  selector?: (state: TState) => TSelected,
+  equalityFn?: (left: TSelected | TState, right: TSelected | TState) => boolean
+): TSelected | TState {
+  const actualSelector = useSelectorWithEquality({
+    cacheKey: store,
+    equalityFn: equalityFn ?? shallow,
+    selector: (state: TState): TSelected | TState => (selector ? selector(state) : state),
+  });
+
+  return useStore(store, actualSelector);
+}
+
+/** Keeps plain selection independent of equality caching. */
+export function useStoreSelectionPlain<TState, TSelected = TState>(
+  store: SelectionStore<TState>,
+  selector?: (state: TState) => TSelected
+): TSelected | TState {
+  const actualSelector = selector ?? identitySelector<TState>;
+
+  return useStore<typeof store, TSelected | TState>(store, actualSelector);
+}