feat(core): Move useAtomValue and useSetAtom to core (pmndrs#989)

* Rename useUpdateAtom to useSetAtom

Move useAtomValue and useSetAtom to core

* Fix test

* Move exports, delete files

* Move impl to respective hooks

* Do not export ResolveType

* Cleanup after merge

* Changes after review

* refactor, add comments

* empty commit

* empty commit

Co-authored-by: daishi <[email protected]>

examples/hacker_news/src/App.tsx

@@ -1,8 +1,7 @@
 import { Suspense } from 'react'
 import { a, useSpring } from '@react-spring/web'
 import Parser from 'html-react-parser'
-import { Provider, atom, useAtom } from 'jotai'
-import { useUpdateAtom } from 'jotai/utils'
+import { Provider, atom, useAtom, useSetAtom } from 'jotai'
 
 type PostData = {
   by: string
@@ -34,9 +33,9 @@ function Id() {
 }
 
 function Next() {
-  // Use `useUpdateAtom` to avoid re-render
+  // Use `useSetAtom` to avoid re-render
   // const [, set] = useAtom(postId)
-  const setPostId = useUpdateAtom(postId)
+  const setPostId = useSetAtom(postId)
   return (
     <button onClick={() => setPostId((id) => id + 1)}>
       <div>→</div>

examples/todos/src/App.tsx

@@ -2,9 +2,8 @@ import type { FormEvent } from 'react'
 import { CloseOutlined } from '@ant-design/icons'
 import { a, useTransition } from '@react-spring/web'
 import { Radio } from 'antd'
-import { Provider, atom, useAtom } from 'jotai'
+import { Provider, atom, useAtom, useSetAtom } from 'jotai'
 import type { PrimitiveAtom } from 'jotai'
-import { useUpdateAtom } from 'jotai/utils'
 
 type Todo = {
   title: string
@@ -76,9 +75,9 @@ const Filtered = (props: FilteredType) => {
 }
 
 const TodoList = () => {
-  // Use `useUpdateAtom` to avoid re-render
+  // Use `useSetAtom` to avoid re-render
   // const [, setTodos] = useAtom(todosAtom)
-  const setTodos = useUpdateAtom(todosAtom)
+  const setTodos = useSetAtom(todosAtom)
   const remove: RemoveFn = (todo) =>
     setTodos((prev) => prev.filter((item) => item !== todo))
   const add = (e: FormEvent<HTMLFormElement>) => {

examples/todos_with_atomFamily/src/App.tsx

@@ -2,8 +2,8 @@ import type { FormEvent } from 'react'
 import { CloseOutlined } from '@ant-design/icons'
 import { a, useTransition } from '@react-spring/web'
 import { Radio } from 'antd'
-import { Provider, atom, useAtom } from 'jotai'
-import { atomFamily, useUpdateAtom } from 'jotai/utils'
+import { Provider, atom, useAtom, useSetAtom } from 'jotai'
+import { atomFamily } from 'jotai/utils'
 import { nanoid } from 'nanoid'
 
 type Param = { id: string; title?: string }
@@ -75,9 +75,9 @@ const Filtered = ({ remove }: { remove: (id: string) => void }) => {
 }
 
 const TodoList = () => {
-  // Use `useUpdateAtom` to avoid re-render
+  // Use `useSetAtom` to avoid re-render
   // const [, setTodos] = useAtom(todosAtom)
-  const setTodos = useUpdateAtom(todosAtom)
+  const setTodos = useSetAtom(todosAtom)
   const remove = (id: string) => {
     setTodos((prev) => prev.filter((item) => item !== id))
     todoAtomFamily.remove({ id })

src/core/useAtom.ts

@@ -1,23 +1,9 @@
-import {
-  useCallback,
-  useContext,
-  useDebugValue,
-  useEffect,
-  useReducer,
-} from 'react'
-import type { Reducer } from 'react'
 import type { Atom, Scope, SetAtom, WritableAtom } from './atom'
-import { getScopeContext } from './contexts'
-import { COMMIT_ATOM, READ_ATOM, SUBSCRIBE_ATOM, WRITE_ATOM } from './store'
-import type { VersionObject } from './store'
+import { useAtomValue } from './useAtomValue'
+import { useSetAtom } from './useSetAtom'
 
 type ResolveType<T> = T extends Promise<infer V> ? V : T
 
-const isWritable = <Value, Update, Result extends void | Promise<void>>(
-  atom: Atom<Value> | WritableAtom<Value, Update, Result>
-): atom is WritableAtom<Value, Update, Result> =>
-  !!(atom as WritableAtom<Value, Update, Result>).write
-
 export function useAtom<Value, Update, Result extends void | Promise<void>>(
   atom: WritableAtom<Value, Update, Result>,
   scope?: Scope
@@ -38,85 +24,9 @@ export function useAtom<Value, Update, Result extends void | Promise<void>>(
     )
     scope = (atom as { scope: Scope }).scope
   }
-
-  const ScopeContext = getScopeContext(scope)
-  const { s: store, w: versionedWrite } = useContext(ScopeContext)
-
-  const getAtomValue = useCallback(
-    (version?: VersionObject) => {
-      // This call to READ_ATOM is the place where derived atoms will actually be
-      // recomputed if needed.
-      const atomState = store[READ_ATOM](atom, version)
-      if ('e' in atomState) {
-        throw atomState.e // read error
-      }
-      if ('p' in atomState) {
-        throw atomState.p // read promise
-      }
-      if ('v' in atomState) {
-        return atomState.v as ResolveType<Value>
-      }
-      throw new Error('no atom value')
-    },
-    [store, atom]
-  )
-
-  // Pull the atoms's state from the store into React state.
-  const [[version, value, atomFromUseReducer], rerenderIfChanged] = useReducer<
-    Reducer<
-      readonly [VersionObject | undefined, ResolveType<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) {
-    rerenderIfChanged(undefined)
-  }
-
-  useEffect(() => {
-    // Call `rerenderIfChanged` whenever this atom is invalidated. Note
-    // that derived atoms may not be recomputed yet.
-    const unsubscribe = store[SUBSCRIBE_ATOM](atom, rerenderIfChanged)
-    rerenderIfChanged(undefined)
-    return unsubscribe
-  }, [store, atom])
-
-  useEffect(() => {
-    store[COMMIT_ATOM](atom, version)
-  })
-
-  const setAtom = useCallback(
-    (update: Update) => {
-      if (isWritable(atom)) {
-        const write = (version?: VersionObject) =>
-          store[WRITE_ATOM](atom, update, version)
-        return versionedWrite ? versionedWrite(write) : write()
-      } else {
-        throw new Error('not writable atom')
-      }
-    },
-    [store, versionedWrite, atom]
-  )
-
-  useDebugValue(value)
-  return [value, setAtom]
+  return [
+    useAtomValue(atom, scope),
+    // We do wrong type assertion here, which results in throwing an error.
+    useSetAtom(atom as WritableAtom<Value, Update, Result>, scope),
+  ]
 }

src/core/useAtomValue.ts

@@ -0,0 +1,87 @@
+import {
+  useCallback,
+  useContext,
+  useDebugValue,
+  useEffect,
+  useReducer,
+} from 'react'
+import type { Reducer } from 'react'
+import type { Atom, Scope } from './atom'
+import { getScopeContext } from './contexts'
+import { COMMIT_ATOM, READ_ATOM, SUBSCRIBE_ATOM } from './store'
+import type { VersionObject } from './store'
+
+type ResolveType<T> = T extends Promise<infer V> ? V : T
+
+export function useAtomValue<Value>(
+  atom: Atom<Value>,
+  scope?: Scope
+): ResolveType<Value> {
+  const ScopeContext = getScopeContext(scope)
+  const { s: store } = useContext(ScopeContext)
+
+  const getAtomValue = useCallback(
+    (version?: VersionObject) => {
+      // This call to READ_ATOM is the place where derived atoms will actually be
+      // recomputed if needed.
+      const atomState = store[READ_ATOM](atom, version)
+      if ('e' in atomState) {
+        throw atomState.e // read error
+      }
+      if ('p' in atomState) {
+        throw atomState.p // read promise
+      }
+      if ('v' in atomState) {
+        return atomState.v as ResolveType<Value>
+      }
+      throw new Error('no atom value')
+    },
+    [store, atom]
+  )
+
+  // Pull the atoms's state from the store into React state.
+  const [[version, value, atomFromUseReducer], rerenderIfChanged] = useReducer<
+    Reducer<
+      readonly [VersionObject | undefined, ResolveType<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) {
+    rerenderIfChanged(undefined)
+  }
+
+  useEffect(() => {
+    // Call `rerenderIfChanged` whenever this atom is invalidated. Note
+    // that derived atoms may not be recomputed yet.
+    const unsubscribe = store[SUBSCRIBE_ATOM](atom, rerenderIfChanged)
+    rerenderIfChanged(undefined)
+    return unsubscribe
+  }, [store, atom])
+
+  useEffect(() => {
+    store[COMMIT_ATOM](atom, version)
+  })
+
+  useDebugValue(value)
+  return value
+}

src/core/useSetAtom.ts

@@ -0,0 +1,31 @@
+import { useCallback, useContext } from 'react'
+import type { Scope, SetAtom, WritableAtom } from '../core/atom'
+import { WRITE_ATOM } from '../core/store'
+import type { VersionObject } from '../core/store'
+import { getScopeContext } from './contexts'
+
+export function useSetAtom<Value, Update, Result extends void | Promise<void>>(
+  atom: WritableAtom<Value, Update, Result>,
+  scope?: Scope
+): SetAtom<Update, Result> {
+  const ScopeContext = getScopeContext(scope)
+  const { s: store, w: versionedWrite } = useContext(ScopeContext)
+  const setAtom = useCallback(
+    (update: Update) => {
+      if (
+        !('write' in atom) &&
+        typeof process === 'object' &&
+        process.env.NODE_ENV !== 'production'
+      ) {
+        // useAtom can pass non writable atom with wrong type assertion,
+        // so we should check here.
+        throw new Error('not writable atom')
+      }
+      const write = (version?: VersionObject) =>
+        store[WRITE_ATOM](atom, update, version)
+      return versionedWrite ? versionedWrite(write) : write()
+    },
+    [store, versionedWrite, atom]
+  )
+  return setAtom as SetAtom<Update, Result>
+}

src/index.ts

@@ -1,6 +1,8 @@
 export { Provider } from './core/Provider'
 export { atom } from './core/atom'
 export { useAtom } from './core/useAtom'
+export { useAtomValue } from './core/useAtomValue'
+export { useSetAtom } from './core/useSetAtom'
 export { createStoreForExport as unstable_createStore } from './core/store'
 export type { Atom, WritableAtom, PrimitiveAtom } from './core/atom'
 export type {

src/utils.ts

@@ -1,6 +1,6 @@
 export { RESET } from './utils/constants'
-export { useUpdateAtom } from './utils/useUpdateAtom'
-export { useAtomValue } from './utils/useAtomValue'
+export { useSetAtom as useUpdateAtom } from 'jotai'
+export { useAtomValue } from 'jotai'
 export { atomWithReset } from './utils/atomWithReset'
 export { useResetAtom } from './utils/useResetAtom'
 export { useReducerAtom } from './utils/useReducerAtom'

src/utils/useAtomCallback.ts

@@ -1,9 +1,7 @@
 import { useCallback, useMemo } from 'react'
-import { atom } from 'jotai'
+import { atom, useSetAtom } from 'jotai'
 import type { Setter, WritableAtom } from 'jotai'
 import type { Scope } from '../core/atom'
-// NOTE importing non-core functions is generally not allowed. this is an exception.
-import { useUpdateAtom } from './useUpdateAtom'
 
 type WriteGetter = Parameters<WritableAtom<unknown, unknown>['write']>[0]
 
@@ -47,7 +45,7 @@ export function useAtomCallback<Result, Arg>(
       ),
     [callback]
   )
-  const invoke = useUpdateAtom(anAtom, scope)
+  const invoke = useSetAtom(anAtom, scope)
   return useCallback(
     (arg: Arg) =>
       new Promise<Result>((resolve, reject) => {

tests/utils/useAtomValue.test.tsx → tests/useAtomValue.test.tsx

@@ -1,7 +1,6 @@
 import { fireEvent, render } from '@testing-library/react'
-import { atom } from 'jotai'
-import { useAtomValue, useUpdateAtom } from 'jotai/utils'
-import { getTestProvider } from '../testUtils'
+import { atom, useAtomValue, useSetAtom } from 'jotai'
+import { getTestProvider } from './testUtils'
 
 const Provider = getTestProvider()
 
@@ -10,7 +9,7 @@ it('useAtomValue basic test', async () => {
 
   const Counter = () => {
     const count = useAtomValue(countAtom)
-    const setCount = useUpdateAtom(countAtom)
+    const setCount = useSetAtom(countAtom)
 
     return (
       <>