feat: useSelector (#6)

* feat: useSelector

* update CHANGELOG

CHANGELOG.md

@@ -2,6 +2,10 @@
 
 ## [Unreleased]
 
+### Added
+
+- feat: useSelector #6
+
 ## [0.5.0] - 2024-09-09
 
 ### Added

examples/02_create/src/app.tsx

@@ -1,4 +1,3 @@
-import { useCallback } from 'react';
 import { create } from 'jotai-zustand';
 
 const useCountStore = create(
@@ -11,8 +10,8 @@ const useCountStore = create(
 );
 
 const Counter = () => {
-  const count = useCountStore(useCallback((state) => state.count, []));
-  const inc = useCountStore(useCallback((state) => state.inc, []));
+  const count = useCountStore((state) => state.count);
+  const inc = useCountStore((state) => state.inc);
 
   return (
     <>

src/create.ts

@@ -1,8 +1,5 @@
-import { useMemo } from 'react';
-import { atom, createStore } from 'jotai/vanilla';
-import { useAtomValue } from 'jotai/react';
-
 import { atomWithActions } from './atomWithActions.js';
+import { useSelector } from './useSelector.js';
 
 export function create<State extends object, Actions extends object>(
   initialState: State,
@@ -11,18 +8,9 @@ export function create<State extends object, Actions extends object>(
     get: () => State,
   ) => Actions,
 ) {
-  const store = createStore();
   const theAtom = atomWithActions(initialState, createActions);
-  const useStore = <Slice>(selector: (state: State & Actions) => Slice) => {
-    const derivedAtom = useMemo(
-      () => atom((get) => selector(get(theAtom))),
-      [selector],
-    );
-    return useAtomValue(derivedAtom, { store });
-  };
-  const useStoreWithGetState = useStore as typeof useStore & {
-    getState: () => State & Actions;
-  };
-  useStoreWithGetState.getState = () => store.get(theAtom);
-  return useStoreWithGetState;
+  return <Slice>(
+    selector: (state: State & Actions) => Slice,
+    equalityFn?: (a: Slice, b: Slice) => boolean,
+  ) => useSelector(theAtom, selector, equalityFn);
 }

src/index.ts

@@ -1,3 +1,4 @@
 export { atomWithStore } from './atomWithStore.js';
 export { atomWithActions } from './atomWithActions.js';
+export { useSelector } from './useSelector.js';
 export { create } from './create.js';

src/useSelector.ts

@@ -0,0 +1,65 @@
+import { useEffect, useReducer } from 'react';
+import type { ReducerWithoutAction } from 'react';
+import type { Atom, ExtractAtomValue } from 'jotai/vanilla';
+import { useStore } from 'jotai/react';
+
+type Store = ReturnType<typeof useStore>;
+
+type Options = Parameters<typeof useStore>[0];
+
+export function useSelector<Value, Slice>(
+  atom: Atom<Value>,
+  selector: (value: Value) => Slice,
+  equalityFn?: (a: Slice, b: Slice) => boolean,
+  options?: Options,
+): Awaited<Slice>;
+
+export function useSelector<AtomType extends Atom<unknown>, Slice>(
+  atom: AtomType,
+  selector: (value: ExtractAtomValue<AtomType>) => Slice,
+  equalityFn?: (a: Slice, b: Slice) => boolean,
+  options?: Options,
+): Awaited<Slice>;
+
+export function useSelector<Value, Slice>(
+  atom: Atom<Value>,
+  selector: (value: Value) => Slice,
+  equalityFn: (a: Slice, b: Slice) => boolean = Object.is,
+  options?: Options,
+) {
+  const store = useStore(options);
+
+  const [[sliceFromReducer, storeFromReducer, atomFromReducer], rerender] =
+    useReducer<
+      ReducerWithoutAction<readonly [Slice, Store, typeof atom]>,
+      undefined
+    >(
+      (prev) => {
+        const nextSlice = selector(store.get(atom));
+        if (
+          equalityFn(prev[0], nextSlice) &&
+          prev[1] === store &&
+          prev[2] === atom
+        ) {
+          return prev;
+        }
+        return [nextSlice, store, atom];
+      },
+      undefined,
+      () => [selector(store.get(atom)), store, atom],
+    );
+
+  let slice = sliceFromReducer;
+  if (storeFromReducer !== store || atomFromReducer !== atom) {
+    rerender();
+    slice = selector(store.get(atom));
+  }
+
+  useEffect(() => {
+    const unsub = store.sub(atom, () => rerender());
+    rerender();
+    return unsub;
+  }, [store, atom]);
+
+  return slice;
+}