fix(core): useAtom should return correct values when changing atoms (#…

…1236)

* add a failing test

* fix useAtom when changing an atom

* minor renaming

src/core/useAtomValue.ts

@@ -40,34 +40,37 @@ export function useAtomValue<Value>(
   )
 
   // Pull the atoms's state from the store into React state.
-  const [[version, value, atomFromUseReducer], rerenderIfChanged] = useReducer<
-    Reducer<
-      readonly [VersionObject | undefined, Awaited<Value>, Atom<Value>],
-      VersionObject | undefined
-    >,
-    undefined
-  >(
-    useCallback(
-      (prev, nextVersion) => {
-        const nextValue = getAtomValue(nextVersion)
-        if (Object.is(prev[1], nextValue) && prev[2] === atom) {
-          return prev // bail out
-        }
-        return [nextVersion, nextValue, atom]
-      },
-      [getAtomValue, atom]
-    ),
-    undefined,
-    () => {
-      // NOTE should/could branch on mount?
-      const initialVersion = undefined
-      const initialValue = getAtomValue(initialVersion)
-      return [initialVersion, initialValue, atom]
-    }
-  )
+  const [[version, valueFromReducer, atomFromReducer], rerenderIfChanged] =
+    useReducer<
+      Reducer<
+        readonly [VersionObject | undefined, Awaited<Value>, Atom<Value>],
+        VersionObject | undefined
+      >,
+      undefined
+    >(
+      useCallback(
+        (prev, nextVersion) => {
+          const nextValue = getAtomValue(nextVersion)
+          if (Object.is(prev[1], nextValue) && prev[2] === atom) {
+            return prev // bail out
+          }
+          return [nextVersion, nextValue, atom]
+        },
+        [getAtomValue, atom]
+      ),
+      undefined,
+      () => {
+        // NOTE should/could branch on mount?
+        const initialVersion = undefined
+        const initialValue = getAtomValue(initialVersion)
+        return [initialVersion, initialValue, atom]
+      }
+    )
 
-  if (atomFromUseReducer !== atom) {
+  let value = valueFromReducer
+  if (atomFromReducer !== atom) {
     rerenderIfChanged(undefined)
+    value = getAtomValue()
   }
 
   useEffect(() => {

tests/basic.test.tsx

@@ -9,7 +9,7 @@ import {
 import { fireEvent, render, waitFor } from '@testing-library/react'
 import ReactDOM, { unstable_batchedUpdates } from 'react-dom'
 import { atom, useAtom } from 'jotai'
-import type { WritableAtom } from 'jotai'
+import type { PrimitiveAtom, WritableAtom } from 'jotai'
 import { getTestProvider } from './testUtils'
 
 const Provider = getTestProvider()
@@ -1021,3 +1021,33 @@ it('onMount is not called when atom value is accessed from writeGetter in derive
   expect(onMount).not.toBeCalled()
   expect(onUnmount).not.toBeCalled()
 })
+
+it('useAtom returns consistent value with input with changing atoms (#1235)', async () => {
+  const countAtom = atom(0)
+  const valueAtoms = [atom(0), atom(1)]
+
+  const Counter = () => {
+    const [count, setCount] = useAtom(countAtom)
+    const [value] = useAtom(valueAtoms[count] as PrimitiveAtom<number>)
+    if (count !== value) {
+      throw new Error('value mismatch')
+    }
+    return (
+      <>
+        <div>count: {count}</div>
+        <button onClick={() => setCount((c) => c + 1)}>button</button>
+      </>
+    )
+  }
+
+  const { getByText, findByText } = render(
+    <Provider>
+      <Counter />
+    </Provider>
+  )
+
+  await findByText('count: 0')
+
+  fireEvent.click(getByText('button'))
+  await findByText('count: 1')
+})