deprecate devstore

src/vanilla.ts

@@ -1,7 +1,11 @@
 export { atom } from './vanilla/atom.ts'
 export type { Atom, WritableAtom, PrimitiveAtom } from './vanilla/atom.ts'
 
-export { createStore, getDefaultStore } from './vanilla/store.ts'
+export {
+  createStore,
+  getDefaultStore,
+  INTERNAL_overrideCreateStore,
+} from './vanilla/store.ts'
 
 export type {
   Getter,

src/vanilla/store.ts

@@ -10,6 +10,7 @@ export type INTERNAL_PrdStore = INTERNAL_Store
 
 // For debugging purpose only
 // This will be removed in the near future
+/* @deprecated Deprecated: Use devstore from the devtools library */
 export type INTERNAL_DevStoreRev4 = {
   dev4_get_internal_weak_map: () => {
     get: (atom: Atom<unknown>) => INTERNAL_AtomState | undefined
@@ -20,6 +21,7 @@ export type INTERNAL_DevStoreRev4 = {
   ) => void
 }
 
+/* @deprecated Deprecated: Use devstore from the devtools library */
 const createDevStoreRev4 = (): INTERNAL_PrdStore & INTERNAL_DevStoreRev4 => {
   let inRestoreAtom = 0
   const storeHooks = INTERNAL_initializeStoreHooks({})
@@ -56,7 +58,10 @@ const createDevStoreRev4 = (): INTERNAL_PrdStore & INTERNAL_DevStoreRev4 => {
   })
   const devStore: INTERNAL_DevStoreRev4 = {
     // store dev methods (these are tentative and subject to change without notice)
-    dev4_get_internal_weak_map: () => atomStateMap,
+    dev4_get_internal_weak_map: () => {
+      console.log('Deprecated: Use devstore from the devtools library')
+      return atomStateMap
+    },
     dev4_get_mounted_atoms: () => debugMountedAtoms,
     dev4_restore_atoms: (values) => {
       const restoreAtom: WritableAtom<null, [], void> = {
@@ -84,12 +89,22 @@ type PrdOrDevStore =
   | INTERNAL_PrdStore
   | (INTERNAL_PrdStore & INTERNAL_DevStoreRev4)
 
+let overiddenCreateStore: typeof createStore | undefined
+
+export function INTERNAL_overrideCreateStore(
+  fn: (prev: typeof createStore | undefined) => typeof createStore,
+): void {
+  overiddenCreateStore = fn(overiddenCreateStore)
+}
+
 export function createStore(): PrdOrDevStore {
+  if (overiddenCreateStore) {
+    return overiddenCreateStore()
+  }
   if (import.meta.env?.MODE !== 'production') {
     return createDevStoreRev4()
   }
-  const store = INTERNAL_buildStore()
-  return store
+  return INTERNAL_buildStore()
 }
 
 let defaultStore: PrdOrDevStore | undefined

tests/vanilla/storedev.test.tsx

@@ -1,48 +1,112 @@
 import { describe, expect, it, vi } from 'vitest'
-import { atom, createStore } from 'jotai/vanilla'
+import { atom } from 'jotai/vanilla'
+import type { Atom, WritableAtom } from 'jotai/vanilla'
+import {
+  INTERNAL_buildStoreRev1 as INTERNAL_buildStore,
+  INTERNAL_initializeStoreHooks,
+} from 'jotai/vanilla/internals'
 import type {
-  INTERNAL_DevStoreRev4,
-  INTERNAL_PrdStore,
-} from 'jotai/vanilla/store'
+  INTERNAL_AtomState,
+  INTERNAL_Store,
+} from 'jotai/vanilla/internals'
 
-describe('[DEV-ONLY] dev-only methods rev4', () => {
+type DevStore = {
+  get_internal_weak_map: () => {
+    get: (atom: Atom<unknown>) => INTERNAL_AtomState | undefined
+  }
+  get_mounted_atoms: () => Set<Atom<unknown>>
+  restore_atoms: (values: Iterable<readonly [Atom<unknown>, unknown]>) => void
+}
+
+const createDevStore = (): INTERNAL_Store & DevStore => {
+  let inRestoreAtom = 0
+  const storeHooks = INTERNAL_initializeStoreHooks({})
+  const atomStateMap = new WeakMap()
+  const mountedAtoms = new WeakMap()
+  const store = INTERNAL_buildStore(
+    atomStateMap,
+    mountedAtoms,
+    undefined,
+    undefined,
+    undefined,
+    undefined,
+    storeHooks,
+    undefined,
+    (atom, get, set, ...args) => {
+      if (inRestoreAtom) {
+        return set(atom, ...(args as any))
+      }
+      return atom.write(get, set, ...(args as any))
+    },
+  )
+  const debugMountedAtoms = new Set<Atom<unknown>>()
+  storeHooks.m.add(undefined, (atom) => {
+    debugMountedAtoms.add(atom)
+    const atomState = atomStateMap.get(atom)
+    // For DevStoreRev4 compatibility
+    ;(atomState as any).m = mountedAtoms.get(atom)
+  })
+  storeHooks.u.add(undefined, (atom) => {
+    debugMountedAtoms.delete(atom)
+    const atomState = atomStateMap.get(atom)
+    // For DevStoreRev4 compatibility
+    delete (atomState as any).m
+  })
+  const devStore: DevStore = {
+    // store dev methods (these are tentative and subject to change without notice)
+    get_internal_weak_map: () => atomStateMap,
+    get_mounted_atoms: () => debugMountedAtoms,
+    restore_atoms: (values) => {
+      const restoreAtom: WritableAtom<null, [], void> = {
+        read: () => null,
+        write: (_get, set) => {
+          ++inRestoreAtom
+          try {
+            for (const [atom, value] of values) {
+              if ('init' in atom) {
+                set(atom as never, value)
+              }
+            }
+          } finally {
+            --inRestoreAtom
+          }
+        },
+      }
+      store.set(restoreAtom)
+    },
+  }
+  return Object.assign(store, devStore)
+}
+
+describe('dev-only methods', () => {
   it('should get atom value', () => {
-    const store = createStore() as INTERNAL_DevStoreRev4 & INTERNAL_PrdStore
-    if (!('dev4_get_internal_weak_map' in store)) {
-      throw new Error('dev methods are not available')
-    }
+    const store = createDevStore()
     const countAtom = atom(0)
     store.set(countAtom, 1)
-    const weakMap = store.dev4_get_internal_weak_map()
+    const weakMap = store.get_internal_weak_map()
     expect(weakMap.get(countAtom)?.v).toEqual(1)
   })
 
   it('should restore atoms and its dependencies correctly', () => {
-    const store = createStore() as INTERNAL_DevStoreRev4 & INTERNAL_PrdStore
-    if (!('dev4_restore_atoms' in store)) {
-      throw new Error('dev methods are not available')
-    }
+    const store = createDevStore()
     const countAtom = atom(0)
     const derivedAtom = atom((get) => get(countAtom) * 2)
     store.set(countAtom, 1)
-    store.dev4_restore_atoms([[countAtom, 2]])
+    store.restore_atoms([[countAtom, 2]])
     expect(store.get(countAtom)).toBe(2)
     expect(store.get?.(derivedAtom)).toBe(4)
   })
 
   it('should restore atoms and call store listeners correctly', () => {
-    const store = createStore() as INTERNAL_DevStoreRev4 & INTERNAL_PrdStore
-    if (!('dev4_restore_atoms' in store)) {
-      throw new Error('dev methods are not available')
-    }
+    const store = createDevStore()
     const countAtom = atom(0)
     const derivedAtom = atom((get) => get(countAtom) * 2)
     const countCb = vi.fn()
     const derivedCb = vi.fn()
     store.set(countAtom, 2)
     const unsubCount = store.sub(countAtom, countCb)
     const unsubDerived = store.sub(derivedAtom, derivedCb)
-    store.dev4_restore_atoms([
+    store.restore_atoms([
       [countAtom, 1],
       [derivedAtom, 2],
     ])
@@ -54,15 +118,12 @@ describe('[DEV-ONLY] dev-only methods rev4', () => {
   })
 
   it('should return all the mounted atoms correctly', () => {
-    const store = createStore() as INTERNAL_DevStoreRev4 & INTERNAL_PrdStore
-    if (!('dev4_get_mounted_atoms' in store)) {
-      throw new Error('dev methods are not available')
-    }
+    const store = createDevStore()
     const countAtom = atom(0)
     const derivedAtom = atom((get) => get(countAtom) * 2)
     const unsub = store.sub(derivedAtom, vi.fn())
     store.set(countAtom, 1)
-    const result = store.dev4_get_mounted_atoms()
+    const result = store.get_mounted_atoms()
     expect(
       Array.from(result)
         .sort((a, b) => Object.keys(a).length - Object.keys(b).length)
@@ -83,26 +144,20 @@ describe('[DEV-ONLY] dev-only methods rev4', () => {
   })
 
   it("should return all the mounted atoms correctly after they're unsubscribed", () => {
-    const store = createStore() as INTERNAL_DevStoreRev4 & INTERNAL_PrdStore
-    if (!('dev4_get_mounted_atoms' in store)) {
-      throw new Error('dev methods are not available')
-    }
+    const store = createDevStore()
     const countAtom = atom(0)
     const derivedAtom = atom((get) => get(countAtom) * 2)
     const unsub = store.sub(derivedAtom, vi.fn())
     store.set(countAtom, 1)
     unsub()
-    const result = store.dev4_get_mounted_atoms()
+    const result = store.get_mounted_atoms()
     expect(Array.from(result)).toStrictEqual([])
   })
 
   it('should restore atoms with custom write function', () => {
-    const store = createStore() as INTERNAL_DevStoreRev4 & INTERNAL_PrdStore
-    if (!('dev4_restore_atoms' in store)) {
-      throw new Error('dev methods are not available')
-    }
+    const store = createDevStore()
     const countAtom = atom(0, () => {})
-    store.dev4_restore_atoms([[countAtom, 1]])
+    store.restore_atoms([[countAtom, 1]])
     expect(store.get(countAtom)).toBe(1)
   })
 })