refactor: deepen store selection module

src/hooks/index.ts

@@ -1,12 +1,4 @@
 export { createResolvedStoreHooks } from "./createResolvedStoreHooks";
 export type { ResolvedStoreBindings, StorePlainHook, StoreValueHook } from "../types";
-export { useSelectorWithEquality } from "./useSelectorWithEquality";
-export type {
-  UseSelectorWithEqualityEqualityFn,
-  UseSelectorWithEqualityOptions,
-  UseSelectorWithEqualityReturn,
-  UseSelectorWithEqualitySelector,
-} from "./useSelectorWithEquality";
 export { createStoreSelectionBindings } from "./storeSelection";
-export { useStoreSelection, useStoreSelectionPlain } from "./storeSelection";
 export type { StoreResolver, StoreSelectionBindings } from "./storeSelection";

src/hooks/storeSelection/createStoreSelectionBindings.ts

@@ -1,4 +1,4 @@
-import { useStoreSelection, useStoreSelectionPlain } from "./useStoreSelection";
+import { useStoreSelection, useStoreSelectionPlain } from "./storeSelection.hooks";
 import type { StoreResolver, StoreSelectionBindings } from "./storeSelection.types";
 
 export function createStoreSelectionBindings<TState>(

src/hooks/storeSelection/index.ts

@@ -1,3 +1,2 @@
 export { createStoreSelectionBindings } from "./createStoreSelectionBindings";
-export { useStoreSelection, useStoreSelectionPlain } from "./useStoreSelection";
 export type { StoreResolver, StoreSelectionBindings } from "./storeSelection.types";

src/hooks/storeSelection/storeSelection.hooks.ts

@@ -0,0 +1,56 @@
+import { useRef } from "react";
+import { useStore } from "zustand";
+import { shallow } from "zustand/shallow";
+import { identitySelector } from "./storeSelection.utils";
+import type { SelectionCache, SelectionStore } from "./storeSelection.types";
+
+function useRetainedStoreSelection<TState, TSelected>(
+  store: SelectionStore<TState>,
+  selector: (state: TState) => TSelected,
+  equalityFn: (left: TSelected, right: TSelected) => boolean
+): TSelected {
+  const cacheRef = useRef<SelectionCache<TSelected>>({
+    hasValue: false,
+    store,
+  });
+
+  function retainEqualSelection(state: TState): TSelected {
+    const nextValue = selector(state);
+    const cache = cacheRef.current;
+
+    if (cache.hasValue && Object.is(cache.store, store) && equalityFn(cache.value, nextValue)) {
+      return cache.value;
+    }
+
+    cacheRef.current = {
+      hasValue: true,
+      store,
+      value: nextValue,
+    };
+
+    return nextValue;
+  }
+
+  return useStore(store, retainEqualSelection);
+}
+
+export function useStoreSelection<TState, TSelected = TState>(
+  store: SelectionStore<TState>,
+  selector?: (state: TState) => TSelected,
+  equalityFn?: (left: TSelected | TState, right: TSelected | TState) => boolean
+): TSelected | TState {
+  return useRetainedStoreSelection(
+    store,
+    (state: TState): TSelected | TState => (selector ? selector(state) : state),
+    equalityFn ?? shallow
+  );
+}
+
+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);
+}

src/hooks/storeSelection/storeSelection.types.ts

@@ -8,6 +8,19 @@ export type SelectionStore<TState> = Pick<
 
 export type StoreResolver<TState> = () => SelectionStore<TState>;
 
+export interface EmptySelectionCache {
+  hasValue: false;
+  store: unknown;
+}
+
+export interface RetainedSelectionCache<TSelected> {
+  hasValue: true;
+  store: unknown;
+  value: TSelected;
+}
+
+export type SelectionCache<TSelected> = EmptySelectionCache | RetainedSelectionCache<TSelected>;
+
 export interface StoreSelectionBindings<TState> {
   useStoreValue: StoreValueHook<TState>;
   useStorePlain: StorePlainHook<TState>;