perf: avoid extra rerender

src/react/useAtomValue.ts

@@ -94,7 +94,9 @@ export function useAtomValue<Value>(atom: Atom<Value>, options?: Options) {
       }
       rerender()
     })
-    rerender()
+    if (!import.meta.env?.USE_STORE2) {
+      rerender()
+    }
     return unsub
   }, [store, atom, delay])
 

src/vanilla/store2.ts

@@ -132,6 +132,7 @@ type AtomState<Value = AnyValue> = {
    * The map value is the epoch number of the dependency.
    */
   readonly d: Map<AnyAtom, number>
+  readonly t: Array<WeakRef<AnyAtom>>
   /**
    * Set of atoms with pending promise that depend on the atom.
    *
@@ -146,6 +147,8 @@ type AtomState<Value = AnyValue> = {
   v?: Value
   /** Atom error */
   e?: AnyError
+  /** Should update */
+  u?: boolean // only if the atom is unmounted but has been written state, the value is true.
 }
 
 const isAtomStateInitialized = <Value>(atomState: AtomState<Value>) =>
@@ -235,7 +238,13 @@ const flushPending = (pending: Pending) => {
     pending[1].clear()
     const functions = new Set(pending[2])
     pending[2].clear()
-    atomStates.forEach((atomState) => atomState.m?.l.forEach((l) => l()))
+    atomStates.forEach((atomState) => {
+      if (atomState.m && atomState.m.l.size) {
+        atomState.m.l.forEach((l) => l())
+      } else {
+        atomState.u = true
+      }
+    })
     functions.forEach((fn) => fn())
   }
 }
@@ -259,13 +268,24 @@ type Store = PrdStore | (PrdStore & DevStoreRev4)
 export type INTERNAL_DevStoreRev4 = DevStoreRev4
 export type INTERNAL_PrdStore = PrdStore
 
