- proxy subscribe lives for the life of the mutable atom

- import types with type keyword
- add tests for updating proxy even when component is not mounted

__tests__/mutableAtom.test.tsx

@@ -1,8 +1,15 @@
 import React from 'react'
 import { act, render, renderHook, waitFor } from '@testing-library/react'
-import { atom, useAtom, useAtomValue, useSetAtom } from 'jotai'
+import {
+  Provider,
+  atom,
+  createStore,
+  useAtom,
+  useAtomValue,
+  useSetAtom,
+} from 'jotai'
 import assert from 'minimalistic-assert'
-import { mutableAtom } from '../src/mutableAtom'
+import { makeMutableAtom, mutableAtom } from '../src/mutableAtom'
 import type { ProxyState } from '../src/mutableAtom/types'
 
 it('should be defined on initial render', async () => {
@@ -274,7 +281,7 @@ it('should reject writing to properties other than `value`', async () => {
   }
   const { result } = renderHook(useTest)
   expect(async () => {
-    await act(() => {
+    await act(async () => {
       result.current.countProxy.value = 1
     })
   }).not.toThrow()
@@ -284,6 +291,47 @@ it('should reject writing to properties other than `value`', async () => {
   }).toThrow() // 'set' on proxy: trap returned falsish for property 'NOT_VALUE'
 })
 
+it('should allow updating even when component is unmounted', async () => {
+  expect.assertions(2)
+  const store = createStore()
+  const countAtom = atom({ value: 0 })
+  let isMounted = false
+  countAtom.onMount = () => {
+    isMounted = true
+  }
+  const mutableCountAtom = makeMutableAtom(countAtom)
+
+  function useTest() {
+    useAtomValue(mutableCountAtom)
+  }
+  function wrapper({ children }: { children: React.ReactNode }) {
+    return <Provider store={store}>{children}</Provider>
+  }
+
+  const { unmount } = renderHook(useTest, { wrapper })
+  await waitFor(() => assert(isMounted))
+  unmount()
+  expect(store.get(countAtom).value).toBe(0)
+  await act(async () => {
+    const countProxy = store.get(mutableCountAtom)
+    countProxy.value++
+  })
+  expect(store.get(countAtom).value).toBe(1)
+})
+
+it('should allow updating even when component has not mounted', async () => {
+  expect.assertions(2)
+  const store = createStore()
+  const countAtom = atom({ value: 0 })
+  const mutableCountAtom = makeMutableAtom(countAtom)
+  expect(store.get(countAtom).value).toBe(0)
+  await act(async () => {
+    const countProxy = store.get(mutableCountAtom)
+    countProxy.value++
+  })
+  expect(store.get(countAtom).value).toBe(1)
+})
+
 it('should correctly handle updates via writable atom', async () => {
   expect.assertions(3)
   const mutableCountAtom = mutableAtom(0)

src/mutableAtom/index.ts

@@ -1,7 +1,7 @@
 import { atom } from 'jotai'
-import type { Getter, Setter } from 'jotai'
+import type { Getter, PrimitiveAtom, Setter } from 'jotai'
 import { proxy, snapshot, subscribe } from 'valtio'
-import {
+import type {
   Options,
   PromiseOrValue,
   ProxyState,
@@ -16,48 +16,39 @@ export function mutableAtom<Value>(
   value: Value,
   options: Options<Value> = defaultOptions
 ) {
-  const { proxyFn } = { ...defaultOptions, ...options }
-
   const valueAtom = atom({ value })
 
   if (process.env.NODE_ENV !== 'production') {
     valueAtom.debugPrivate = true
   }
+  return makeMutableAtom(valueAtom, options)
+}
+
+export function makeMutableAtom<Value>(
+  valueAtom: PrimitiveAtom<Wrapped<Value>>,
+  options: Options<Value> = defaultOptions
+) {
+  const { proxyFn } = { ...defaultOptions, ...options }
 
   const storeAtom = atom(
     () =>
       ({
-        isMounted: false,
+        hasMounted: false,
         proxyState: null,
         unsubscribe: null,
       } as Store<Value>),
-    (get, set, isOnMount: boolean) => {
-      if (isOnMount) {
-        createProxyState(get, (fn) => fn(set))
-      } else {
-        onAtomUnmount(get)
-      }
+    (get, set) => {
+      // switch to synchronous imperative updates on mount
+      createProxyState(get, (fn) => fn(set))
     }
   )
 
-  storeAtom.onMount = (setOnMount) => {
-    // switch to synchronous imperative updates on mount
-    setOnMount(true)
-    return () => setOnMount(false)
-  }
+  storeAtom.onMount = (setInit) => setInit()
 
   if (process.env.NODE_ENV !== 'production') {
     storeAtom.debugPrivate = true
   }
 
-  /**
-   * unsubscribe on atom unmount
-   */
-  function onAtomUnmount(get: Getter) {
-    get(storeAtom).unsubscribe?.()
-    get(storeAtom).isMounted = false
-  }
-
   /**
    * sync the proxy state with the atom
    */
@@ -80,22 +71,18 @@ export function mutableAtom<Value>(
     const { value } = get(valueAtom)
     store.proxyState ??= proxyFn({ value })
     store.proxyState.value = value
-    const unsubscribe = subscribe(store.proxyState, onChange(get, setCb), true)
     store.unsubscribe?.()
-    store.unsubscribe = () => {
-      store.unsubscribe = null
-      unsubscribe()
-    }
-    store.isMounted = true
+    store.unsubscribe = subscribe(store.proxyState, onChange(get, setCb), true)
+    store.hasMounted = true
     return store.proxyState
   }
 
   /**
    * return the proxy if it exists, otherwise create and subscribe to it
    */
   function ensureProxyState(get: Getter, setCb: SetCb) {
-    const { isMounted, proxyState } = get(storeAtom)
-    if (proxyState === null || !isMounted) {
+    const { hasMounted, proxyState } = get(storeAtom)
+    if (proxyState === null || !hasMounted) {
       return createProxyState(get, setCb)
     }
     return proxyState
@@ -104,21 +91,13 @@ export function mutableAtom<Value>(
   /**
    * wrap the proxy state in a proxy to ensure rerender on value change
    */
-  function wrapProxyState(
-    proxyState: ProxyState<Value>,
-    get: Getter,
-    setCb: SetCb
-  ) {
+  function wrapProxyState(proxyState: ProxyState<Value>) {
     return new Proxy(proxyState, {
       get: (target, property) => {
-        if (property === 'value') {
-          ensureProxyState(get, setCb)
-        }
         return target[property as keyof ProxyState<Value>]
       },
       set(target, property, value) {
         if (property === 'value') {
-          ensureProxyState(get, setCb)
           target[property] = value
           return true
         }
@@ -135,7 +114,7 @@ export function mutableAtom<Value>(
       get(valueAtom) // subscribe to value updates
       const setCb = makeSetCb(setSelf)
       const proxyState = ensureProxyState(get, setCb)
-      return wrapProxyState(proxyState, get, setCb)
+      return wrapProxyState(proxyState)
     },
     (get, set, writeFn: WriteFn) => writeFn(get, set)
   )
@@ -162,7 +141,7 @@ const makeSetCb =
 
 /**
  * delays execution until next microtask
- * */
+ */
 function defer(fn?: () => PromiseOrValue<void>) {
   return Promise.resolve().then(fn)
 }

src/mutableAtom/types.ts

@@ -24,7 +24,7 @@ export type SetSelf<Args extends unknown[]> = SetAtom<Args, void>
 
 export type Store<Value> = {
   unsubscribe: CleanupFn | null
-  isMounted: boolean
+  hasMounted: boolean
   proxyState: ProxyState<Value> | null
 }