expose building blocks

pmndrs · Feb 4, 2025 · 81fb3b3 · 81fb3b3
1 parent 80f0472
commit 81fb3b3
Show file tree

Hide file tree

Showing 5 changed files with 31 additions and 14 deletions.
diff --git a/.github/workflows/test-multiple-builds.yml b/.github/workflows/test-multiple-builds.yml
@@ -43,7 +43,7 @@ jobs:
           sed -i~ "s/resolve('\.\/src\(.*\)\.ts')/resolve('\.\/dist\1.js')/" vitest.config.mts
           sed -i~ "s/import { useResetAtom } from 'jotai\/react\/utils'/const { useResetAtom } = require('..\/..\/..\/dist\/react\/utils.js')/" tests/react/utils/useResetAtom.test.tsx
           sed -i~ "s/import { RESET, atomWithReducer, atomWithReset } from 'jotai\/vanilla\/utils'/const { RESET, atomWithReducer, atomWithReset } = require('..\/..\/..\/dist\/vanilla\/utils.js')/" tests/react/utils/useResetAtom.test.tsx
-          perl -i~ -0777 -pe "s/import {[^}]+} from 'jotai\/vanilla\/internals'/const { INTERNAL_buildStore, INTERNAL_createStoreArgs, INTERNAL_initializeStoreHooks, INTERNAL_getStoreArgsRev1: INTERNAL_getStoreArgs } = require('..\/..\/dist\/vanilla\/internals.js')/g" tests/vanilla/store.test.tsx tests/vanilla/derive.test.tsx tests/vanilla/effect.test.ts
+          perl -i~ -0777 -pe "s/import {[^}]+} from 'jotai\/vanilla\/internals'/const { INTERNAL_buildStore, INTERNAL_createStoreArgs, INTERNAL_initializeStoreHooks, INTERNAL_getStoreArgsRev1: INTERNAL_getStoreArgs, INTERNAL_createStoreArgsRev1: INTERNAL_createStoreArgs } = require('..\/..\/dist\/vanilla\/internals.js')/g" tests/vanilla/store.test.tsx tests/vanilla/derive.test.tsx tests/vanilla/effect.test.ts
       - name: Patch for ESM
         if: ${{ matrix.build == 'esm' }}
         run: |

diff --git a/src/vanilla/internals.ts b/src/vanilla/internals.ts
@@ -822,16 +822,17 @@ const createStoreArgs: (...storeArgs: Partial<StoreArgs>) => StoreArgs = (
   atomOnMount,
 ]
 
