Counter.tsx

import {view} from './jotai-view'
import {atom} from 'jotai'

const countAtom = atom(0)

const Counter = view((props, get, set) => {
  return (
    <button onClick={() => set(countAtom, get(countAtom) + 1)}>
      count is {get(countAtom)}
    </button>
  )
})

jotai-view.ts

import {Atom, WritableAtom, useStore} from 'jotai'
import {ReactNode, useEffect, useMemo, useState} from 'react'

type Getter = <Value>(atom: Atom<Value>) => Value
type Setter = <Value, Args extends unknown[], Result>(
  atom: WritableAtom<Value, Args, Result>,
  ...args: Args
) => Result

export function view<Props>(
  render: (props: Props, get: Getter, set: Setter) => ReactNode
) {
  return (props: Props) => {
    const dependencies = useMemo(() => new Set<Atom<unknown>>(), [])
    const tracking = useMemo(() => new Map<Atom<unknown>, () => void>(), [])
    const [, setRevision] = useState(0)
    const store = useStore()
    const get = <Value>(atom: Atom<Value>) => {
      dependencies.add(atom)
      return store.get(atom)
    }
    useEffect(() => {
      for (const dependency of dependencies) {
        if (tracking.has(dependency)) return
        tracking.set(
          dependency,
          store.sub(dependency, () => setRevision(r => r + 1))
        )
      }
      for (const [dependency, cancel] of tracking.entries()) {
        if (!dependencies.has(dependency)) cancel()
        tracking.delete(dependency)
      }
      return () => dependencies.clear()
    })
    useEffect(() => {
      return () => {
        for (const cancel of tracking.values()) cancel()
      }
    }, [])
    return render(props, get, store.set)
  }
}