tests following the change

src/vanilla.ts

@@ -6,8 +6,12 @@ export {
   getDefaultStore,
 
   // Internal functions (subject to change without notice)
+  // In case you rely on them, be sure to pin the version
   INTERNAL_getSecretStoreMethods,
   INTERNAL_buildStore,
+  INTERNAL_createBatch,
+  INTERNAL_addBatchFunc,
+  INTERNAL_flushBatch,
 } from './vanilla/store.ts'
 
 export type {

src/vanilla/store.ts

@@ -194,7 +194,7 @@ const addBatchFunc = (
   batch: Batch,
   priority: BatchPriority,
   fn: () => void,
-) => {
+): void => {
   batch[priority].add(fn)
 }
 
@@ -227,7 +227,7 @@ const addBatchAtomDependent = (
 const getBatchAtomDependents = (batch: Batch, atom: AnyAtom) =>
   batch.D.get(atom)
 
-const flushBatch = (batch: Batch) => {
+const flushBatch = (batch: Batch): void => {
   let error: AnyError
   let hasError = false
   const call = (fn: () => void) => {
@@ -267,7 +267,7 @@ const getSecretStoreMethods = (store: Store): SecretStoreMethods =>
   store[INTERNAL_STORE_METHODS]
 
 type SecretStoreMethods = readonly [
-  ensureAtomState: <Value>(atom: Atom<Value>) => AtomState<Value>,
+  ensureAtomState: Parameters<BuildStore>[0],
   readAtomState: <Value>(
     batch: Batch | undefined,
     atom: Atom<Value>,
@@ -300,7 +300,10 @@ type Store = {
 }
 
 type BuildStore = (
-  ensureAtomState: <Value>(atom: Atom<Value>) => AtomState<Value>,
+  ensureAtomState: <Value>(
+    atom: Atom<Value>,
+    onInit?: (store: Store) => void,
+  ) => AtomState<Value>,
   atomRead?: <Value>(
     atom: Atom<Value>,
     ...params: Parameters<Atom<Value>['read']>
@@ -757,7 +760,9 @@ const deriveDevStoreRev4 = (store: Store): Store & DevStoreRev4 => {
   const [ensureAtomState] = store[INTERNAL_STORE_METHODS]
   let inRestoreAtom = 0
   const newEnsureAtomState: typeof ensureAtomState = (atom) => {
-    const atomState = ensureAtomState(atom)
+    const atomState = ensureAtomState(atom, () => {
+      atom.unstable_onInit?.(derivedStore)
+    })
     const originalMounted = atomState.h
     atomState.h = (batch) => {
       originalMounted?.(batch)
@@ -821,15 +826,18 @@ type PrdOrDevStore = Store | (Store & DevStoreRev4)
 
 export const createStore = (): PrdOrDevStore => {
   const atomStateMap = new WeakMap()
-  const ensureAtomState = <Value>(atom: Atom<Value>) => {
+  const ensureAtomState: Parameters<BuildStore>[0] = (
+    atom,
+    onInit = (s) => atom.unstable_onInit?.(s),
+  ) => {
     if (import.meta.env?.MODE !== 'production' && !atom) {
       throw new Error('Atom is undefined or null')
     }
     let atomState = atomStateMap.get(atom)
     if (!atomState) {
       atomState = { d: new Map(), p: new Set(), n: 0 }
       atomStateMap.set(atom, atomState)
-      atom.unstable_onInit?.(store)
+      onInit(store)
     }
     return atomState
   }
@@ -861,6 +869,9 @@ export const getDefaultStore = (): PrdOrDevStore => {
 export const INTERNAL_getSecretStoreMethods: typeof getSecretStoreMethods =
   getSecretStoreMethods
 export const INTERNAL_buildStore: typeof buildStore = buildStore
+export const INTERNAL_createBatch: typeof createBatch = createBatch
+export const INTERNAL_addBatchFunc: typeof addBatchFunc = addBatchFunc
+export const INTERNAL_flushBatch: typeof flushBatch = flushBatch
 
 // Internal types (subject to change without notice)
 export type INTERNAL_AtomState = AtomState

tests/vanilla/unstable_derive.test.tsx → tests/vanilla/derive.test.tsx

@@ -1,8 +1,26 @@
 import { describe, expect, it, vi } from 'vitest'
-import { atom, createStore } from 'jotai/vanilla'
+import {
+  atom,
+  INTERNAL_buildStore as buildStore,
+  createStore,
+  INTERNAL_getSecretStoreMethods as getSecretStoreMethods,
+} from 'jotai/vanilla'
 import type { Atom, Getter } from 'jotai/vanilla'
 
-describe('unstable_derive for scoping atoms', () => {
+type EnsureAtomState = ReturnType<typeof getSecretStoreMethods>[0]
+
+const deriveStore = (
+  store: ReturnType<typeof createStore>,
+  enhanceEnsureAtomState: (fn: EnsureAtomState) => EnsureAtomState,
+  atomRead?: Parameters<typeof buildStore>[1],
+) => {
+  const [ensureAtomState] = getSecretStoreMethods(store)
+  const newEnsureAtomState = enhanceEnsureAtomState(ensureAtomState)
+  const derivedStore = buildStore(newEnsureAtomState, atomRead)
+  return derivedStore
+}
+
+describe('derive for scoping atoms', () => {
   /**
    * a
    * S1[a]: a1
@@ -13,38 +31,20 @@ describe('unstable_derive for scoping atoms', () => {
     const scopedAtoms = new Set<Atom<unknown>>([a])
 
     const store = createStore()
-    const derivedStore = store.unstable_derive(
-      (
-        getAtomState,
-        setAtomState,
-        atomRead,
-        atomWrite,
-        atomOnInit,
-        atomOnMount,
-      ) => {
-        const scopedAtomStateMap = new WeakMap()
-        return [
-          (atom) => {
-            if (scopedAtoms.has(atom)) {
-              return scopedAtomStateMap.get(atom)
-            }
-            return getAtomState(atom)
-          },
-          (atom, atomState) => {
-            if (scopedAtoms.has(atom)) {
-              scopedAtomStateMap.set(atom, atomState)
-            } else {
-              setAtomState(atom, atomState)
-            }
-            return atomState
-          },
-          atomRead,
-          atomWrite,
-          atomOnInit,
-          atomOnMount,
-        ]
-      },
-    )
+    const scopedAtomStateMap = new WeakMap()
+    const derivedStore = deriveStore(store, (ensureAtomState) => {
+      return (atom) => {
+        if (scopedAtoms.has(atom)) {
+          let atomState = scopedAtomStateMap.get(atom)
+          if (!atomState) {
+            atomState = { d: new Map(), p: new Set(), n: 0 }
+            scopedAtomStateMap.set(atom, atomState)
+          }
+          return atomState
+        }
+        return ensureAtomState(atom)
+      }
+    })
 
     expect(store.get(a)).toBe('a')
     expect(derivedStore.get(a)).toBe('a')
@@ -69,38 +69,20 @@ describe('unstable_derive for scoping atoms', () => {
     const scopedAtoms = new Set<Atom<unknown>>([a])
 
     const store = createStore()
-    const derivedStore = store.unstable_derive(
-      (
-        getAtomState,
-        setAtomState,
-        atomRead,
-        atomWrite,
-        atomOnInit,
-        atomOnMount,
-      ) => {
-        const scopedAtomStateMap = new WeakMap()
-        return [
-          (atom) => {
-            if (scopedAtoms.has(atom)) {
-              return scopedAtomStateMap.get(atom)
-            }
-            return getAtomState(atom)
-          },
-          (atom, atomState) => {
-            if (scopedAtoms.has(atom)) {
-              scopedAtomStateMap.set(atom, atomState)
-            } else {
-              setAtomState(atom, atomState)
-            }
-            return atomState
-          },
-          atomRead,
-          atomWrite,
-          atomOnInit,
-          atomOnMount,
-        ]
-      },
-    )
+    const scopedAtomStateMap = new WeakMap()
+    const derivedStore = deriveStore(store, (ensureAtomState) => {
+      return (atom) => {
+        if (scopedAtoms.has(atom)) {
+          let atomState = scopedAtomStateMap.get(atom)
+          if (!atomState) {
+            atomState = { d: new Map(), p: new Set(), n: 0 }
+            scopedAtomStateMap.set(atom, atomState)
+          }
+          return atomState
+        }
+        return ensureAtomState(atom)
+      }
+    })
 
     expect(store.get(c)).toBe('ab')
     expect(derivedStore.get(c)).toBe('ab')
@@ -124,46 +106,33 @@ describe('unstable_derive for scoping atoms', () => {
 
     function makeStores() {
       const store = createStore()
-      const derivedStore = store.unstable_derive(
-        (
-          getAtomState,
-          setAtomState,
-          atomRead,
-          atomWrite,
-          atomOnInit,
-          atomOnMount,
-        ) => {
-          const scopedAtomStateMap = new WeakMap()
-          return [
-            (atom) => {
-              if (scopedAtoms.has(atom)) {
-                return scopedAtomStateMap.get(atom)
-              }
-              return getAtomState(atom)
-            },
-            (atom, atomState) => {
-              if (scopedAtoms.has(atom)) {
+      const scopedAtomStateMap = new WeakMap()
+      const derivedStore = deriveStore(
+        store,
+        (ensureAtomState) => {
+          return (atom) => {
+            if (scopedAtoms.has(atom)) {
+              let atomState = scopedAtomStateMap.get(atom)
+              if (!atomState) {
+                atomState = { d: new Map(), p: new Set(), n: 0 }
                 scopedAtomStateMap.set(atom, atomState)
-              } else {
-                setAtomState(atom, atomState)
               }
               return atomState
-            },
-            (a, get, options) => {
-              const myGet: Getter = (aa) => {
-                if (scopedAtoms.has(aa)) {
-                  scopedAtoms.add(a) // Is this too naive?
-                }
-                return get(aa)
-              }
-              return atomRead(a, myGet, options)
-            },
-            atomWrite,
-            atomOnInit,
-            atomOnMount,
-          ]
+            }
+            return ensureAtomState(atom)
+          }
+        },
+        (a, get, options) => {
+          const myGet: Getter = (aa) => {
+            if (scopedAtoms.has(aa)) {
+              scopedAtoms.add(a) // Is this too naive?
+            }
+            return get(aa)
+          }
+          return a.read(myGet, options)
         },
       )
+
       expect(store.get(b)).toBe('a')
       expect(derivedStore.get(b)).toBe('a')
       return { store, derivedStore }
@@ -215,35 +184,18 @@ describe('unstable_derive for scoping atoms', () => {
 it('should pass the correct store instance to the atom initializer', () => {
   expect.assertions(2)
   const baseStore = createStore()
-  const derivedStore = baseStore.unstable_derive(
-    (
-      getAtomState,
-      setAtomState,
-      atomRead,
-      atomWrite,
-      atomOnInit,
-      atomOnMount,
-    ) => {
-      const initializedAtoms = new WeakSet()
-      return [
-        (atom) => {
-          if (!initializedAtoms.has(atom)) {
-            return undefined
-          }
-          return getAtomState(atom)
-        },
-        (atom, atomState) => {
-          initializedAtoms.add(atom)
-          setAtomState(atom, atomState)
-          return atomState
-        },
-        atomRead,
-        atomWrite,
-        atomOnInit,
-        atomOnMount,
-      ]
-    },
-  )
+  const initializedAtoms = new WeakSet()
+  const derivedStore = deriveStore(baseStore, (ensureAtomState) => {
+    return (atom) => {
+      if (!initializedAtoms.has(atom)) {
+        initializedAtoms.add(atom)
+        const atomState = ensureAtomState(atom, () => {})
+        atom.unstable_onInit?.(derivedStore)
+        return atomState
+      }
+      return ensureAtomState(atom)
+    }
+  })
   const a = atom(null)
   a.unstable_onInit = (store) => {
     expect(store).toBe(baseStore)

tests/vanilla/effect.test.ts

@@ -1,10 +1,13 @@
 import { expect, it, vi } from 'vitest'
 import type { Atom, Getter, Setter } from 'jotai/vanilla'
-import { atom, createStore } from 'jotai/vanilla'
+import {
+  atom,
+  createStore,
+  INTERNAL_getSecretStoreMethods as getSecretStoreMethods,
+} from 'jotai/vanilla'
 
 type Store = ReturnType<typeof createStore>
-type GetAtomState = Parameters<Parameters<Store['unstable_derive']>[0]>[0]
-type AtomState = NonNullable<ReturnType<GetAtomState>>
+type AtomState = ReturnType<ReturnType<typeof getSecretStoreMethods>[0]>
 type AnyAtom = Atom<unknown>
 type Batch = Parameters<NonNullable<AtomState['u']>>[0]
 
@@ -115,28 +118,9 @@ function ensureBatchChannel(batch: BatchWithSyncEffect) {
   return batch[syncEffectChannelSymbol]!
 }
 
-const getAtomStateMap = new WeakMap<Store, GetAtomState>()
-
-/**
- * HACK: steal atomState to synchronously determine if
- * the atom is mounted
- * We return null to cause the buildStore(...args) to throw
- * to abort creating a derived store
- */
 function getAtomState(store: Store, atom: AnyAtom): AtomState {
-  let getAtomStateFn = getAtomStateMap.get(store)
-  if (!getAtomStateFn) {
-    try {
-      store.unstable_derive((...storeArgs) => {
-        getAtomStateFn = storeArgs[0]
-        return null as any
-      })
-    } catch {
-      // expect error
-    }
-    getAtomStateMap.set(store, getAtomStateFn!)
-  }
-  return getAtomStateFn!(atom)!
+  const [ensureAtomState] = getSecretStoreMethods(store)
+  return ensureAtomState(atom)!
 }
 
 it('fires after recomputeDependents and before atom listeners', async function test() {