compat: comply with new buildingBlocks api

package.json

@@ -85,7 +85,7 @@
     "eslint-plugin-prettier": "^5.2.3",
     "happy-dom": "^15.11.7",
     "jiti": "^2.4.2",
-    "jotai": "2.19.0",
+    "jotai": "https://pkg.csb.dev/pmndrs/jotai/commit/ac38969c/jotai",
     "jotai-effect": "link:",
     "prettier": "^3.4.2",
     "react": "19.2.1",
@@ -98,7 +98,7 @@
     "vitest": "^3.0.3"
   },
   "peerDependencies": {
-    "jotai": ">=2.16.0"
+    "jotai": ">=2.20.0"
   },
   "engines": {
     "node": ">=12.20.0"

src/atomEffect.ts

@@ -4,12 +4,12 @@ import type {
   INTERNAL_AtomState as AtomState,
   INTERNAL_MountedMap as MountedMap,
   INTERNAL_StoreHooks as StoreHooks,
-  INTERNAL_buildStoreRev2 as buildStore,
+  INTERNAL_buildStoreRev3 as buildStore,
 } from 'jotai/vanilla/internals'
 import {
-  INTERNAL_getBuildingBlocksRev2 as getBuildingBlocks,
+  INTERNAL_getBuildingBlocksRev3 as getBuildingBlocks,
   INTERNAL_hasInitialValue as hasInitialValue,
-  INTERNAL_initializeStoreHooksRev2 as initializeStoreHooks,
+  INTERNAL_initializeStoreHooksRev3 as initializeStoreHooks,
   INTERNAL_isAtomStateInitialized as isAtomStateInitialized,
   INTERNAL_returnAtomValue as returnAtomValue,
 } from 'jotai/vanilla/internals'
@@ -50,6 +50,19 @@ export function atomEffect(effect: Effect): Atom<void> & { effect: Effect } {
   effectAtom.effect = effect
 
   effectAtom.INTERNAL_onInit = (store) => {
+    const buildingBlocks = getBuildingBlocks(store)
+    const mountedMap = buildingBlocks[1]
+    const changedAtoms = buildingBlocks[3]
+    const storeHooks = initializeStoreHooks(buildingBlocks[6])
+    const ensureAtomState = buildingBlocks[11]
+    const flushCallbacks = buildingBlocks[12]
+    const recomputeInvalidatedAtoms = buildingBlocks[13]
+    const readAtomState = buildingBlocks[14]
+    const invalidateDependents = buildingBlocks[15]
+    const writeAtomState = buildingBlocks[16]
+    const mountDependencies = buildingBlocks[17]
+    const setAtomStateValueOrPromise = buildingBlocks[20]
+
     const deps = new Set<AnyAtom>()
     let inProgress = 0
     let isRecursing = false
@@ -69,10 +82,10 @@ export function atomEffect(effect: Effect): Atom<void> & { effect: Effect } {
           return store.get(a)
         }
         if (a === (effectAtom as AnyAtom)) {
-          const aState = ensureAtomState(store, a)
+          const aState = ensureAtomState(buildingBlocks, store, a)
           if (!isAtomStateInitialized(aState)) {
             if (hasInitialValue(a)) {
-              setAtomStateValueOrPromise(store, a, a.init)
+              setAtomStateValueOrPromise(buildingBlocks, store, a, a.init)
             } else {
               // NOTE invalid derived atoms can reach here
               throw new Error('no atom init')
@@ -81,7 +94,7 @@ export function atomEffect(effect: Effect): Atom<void> & { effect: Effect } {
           return returnAtomValue(aState)
         }
         // a !== atom
-        const aState = readAtomState(store, a)
+        const aState = readAtomState(buildingBlocks, store, a)
         try {
           return returnAtomValue(aState)
         } finally {
@@ -91,9 +104,9 @@ export function atomEffect(effect: Effect): Atom<void> & { effect: Effect } {
             deps.add(a)
           } else {
             if (mountedMap.has(a)) {
-              mountDependencies(store, effectAtom)
-              recomputeInvalidatedAtoms(store)
-              flushCallbacks(store)
+              mountDependencies(buildingBlocks, store, effectAtom)
+              recomputeInvalidatedAtoms(buildingBlocks, store)
+              flushCallbacks(buildingBlocks, store)
             }
           }
         }
@@ -105,7 +118,7 @@ export function atomEffect(effect: Effect): Atom<void> & { effect: Effect } {
         a: WritableAtom<V, As, R>,
         ...args: As
       ) => {
-        const aState = ensureAtomState(store, a)
+        const aState = ensureAtomState(buildingBlocks, store, a)
         try {
           ++inProgress
           if (a === (effectAtom as AnyAtom)) {
@@ -115,21 +128,21 @@ export function atomEffect(effect: Effect): Atom<void> & { effect: Effect } {
             }
             const prevEpochNumber = aState.n
             const v = args[0] as V
-            setAtomStateValueOrPromise(store, a, v)
-            mountDependencies(store, a)
+            setAtomStateValueOrPromise(buildingBlocks, store, a, v)
+            mountDependencies(buildingBlocks, store, a)
             if (prevEpochNumber !== aState.n) {
               changedAtoms.add(a)
               storeHooks.c?.(a)
-              invalidateDependents(store, a)
+              invalidateDependents(buildingBlocks, store, a)
             }
             return undefined as R
           } else {
-            return writeAtomState(store, a, ...args)
+            return writeAtomState(buildingBlocks, store, a, args)
           }
         } finally {
           if (!isSync) {
-            recomputeInvalidatedAtoms(store)
-            flushCallbacks(store)
+            recomputeInvalidatedAtoms(buildingBlocks, store)
+            flushCallbacks(buildingBlocks, store)
           }
           --inProgress
         }
@@ -144,10 +157,10 @@ export function atomEffect(effect: Effect): Atom<void> & { effect: Effect } {
         }
         try {
           isRecursing = true
-          mountDependencies(store, effectAtom)
+          mountDependencies(buildingBlocks, store, effectAtom)
           return setter(a, ...args)
         } finally {
-          recomputeInvalidatedAtoms(store)
+          recomputeInvalidatedAtoms(buildingBlocks, store)
           isRecursing = false
           if (hasChanged) {
             hasChanged = false
@@ -179,28 +192,18 @@ export function atomEffect(effect: Effect): Atom<void> & { effect: Effect } {
       } finally {
         isSync = false
         deps.forEach((depAtom) => {
-          atomState.d.set(depAtom, ensureAtomState(store, depAtom).n)
+          atomState.d.set(
+            depAtom,
+            ensureAtomState(buildingBlocks, store, depAtom).n
+          )
         })
-        mountDependencies(store, effectAtom)
-        recomputeInvalidatedAtoms(store)
+        mountDependencies(buildingBlocks, store, effectAtom)
+        recomputeInvalidatedAtoms(buildingBlocks, store)
       }
     }
 
-    const buildingBlocks = getBuildingBlocks(store)
-    const mountedMap = buildingBlocks[1]
-    const changedAtoms = buildingBlocks[3]
-    const storeHooks = initializeStoreHooks(buildingBlocks[6])
-    const ensureAtomState = buildingBlocks[11]
-    const flushCallbacks = buildingBlocks[12]
-    const recomputeInvalidatedAtoms = buildingBlocks[13]
-    const readAtomState = buildingBlocks[14]
-    const invalidateDependents = buildingBlocks[15]
-    const writeAtomState = buildingBlocks[16]
-    const mountDependencies = buildingBlocks[17]
-    const setAtomStateValueOrPromise = buildingBlocks[20]
-
     const atomEffectChannel = ensureAtomEffectChannel(store, storeHooks)
-    const atomState = ensureAtomState(store, effectAtom)
+    const atomState = ensureAtomState(buildingBlocks, store, effectAtom)
     // initialize atomState
     atomState.v = undefined
 

src/withAtomEffect.ts

@@ -1,7 +1,7 @@
 import type { Atom, WritableAtom } from 'jotai/vanilla'
 import {
-  INTERNAL_getBuildingBlocksRev2 as getBuildingBlocks,
-  INTERNAL_initializeStoreHooksRev2 as initializeStoreHooks,
+  INTERNAL_getBuildingBlocksRev3 as getBuildingBlocks,
+  INTERNAL_initializeStoreHooksRev3 as initializeStoreHooks,
 } from 'jotai/vanilla/internals'
 import type { Effect, GetterWithPeek, SetterWithRecurse } from './atomEffect'
 import { atomEffect } from './atomEffect'
@@ -82,39 +82,43 @@ export function withAtomEffect<T extends Atom<unknown>>(
       })
       effectAtom.debugPrivate = true
     }
-    const effectAtomState = ensureAtomState(store, effectAtom)
-    const targetWithEffectAtomState = ensureAtomState(store, targetWithEffect)
+    const effectAtomState = ensureAtomState(buildingBlocks, store, effectAtom)
+    const targetWithEffectAtomState = ensureAtomState(
+      buildingBlocks,
+      store,
+      targetWithEffect
+    )
 
     storeHooks.c.add(targetWithEffect, function atomChanged() {
       if (isSubscribed) {
         invalidatedAtoms.set(effectAtom, effectAtomState.n)
         effectAtomState.d.set(targetWithEffect, targetWithEffectAtomState.n - 1)
-        readAtomState(store, effectAtom)
-        mountDependencies(store, effectAtom)
+        readAtomState(buildingBlocks, store, effectAtom)
+        mountDependencies(buildingBlocks, store, effectAtom)
         invalidatedAtoms.delete(effectAtom)
         effectAtomState.d.delete(targetWithEffect)
       }
     })
     storeHooks.m.add(targetWithEffect, function mountEffect() {
-      const atomState = ensureAtomState(store, targetWithEffect)
+      const atomState = ensureAtomState(buildingBlocks, store, targetWithEffect)
       const { n } = atomState
       // Defer effect mount to the next flush `f` so nested mount waves do not replace mounted maps
       // after inner passes have populated mounted.t, which can strand invalidation edges (#3292).
       const unsubFlush = storeHooks.f.add(() => {
         unsubFlush()
-        mountAtom(store, effectAtom)
+        mountAtom(buildingBlocks, store, effectAtom)
         if (n !== atomState.n) {
           const unsubPost = storeHooks.f.add(() => {
             unsubPost()
-            invalidateDependents(store, targetWithEffect)
+            invalidateDependents(buildingBlocks, store, targetWithEffect)
           })
         }
-        flushCallbacks(store)
+        flushCallbacks(buildingBlocks, store)
       })
     })
     storeHooks.u.add(targetWithEffect, function unmountEffect() {
-      unmountAtom(store, effectAtom)
-      flushCallbacks(store)
+      unmountAtom(buildingBlocks, store, effectAtom)
+      flushCallbacks(buildingBlocks, store)
     })
     storeHooks.f.add(function flushEffect() {
       inProgress = false

tests/atomEffect.test.ts

@@ -4,8 +4,8 @@ import { act, render } from '@testing-library/react'
 import { Provider, useAtomValue } from 'jotai/react'
 import { atom, createStore } from 'jotai/vanilla'
 import {
-  INTERNAL_getBuildingBlocksRev2 as getBuildingBlocks,
-  INTERNAL_initializeStoreHooksRev2 as initializeStoreHooks,
+  INTERNAL_getBuildingBlocksRev3 as getBuildingBlocks,
+  INTERNAL_initializeStoreHooksRev3 as initializeStoreHooks,
 } from 'jotai/vanilla/internals'
 import { describe, expect, it, vi } from 'vitest'
 import { atomEffect } from '../src/atomEffect'
@@ -1009,27 +1009,27 @@ it('gets the right internals from the store', function test() {
   ) // storeHooks
   expect(buildingBlocks[11].name.endsWith('ensureAtomState')).toBe(true)
   expect(buildingBlocks[11]).toBeInstanceOf(Function)
-  expect(buildingBlocks[11]).toHaveLength(2)
+  expect(buildingBlocks[11]).toHaveLength(3)
   expect(buildingBlocks[12].name.endsWith('flushCallbacks')).toBe(true)
   expect(buildingBlocks[12]).toBeInstanceOf(Function)
-  expect(buildingBlocks[12]).toHaveLength(1)
+  expect(buildingBlocks[12]).toHaveLength(2)
   expect(buildingBlocks[13].name.endsWith('recomputeInvalidatedAtoms')).toBe(
     true
   )
   expect(buildingBlocks[13]).toBeInstanceOf(Function)
-  expect(buildingBlocks[13]).toHaveLength(1)
+  expect(buildingBlocks[13]).toHaveLength(2)
   expect(buildingBlocks[14].name.endsWith('readAtomState')).toBe(true)
   expect(buildingBlocks[14]).toBeInstanceOf(Function)
-  expect(buildingBlocks[14]).toHaveLength(2)
+  expect(buildingBlocks[14]).toHaveLength(3)
   expect(buildingBlocks[15].name.endsWith('invalidateDependents')).toBe(true)
   expect(buildingBlocks[15]).toBeInstanceOf(Function)
-  expect(buildingBlocks[15]).toHaveLength(2)
+  expect(buildingBlocks[15]).toHaveLength(3)
   expect(buildingBlocks[16].name.endsWith('writeAtomState')).toBe(true)
   expect(buildingBlocks[16]).toBeInstanceOf(Function)
-  expect(buildingBlocks[16]).toHaveLength(2)
+  expect(buildingBlocks[16]).toHaveLength(4)
   expect(buildingBlocks[17].name.endsWith('mountDependencies')).toBe(true)
   expect(buildingBlocks[17]).toBeInstanceOf(Function)
-  expect(buildingBlocks[17]).toHaveLength(2)
+  expect(buildingBlocks[17]).toHaveLength(3)
 })
 
 it('should not run the effect when the effectAtom is unmounted', function test() {

tests/test-utils.ts

@@ -3,9 +3,9 @@ import type {
   INTERNAL_Store as Store,
 } from 'jotai/vanilla/internals'
 import {
-  INTERNAL_buildStoreRev2 as buildStore,
-  INTERNAL_getBuildingBlocksRev2 as getBuildingBlocks,
-  INTERNAL_initializeStoreHooksRev2 as initializeStoreHooks,
+  INTERNAL_buildStoreRev3 as buildStore,
+  INTERNAL_getBuildingBlocksRev3 as getBuildingBlocks,
+  INTERNAL_initializeStoreHooksRev3 as initializeStoreHooks,
 } from 'jotai/vanilla/internals'
 
 //
@@ -35,8 +35,8 @@ export function createDebugStore(
 ): DebugStore {
   const buildingBlocks: BuildingBlocks = [...getBuildingBlocks(buildStore())]
   const ensureAtomState = buildingBlocks[11]
-  buildingBlocks[11] = (store, atom) =>
-    Object.assign(ensureAtomState(store, atom), { label: atom.debugLabel })
+  buildingBlocks[11] = (bb, store, atom) =>
+    Object.assign(ensureAtomState(bb, store, atom), { label: atom.debugLabel })
   const debugStore = buildStore(...buildingBlocks) as DebugStore
   debugStore.name = name
   debugStore.state = {