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]>

Author: arjunvegda

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.

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
   }

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()
+  })
 })