fix(vanilla/utils): improve selectAtom for more use cases

pmndrs · Feb 3, 2023 · c22d9f9 · c22d9f9
1 parent 0c195b0
commit c22d9f9
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 8 deletions.
diff --git a/docs/utilities/select.mdx b/docs/utilities/select.mdx
@@ -12,7 +12,7 @@ Ref: https://github.com/pmndrs/jotai/issues/36
 ```ts
 function selectAtom<Value, Slice>(
   anAtom: Atom<Value>,
-  selector: (v: Value) => Slice,
+  selector: (v: Value, prevSlice?: Slice) => Slice,
   equalityFn: (a: Slice, b: Slice) => boolean = Object.is
 ): Atom<Slice>
 ```

diff --git a/src/vanilla/utils/selectAtom.ts b/src/vanilla/utils/selectAtom.ts
@@ -17,19 +17,19 @@ const memo3 = <T>(
 
 export function selectAtom<Value, Slice>(
   anAtom: Atom<Promise<Value>>,
-  selector: (v: Awaited<Value>) => Slice,
+  selector: (v: Awaited<Value>, prevSlice?: Slice) => Slice,
   equalityFn?: (a: Slice, b: Slice) => boolean
 ): Atom<Promise<Slice>>
 
 export function selectAtom<Value, Slice>(
   anAtom: Atom<Value>,
-  selector: (v: Awaited<Value>) => Slice,
+  selector: (v: Awaited<Value>, prevSlice?: Slice) => Slice,
   equalityFn?: (a: Slice, b: Slice) => boolean
 ): Atom<Slice>
 
 export function selectAtom<Value, Slice>(
   anAtom: Atom<Value>,
-  selector: (v: Awaited<Value>) => Slice,
+  selector: (v: Awaited<Value>, prevSlice?: Slice) => Slice,
   equalityFn: (a: Slice, b: Slice) => boolean = Object.is
 ) {
   return memo3(
@@ -39,11 +39,11 @@ export function selectAtom<Value, Slice>(
         Awaited<Value>,
         Slice | typeof EMPTY
       ]) => {
-        const slice = selector(value)
-        if (prevSlice !== EMPTY && equalityFn(prevSlice, slice)) {
-          return prevSlice
+        if (prevSlice === EMPTY) {
+          return selector(value)
         }
-        return slice
+        const slice = selector(value, prevSlice)
+        return equalityFn(prevSlice, slice) ? prevSlice : slice
       }
       const derivedAtom: Atom<Slice | Promise<Slice> | typeof EMPTY> & {
         init?: typeof EMPTY