useJotai.ts

import {Atom, useStore} from 'jotai'
import {useCallback, useEffect, useReducer, useRef} from 'react'

/*
const countAtom = atom(0)
function Counter() {
  const [get, set] = useJotai()
  return (
    <button onClick={() => set(countAtom, get(countAtom) + 1)}>
      count is {get(countAtom)}
    </button>
  )
}
*/

export function useJotai() {
  const store = useStore()
  const [, forceUpdate] = useReducer(r => r + 1, 0)
  const dependencies = useRef(new Set<Atom<unknown>>()).current
  const tracking = useRef(new Map<Atom<unknown>, () => void>()).current
  const get = useCallback(
    <Value>(atom: Atom<Value>) => {
      dependencies.add(atom)
      return store.get(atom)
    },
    [store]
  )
  dependencies.clear()
  useEffect(() => {
    for (const dependency of dependencies) {
      if (tracking.has(dependency)) continue
      tracking.set(dependency, store.sub(dependency, forceUpdate))
    }
    for (const [dependency, cancel] of tracking) {
      if (dependencies.has(dependency)) continue
      cancel()
      tracking.delete(dependency)
    }
  })
  useEffect(() => {
    return () => {
      for (const cancel of tracking.values()) cancel()
      tracking.clear()
    }
  }, [tracking])
  return [get, store.set] as const
}