feat: sync jotai renders with useState renders (pmndrs#827) (pmndrs#841)

* feat: test case * a workaround for useReducer weird behavior * fix a test for react18 * fix bail out check Co-authored-by: daishi <[email protected]>
imwillhayes · Nov 24, 2021 · f315f87 · f315f87
1 parent 2c0791a
commit f315f87
Show file tree

Hide file tree

Showing 2 changed files with 59 additions and 11 deletions.
diff --git a/src/core/useAtom.ts b/src/core/useAtom.ts
@@ -54,7 +54,29 @@ export function useAtom<Value, Update, Result extends void | Promise<void>>(
     throw new Error('no atom value')
   }, [store, atom])
 
-  const [value, forceUpdate] = useReducer(getAtomValue, undefined, getAtomValue)
+  const [[value, atomFromUseReducer], forceUpdate] = useReducer(
+    useCallback(
+      (prev) => {
+        const nextValue = getAtomValue()
+        if (Object.is(prev[0], nextValue) && prev[1] === atom) {
+          return prev // bail out
+        }
+        return [nextValue, atom]
+      },
+      [getAtomValue, atom]
+    ),
+    undefined,
+    () => {
+      const initialValue = getAtomValue()
+      return [initialValue, atom]
+    }
+  )
+
+  if (atomFromUseReducer !== atom) {
+    // Note: This seems like a useReducer bug, doesn't it?
+    // https://github.com/pmndrs/jotai/issues/827
+    forceUpdate()
+  }
 
   useEffect(() => {
     const unsubscribe = store[SUBSCRIBE_ATOM](atom, forceUpdate)

diff --git a/tests/basic.test.tsx b/tests/basic.test.tsx
@@ -788,18 +788,10 @@ it('changes atom from parent (#273, #275)', async () => {
   await findByText('commits: 1, id: a')
 
   fireEvent.click(getByText('atom b'))
-  if (IS_REACT18) {
-    await findByText('commits: 3, id: b')
-  } else {
-    await findByText('commits: 2, id: b')
-  }
+  await findByText('commits: 2, id: b')
 
   fireEvent.click(getByText('atom a'))
-  if (IS_REACT18) {
-    await findByText('commits: 5, id: a')
-  } else {
-    await findByText('commits: 3, id: a')
-  }
+  await findByText('commits: 3, id: a')
 })
 
 it('should be able to use a double derived atom twice and useEffect (#373)', async () => {
@@ -896,3 +888,37 @@ it('async chain for multiple sync and async atoms (#443)', async () => {
   await findByText('loading')
   await findByText('count: 3')
 })
+
+it('sync re-renders with useState re-renders (#827)', async () => {
+  const atom0 = atom('atom0')
+  const atom1 = atom('atom1')
+  const atom2 = atom('atom2')
+  const atoms = [atom0, atom1, atom2]
+
+  const App = () => {
+    const [currentAtomIndex, setCurrentAtomIndex] = useState(0)
+    const rotateAtoms = () => {
+      setCurrentAtomIndex((prev) => (prev + 1) % atoms.length)
+    }
+    const [atomValue] = useAtom(atoms[currentAtomIndex] as typeof atoms[number])
+
+    return (
+      <>
+        <span>commits: {useCommitCount()}</span>
+        <h1>{atomValue}</h1>
+        <button onClick={rotateAtoms}>rotate</button>
+      </>
+    )
+  }
+  const { findByText, getByText } = render(
+    <Provider>
+      <App />
+    </Provider>
+  )
+
+  await findByText('commits: 1')
+  fireEvent.click(getByText('rotate'))
+  await findByText('commits: 2')
+  fireEvent.click(getByText('rotate'))
+  await findByText('commits: 3')
+})