feat: add back dev4_restore_atoms to store v2 (#2509)

* feat: add back dev4_restore_atoms to store v2 * refactor: remove internal types for store1 * fix: add primitive atom check when restoring Co-authored-by: Daishi Kato <[email protected]> * refactor: adjust to new changes * test: add failing test for store.sub not being called afer restore * test: simplify failing test * fix: add atoms to pending during restore * Update src/vanilla/store2.ts --------- Co-authored-by: Daishi Kato <[email protected]>
pmndrs · May 21, 2024 · a685dfc · a685dfc
1 parent c9b504f
commit a685dfc
Show file tree

Hide file tree

Showing 3 changed files with 54 additions and 3 deletions.
diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts
@@ -153,9 +153,6 @@ type PrdStore = {
 }
 type Store = PrdStore & Partial<DevStoreRev2>
 
-export type INTERNAL_DevStoreRev2 = DevStoreRev2
-export type INTERNAL_PrdStore = PrdStore
-
 /**
  * Create a new store. Each store is an independent, isolated universe of atom
  * states.

diff --git a/src/vanilla/store2.ts b/src/vanilla/store2.ts
@@ -247,6 +247,7 @@ type DevStoreRev4 = {
  key: K,
  fn: PrdStore[K],
  ) => void
+ dev4_restore_atoms: (values: Iterable<readonly [AnyAtom, AnyValue]>) => void
 }
 
 type PrdStore = {
@@ -679,6 +680,23 @@ export const createStore = (): Store => {
  dev4_override_method: (key, fn) => {
  ;(store as any)[key] = fn
  },
+ dev4_restore_atoms: (values) => {
+ const pending = createPending()
+ for (const [atom, value] of values) {
+ if (hasInitialValue(atom)) {
+ const aState = getAtomState(atom)
+ const hasPrevValue = 'v' in aState
+ const prevValue = aState.v
+ setAtomStateValueOrPromise(atom, aState, value)
+ mountDependencies(pending, atom, aState)
+ if (!hasPrevValue || !Object.is(prevValue, aState.v)) {
+ addPendingAtom(pending, atom, aState)
+ recomputeDependents(pending, atom)
+ }
+ }
+ }
+ flushPending(pending)
+ },
  }
  return store
  }

diff --git a/tests/vanilla/storedev.test.tsx b/tests/vanilla/storedev.test.tsx
@@ -136,4 +136,40 @@ describe.skipIf(!IS_DEV_STORE2)('[DEV-ONLY] dev-only methods rev4', () => {
  const weakMap = store.dev4_get_internal_weak_map()
  expect(weakMap.get(countAtom)?.v).toEqual(1)
  })
+
+ it('should restore atoms and its dependencies correctly', () => {
+ const store = createStore() as any
+ if (!('dev4_restore_atoms' in store)) {
+ throw new Error('dev methods are not available')
+ }
+ const countAtom = atom(0)
+ const derivedAtom = atom((get) => get(countAtom) * 2)
+ store.set(countAtom, 1)
+ store.dev4_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 any
+ if (!('dev4_restore_atoms' in store)) {
+ throw new Error('dev methods are not available')
+ }
+ 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([
+ [countAtom, 1],
+ [derivedAtom, 2],
+ ])
+
+ expect(countCb).toHaveBeenCalled()
+ expect(derivedCb).toHaveBeenCalled()
+ unsubCount()
+ unsubDerived()
+ })
 })