-const buildStore = (...storeArgs: StoreArgs): Store => {
-  const [
+const buildStore = (
+  storeArgs: StoreArgs,
+  [
     flushCallbacks,
     recomputeInvalidatedAtoms,
     readAtomState,
     writeAtomState,
     mountAtom,
     unmountAtom,
-  ] = createBuildingBlocks(storeArgs, () => store as Store)
-
+  ]: BuildingBlocks,
+): Store => {
   const readAtom = <Value>(atom: Atom<Value>): Value =>
     returnAtomValue(readAtomState(atom))
 
@@ -872,11 +873,12 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
 // Export internal functions
 //
 
-export const INTERNAL_createStoreArgs: typeof createStoreArgs = createStoreArgs
 export const INTERNAL_buildStore: typeof buildStore = buildStore
-export const INTERNAL_getStoreArgsRev1: typeof getStoreArgs = getStoreArgs
 export const INTERNAL_createBuildingBlocksRev1: typeof createBuildingBlocks =
   createBuildingBlocks
+export const INTERNAL_createStoreArgsRev1: typeof createStoreArgs =
+  createStoreArgs
+export const INTERNAL_getStoreArgsRev1: typeof getStoreArgs = getStoreArgs
 export const INTERNAL_initializeStoreHooks: typeof initializeStoreHooks =
   initializeStoreHooks
 

diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts
@@ -1,7 +1,8 @@
 import type { Atom, WritableAtom } from './atom.ts'
 import {
   INTERNAL_buildStore,
-  INTERNAL_createStoreArgs,
+  INTERNAL_createBuildingBlocksRev1 as INTERNAL_createBuildingBlocks,
+  INTERNAL_createStoreArgsRev1 as INTERNAL_createStoreArgs,
   INTERNAL_initializeStoreHooks,
 } from './internals.ts'
 import type { INTERNAL_AtomState } from './internals.ts'
@@ -33,7 +34,7 @@ const createDevStoreRev4 = (): INTERNAL_PrdStore & INTERNAL_DevStoreRev4 => {
   const storeHooks = INTERNAL_initializeStoreHooks({})
   const atomStateMap = new WeakMap()
   const mountedAtoms = new WeakMap()
-  const store = INTERNAL_buildStore(
+  const storeArgs = INTERNAL_createStoreArgs(
     atomStateMap,
     mountedAtoms,
     new WeakMap(),
@@ -51,6 +52,8 @@ const createDevStoreRev4 = (): INTERNAL_PrdStore & INTERNAL_DevStoreRev4 => {
     (atom, ...params) => atom.unstable_onInit?.(...params),
     (atom, ...params) => atom.onMount?.(...params),
   )
+  const buildingBlocks = INTERNAL_createBuildingBlocks(storeArgs, () => store)
+  const store = INTERNAL_buildStore(storeArgs, buildingBlocks)
   const debugMountedAtoms = new Set<Atom<unknown>>()
   storeHooks.m.add(undefined, (atom) => {
     debugMountedAtoms.add(atom)
@@ -98,7 +101,9 @@ export const createStore = (): PrdOrDevStore => {
   if (import.meta.env?.MODE !== 'production') {
     return createDevStoreRev4()
   }
-  const store = INTERNAL_buildStore(...INTERNAL_createStoreArgs())
+  const storeArgs = INTERNAL_createStoreArgs()
+  const buildingBlocks = INTERNAL_createBuildingBlocks(storeArgs, () => store)
+  const store = INTERNAL_buildStore(storeArgs, buildingBlocks)
   return store
 }
 

diff --git a/tests/vanilla/derive.test.tsx b/tests/vanilla/derive.test.tsx
@@ -3,7 +3,8 @@ import { atom, createStore } from 'jotai/vanilla'
 import type { Atom } from 'jotai/vanilla'
 import {
   INTERNAL_buildStore,
-  INTERNAL_createStoreArgs,
+  INTERNAL_createBuildingBlocksRev1 as INTERNAL_createBuildingBlocks,
+  INTERNAL_createStoreArgsRev1 as INTERNAL_createStoreArgs,
   INTERNAL_getStoreArgsRev1 as INTERNAL_getStoreArgs,
 } from 'jotai/vanilla/internals'
 
@@ -18,7 +19,11 @@ const deriveStore = (
   const newStoreArgs = INTERNAL_createStoreArgs(
     enhanceAtomStateMap(atomStateMap),
   )
-  const derivedStore = (INTERNAL_buildStore as any)(...newStoreArgs)
+  const buildingBlocks = INTERNAL_createBuildingBlocks(
+    newStoreArgs,
+    () => derivedStore,
+  )
+  const derivedStore = INTERNAL_buildStore(newStoreArgs, buildingBlocks)
   return derivedStore
 }
 

diff --git a/tests/vanilla/store.test.tsx b/tests/vanilla/store.test.tsx
@@ -4,7 +4,8 @@ import { atom, createStore } from 'jotai/vanilla'
 import type { Atom, Getter, PrimitiveAtom } from 'jotai/vanilla'
 import {
   INTERNAL_buildStore,
-  INTERNAL_createStoreArgs,
+  INTERNAL_createBuildingBlocksRev1 as INTERNAL_createBuildingBlocks,
+  INTERNAL_createStoreArgsRev1 as INTERNAL_createStoreArgs,
   INTERNAL_getStoreArgsRev1 as INTERNAL_getStoreArgs,
 } from 'jotai/vanilla/internals'
 
@@ -19,7 +20,11 @@ const deriveStore = (
   const newStoreArgs = INTERNAL_createStoreArgs(
     enhanceAtomStateMap(atomStateMap),
   )
-  const derivedStore = (INTERNAL_buildStore as any)(...newStoreArgs)
+  const buildingBlocks = INTERNAL_createBuildingBlocks(
+    newStoreArgs,
+    () => derivedStore,
+  )
+  const derivedStore = INTERNAL_buildStore(newStoreArgs, buildingBlocks)
   return derivedStore
 }