+const addAtomRef = (atomState: AtomState, atom: AnyAtom) => {
+  if (atomState.m) {
+    atomState.m?.t.add(atom)
+  } else {
+    const index = atomState.t.findIndex((v) => v.deref() === atom)
+    if (index === -1) {
+      atomState.t.push(new WeakRef(atom))
+    }
+  }
+}
+
 export const createStore = (): Store => {
   const atomStateMap = new WeakMap<AnyAtom, AtomState>()
 
   const getAtomState = <Value>(atom: Atom<Value>) => {
     let atomState = atomStateMap.get(atom) as AtomState<Value> | undefined
     if (!atomState) {
-      atomState = { d: new Map(), p: new Set(), n: 0 }
+      atomState = { d: new Map(), p: new Set(), n: 0, t: [] }
       atomStateMap.set(atom, atomState)
     }
     return atomState
@@ -335,7 +355,7 @@ export const createStore = (): Store => {
     if (continuablePromise) {
       addPendingContinuablePromiseToDependency(atom, continuablePromise, aState)
     }
-    aState.m?.t.add(atom)
+    addAtomRef(aState, atom)
     if (pending) {
       addPendingDependent(pending, a, atom)
     }
@@ -456,7 +476,11 @@ export const createStore = (): Store => {
   const recomputeDependents = (pending: Pending, atom: AnyAtom) => {
     const getDependents = (a: AnyAtom): Set<AnyAtom> => {
       const aState = getAtomState(a)
-      const dependents = new Set(aState.m?.t)
+      const dependents: Set<AnyAtom> = new Set(
+        aState.m?.t.size
+          ? aState.m.t
+          : aState.t.map((v) => v.deref()!).filter(Boolean),
+      )
       for (const atomWithPendingContinuablePromise of aState.p) {
         dependents.add(atomWithPendingContinuablePromise)
       }
@@ -575,8 +599,9 @@ export const createStore = (): Store => {
     if (atomState.m && !getPendingContinuablePromise(atomState)) {
       for (const a of atomState.d.keys()) {
         if (!atomState.m.d.has(a)) {
-          const aMounted = mountAtom(pending, a)
-          aMounted.t.add(atom)
+          mountAtom(pending, a)
+          const aState = getAtomState(a)
+          addAtomRef(aState, atom)
           atomState.m.d.add(a)
         }
       }
@@ -597,8 +622,9 @@ export const createStore = (): Store => {
       readAtomState(pending, atom)
       // mount dependencies first
       for (const a of atomState.d.keys()) {
-        const aMounted = mountAtom(pending, a)
-        aMounted.t.add(atom)
+        mountAtom(pending, a)
+        const aState = getAtomState(a)
+        addAtomRef(aState, atom)
       }
       // mount self
       atomState.m = {
@@ -610,9 +636,10 @@ export const createStore = (): Store => {
         const mounted = atomState.m
         const { onMount } = atom
         addPendingFunction(pending, () => {
-          const onUnmount = onMount((...args) =>
-            writeAtomState(pending, atom, ...args),
-          )
+          const onUnmount = onMount((...args) => {
+            atomState.u = true
+            return writeAtomState(pending, atom, ...args)
+          })
           if (onUnmount) {
             mounted.u = onUnmount
           }
@@ -627,6 +654,7 @@ export const createStore = (): Store => {
     atom: AnyAtom,
   ): Mounted | undefined => {
     const atomState = getAtomState(atom)
+    atomState.t.length = 0
     if (
       atomState.m &&
       !atomState.m.l.size &&
@@ -660,6 +688,11 @@ export const createStore = (): Store => {
     flushPending(pending)
     const listeners = mounted.l
     listeners.add(listener)
+    const atomState = getAtomState(atom)
+    if (atomState.u) {
+      listener()
+      atomState.u = false
+    }
     return () => {
       listeners.delete(listener)
       const pending = createPending()

tests/react/optimization.test.tsx

@@ -271,3 +271,26 @@ it('no extra rerenders after commit with derived atoms (#1213)', async () => {
   })
   expect(renderCount1).toBe(renderCount1AfterCommit)
 })
+
+it.skipIf(!import.meta.env?.USE_STORE2)(
+  'no extra rerenders after mounted (#2558)',
+  async () => {
+    const baseAtom = atom(0)
+
+    let renderCount = 0
+
+    const App = () => {
+      const [count] = useAtom(baseAtom)
+      ++renderCount
+      return <div>count: {count}</div>
+    }
+
+    render(
+      <>
+        <App />
+      </>,
+    )
+
+    expect(renderCount).toBe(1)
+  },
+)

tests/vanilla/memoryleaks.test.ts

@@ -22,6 +22,9 @@ describe('test memory leaks (get & set only)', () => {
     const detector2 = new LeakDetector(store.get(derivedAtom))
     objAtom = undefined
     derivedAtom = undefined
+    // WeapRef is different from WeakMap, we should call isLeaking twice to make garbage-collected success.
+    await detector1.isLeaking()
+    await detector2.isLeaking()
     expect(await detector1.isLeaking()).toBe(false)
     expect(await detector2.isLeaking()).toBe(false)
   })
@@ -35,6 +38,7 @@ describe('test memory leaks (get & set only)', () => {
       const detector = new LeakDetector(depAtom)
       store.get(depAtom)
       depAtom = undefined
+      await detector.isLeaking()
       await expect(detector.isLeaking()).resolves.toBe(false)
     },
   )
@@ -49,6 +53,7 @@ describe('test memory leaks (get & set only)', () => {
       }))
       const detector = new LeakDetector(store.get(derivedAtom))
       derivedAtom = undefined
+      await detector.isLeaking()
       expect(await detector.isLeaking()).toBe(false)
     },
   )
@@ -63,6 +68,7 @@ describe('test memory leaks (with subscribe)', () => {
     unsub()
     unsub = undefined
     objAtom = undefined
+    await detector.isLeaking()
     expect(await detector.isLeaking()).toBe(false)
   })
 
@@ -79,6 +85,8 @@ describe('test memory leaks (with subscribe)', () => {
     unsub = undefined
     objAtom = undefined
     derivedAtom = undefined
+    await detector1.isLeaking()
+    await detector2.isLeaking()
     expect(await detector1.isLeaking()).toBe(false)
     expect(await detector2.isLeaking()).toBe(false)
   })
@@ -94,6 +102,7 @@ describe('test memory leaks (with subscribe)', () => {
     unsub()
     unsub = undefined
     derivedAtom = undefined
+    await detector.isLeaking()
     expect(await detector.isLeaking()).toBe(false)
   })
 })