From 2e990a42632092b745854ee2fae8c2713abb0a49 Mon Sep 17 00:00:00 2001
From: Daishi Kato <dai-shi@users.noreply.github.com>
Date: Tue, 10 Sep 2024 08:43:21 +0900
Subject: [PATCH] feat: useSelector (#6)

* feat: useSelector

* update CHANGELOG
---
 CHANGELOG.md                   |  4 +++
 examples/02_create/src/app.tsx |  5 ++-
 src/create.ts                  | 22 +++---------
 src/index.ts                   |  1 +
 src/useSelector.ts             | 65 ++++++++++++++++++++++++++++++++++
 5 files changed, 77 insertions(+), 20 deletions(-)
 create mode 100644 src/useSelector.ts

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 8fed4b7..353e617 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,10 @@
 
 ## [Unreleased]
 
+### Added
+
+- feat: useSelector #6
+
 ## [0.5.0] - 2024-09-09
 
 ### Added
diff --git a/examples/02_create/src/app.tsx b/examples/02_create/src/app.tsx
index 97109b5..7180be1 100644
--- a/examples/02_create/src/app.tsx
+++ b/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 (
     <>
diff --git a/src/create.ts b/src/create.ts
index 9adc614..e9ca64e 100644
--- a/src/create.ts
+++ b/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);
 }
diff --git a/src/index.ts b/src/index.ts
index 980a8ad..36f55c2 100644
--- a/src/index.ts
+++ b/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';
diff --git a/src/useSelector.ts b/src/useSelector.ts
new file mode 100644
index 0000000..2dde33d
--- /dev/null
+++ b/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;
+}