From 8c9258fe2726a9e735d813763cdb41f4b1e0ca94 Mon Sep 17 00:00:00 2001
From: David Maskasky <dmaskasky@gmail.com>
Date: Sun, 17 Nov 2024 02:37:20 -0800
Subject: [PATCH] fix: dedupe subscribers in flushPending

---
 src/vanilla/store.ts              |  9 +++++++--
 tests/vanilla/dependency.test.tsx | 24 ++++++++++++++++++++++--
 2 files changed, 29 insertions(+), 4 deletions(-)

diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts
index 697e52e830..c54cb210e7 100644
--- a/src/vanilla/store.ts
+++ b/src/vanilla/store.ts
@@ -213,11 +213,16 @@ const flushPending = (pending: Pending) => {
   }
   while (pending[1].size || pending[2].size) {
     pending[0].clear()
-    const atomStates = new Set(pending[1].values())
+    const atomListeners = new Set<() => void>()
+    for (const aState of pending[1].values()) {
+      if (aState.m) {
+        aState.m.l.forEach((listener) => atomListeners.add(listener))
+      }
+    }
     pending[1].clear()
     const functions = new Set(pending[2])
     pending[2].clear()
-    atomStates.forEach((atomState) => atomState.m?.l.forEach(call))
+    atomListeners.forEach(call)
     functions.forEach(call)
   }
   if (hasError) {
diff --git a/tests/vanilla/dependency.test.tsx b/tests/vanilla/dependency.test.tsx
index 98e0f18f84..31a4146f7e 100644
--- a/tests/vanilla/dependency.test.tsx
+++ b/tests/vanilla/dependency.test.tsx
@@ -1,5 +1,9 @@
 import { expect, it, vi } from 'vitest'
 import { atom, createStore } from 'jotai/vanilla'
+import type {
+  INTERNAL_DevStoreRev4,
+  INTERNAL_PrdStore,
+} from 'jotai/vanilla/store'
 
 it('can propagate updates with async atom chains', async () => {
   const store = createStore()
@@ -349,16 +353,32 @@ it('can read sync derived atom in write without initializing', () => {
 
 it.only('batches sync writes', () => {
   const a = atom(1)
-  const b = atom(2)
+  a.debugLabel = 'a'
+  const b = atom(1)
+  b.debugLabel = 'b'
   const fetch = vi.fn((va, vb) => va * 10 + vb)
   const c = atom((get) => fetch(get(a), get(b)))
+  c.debugLabel = 'c'
   const d = atom(null, (_, set) => {
     set(a, (v) => ++v)
     set(b, (v) => ++v)
   })
-  const store = createStore()
+  d.debugLabel = 'd'
+  const store = createStore() as INTERNAL_DevStoreRev4 & INTERNAL_PrdStore
+  const getAtomState = store.dev4_get_internal_weak_map().get
+  store.sub(a, () => {})
+  store.sub(b, () => {})
   const cListener = vi.fn()
   store.sub(c, cListener)
+  store.sub(d, () => {})
+  const aState = getAtomState(a) as any
+  aState.label = 'a'
+  const bState = getAtomState(b) as any
+  bState.label = 'b'
+  const cState = getAtomState(c) as any
+  cState.label = 'c'
+  const dState = getAtomState(d) as any
+  dState.label = 'd'
   cListener.mockClear()
   fetch.mockClear()
   store.set(d)