Merge branch 'main' into fix/bug-2554

docs/guides/nextjs.mdx

@@ -87,11 +87,15 @@ Jotai provides SWC plugins for better DX while developing with Next.js. [Find mo
 
 #### HN Posts
 
+Page Router demo:
+
 <Stackblitz
   id="stackblitz-starters-cnz9lg"
   file="store%2Findex.ts,pages%2F_app.tsx,pages%2Findex.tsx"
 />
 
+App Router [demo on Stackblitz](https://stackblitz.com/edit/jotai-nextjs-app-router-demo?file=store%2Findex.ts,app%2Flayout.tsx,components%2FPost.tsx,app%2Fpage.tsx)
+
 #### Next.js repo
 
 ```bash

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.
@@ -400,11 +397,11 @@ export const createStore = (): Store => {
 
   const readAtomState = <Value>(
     atom: Atom<Value>,
-    force?: boolean,
+    force?: (a: AnyAtom) => boolean,
   ): AtomState<Value> => {
     // See if we can skip recomputing this atom.
     const atomState = getAtomState(atom)
-    if (!force && atomState) {
+    if (!force?.(atom) && atomState) {
       // If the atom is mounted, we can use the cache.
       // because it should have been updated by dependencies.
       if (mountedMap.has(atom)) {
@@ -418,7 +415,7 @@ export const createStore = (): Store => {
           if (a === atom) {
             return true
           }
-          const aState = readAtomState(a)
+          const aState = readAtomState(a, force)
           // Check if the atom state is unchanged, or
           // check the atom value in case only dependencies are changed
           return aState === s || isEqualAtomValue(aState, s)
@@ -445,7 +442,7 @@ export const createStore = (): Store => {
         throw new Error('no atom init')
       }
       // a !== atom
-      const aState = readAtomState(a)
+      const aState = readAtomState(a, force)
       nextDependencies.set(a, aState)
       return returnAtomValue(aState)
     }
@@ -534,6 +531,7 @@ export const createStore = (): Store => {
     // Step 2: use the topsorted atom list to recompute all affected atoms
     // Track what's changed, so that we can short circuit when possible
     const changedAtoms = new Set<AnyAtom>([atom])
+    const isMarked = (a: AnyAtom) => markedAtoms.has(a)
     for (let i = topsortedAtoms.length - 1; i >= 0; --i) {
       const a = topsortedAtoms[i]!
       const prevAtomState = getAtomState(a)
@@ -548,12 +546,13 @@ export const createStore = (): Store => {
         }
       }
       if (hasChangedDeps) {
-        const nextAtomState = readAtomState(a, true)
+        const nextAtomState = readAtomState(a, isMarked)
         addPendingDependent(a, nextAtomState)
         if (!isEqualAtomValue(prevAtomState, nextAtomState)) {
           changedAtoms.add(a)
         }
       }
+      markedAtoms.delete(a)
     }
   }
 

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 = {
@@ -346,11 +347,11 @@ export const createStore = (): Store => {
   const readAtomState = <Value>(
     pending: Pending | undefined,
     atom: Atom<Value>,
-    force?: true,
+    force?: (a: AnyAtom) => boolean,
   ): AtomState<Value> => {
     // See if we can skip recomputing this atom.
     const atomState = getAtomState(atom)
-    if (!force && isAtomStateInitialized(atomState)) {
+    if (!force?.(atom) && isAtomStateInitialized(atomState)) {
       // If the atom is mounted, we can use the cache.
       // because it should have been updated by dependencies.
       if (atomState.m) {
@@ -363,7 +364,7 @@ export const createStore = (): Store => {
           ([a, n]) =>
             // Recursively, read the atom state of the dependency, and
             // check if the atom epoch number is unchanged
-            readAtomState(pending, a).n === n,
+            readAtomState(pending, a, force).n === n,
         )
       ) {
         return atomState
@@ -386,7 +387,7 @@ export const createStore = (): Store => {
         return returnAtomValue(aState)
       }
       // a !== atom
-      const aState = readAtomState(pending, a)
+      const aState = readAtomState(pending, a, force)
       if (isSync) {
         addDependency(pending, atom, a, aState)
       } else {
@@ -498,6 +499,7 @@ export const createStore = (): Store => {
     // Step 2: use the topsorted atom list to recompute all affected atoms
     // Track what's changed, so that we can short circuit when possible
     const changedAtoms = new Set<AnyAtom>([atom])
+    const isMarked = (a: AnyAtom) => markedAtoms.has(a)
     for (let i = topsortedAtoms.length - 1; i >= 0; --i) {
       const a = topsortedAtoms[i]!
       const aState = getAtomState(a)
@@ -511,13 +513,14 @@ export const createStore = (): Store => {
         }
       }
       if (hasChangedDeps) {
-        readAtomState(pending, a, true)
+        readAtomState(pending, a, isMarked)
         mountDependencies(pending, a, aState)
         if (!hasPrevValue || !Object.is(prevValue, aState.v)) {
           addPendingAtom(pending, a, aState)
           changedAtoms.add(a)
         }
       }
+      markedAtoms.delete(a)
     }
   }
 
@@ -679,6 +682,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/dependency.test.tsx

@@ -153,3 +153,34 @@ it('keeps atoms mounted between recalculations', async () => {
     unmounted: 0,
   })
 })
+
+it('should not provide stale values to conditional dependents', () => {
+  const dataAtom = atom<number[]>([100])
+  const hasFilterAtom = atom(false)
+  const filteredAtom = atom((get) => {
+    const data = get(dataAtom)
+    const hasFilter = get(hasFilterAtom)
+    if (hasFilter) {
+      return []
+    } else {
+      return data
+    }
+  })
+  const stageAtom = atom((get) => {
+    const hasFilter = get(hasFilterAtom)
+    if (hasFilter) {
+      const filtered = get(filteredAtom)
+      return filtered.length === 0 ? 'is-empty' : 'has-data'
+    } else {
+      return 'no-filter'
+    }
+  })
+
+  const store = createStore()
+  store.sub(filteredAtom, () => undefined)
+  store.sub(stageAtom, () => undefined)
+
+  expect(store.get(stageAtom), 'should start without filter').toBe('no-filter')
+  store.set(hasFilterAtom, true)
+  expect(store.get(stageAtom), 'should update').toBe('is-empty')
+})

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