From f6be2e06d380a1d6bd8726f047ad3d1aa3d614d9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Maciej=20Kwa=C5=9Bniak?= <contact@exelord.com>
Date: Sun, 8 Sep 2024 23:24:31 +0200
Subject: [PATCH] fix: Don't assume iterators will always break on specific
 iteration

---
 packages/set/src/index.ts       | 32 ++++++++++++++++++--------------
 packages/set/test/index.test.ts | 26 ++++++++++++++++----------
 2 files changed, 34 insertions(+), 24 deletions(-)

diff --git a/packages/set/src/index.ts b/packages/set/src/index.ts
index 6dcb7d8d4..0a9dd191a 100644
--- a/packages/set/src/index.ts
+++ b/packages/set/src/index.ts
@@ -28,35 +28,39 @@ export class ReactiveSet<T> extends Set<T> {
     this.#triggers.track($KEYS);
     return super.size;
   }
+
   has(v: T): boolean {
     this.#triggers.track(v);
     return super.has(v);
   }
 
-  *keys(): IterableIterator<T> {
-    for (const key of super.keys()) {
-      this.#triggers.track(key);
-      yield key;
-    }
-    this.#triggers.track($KEYS);
+  keys(): IterableIterator<T> {
+    return this.values();
   }
-  values(): IterableIterator<T> {
-    return this.keys();
+
+  *values(): IterableIterator<T> {
+    this.#triggers.track($KEYS);
+
+    for (const value of super.values()) {
+      yield value;
+    }
   }
+
   *entries(): IterableIterator<[T, T]> {
-    for (const key of super.keys()) {
-      this.#triggers.track(key);
-      yield [key, key];
-    }
     this.#triggers.track($KEYS);
+
+    for (const entry of super.entries()) {
+      yield entry;
+    }
   }
 
   [Symbol.iterator](): IterableIterator<T> {
     return this.values();
   }
-  forEach(callbackfn: (value: T, value2: T, set: this) => void) {
+
+  forEach(callbackfn: (value1: T, value2: T, set: Set<T>) => void, thisArg?: any): void {
     this.#triggers.track($KEYS);
-    super.forEach(callbackfn as any);
+    super.forEach(callbackfn, thisArg);
   }
 
   // writes
diff --git a/packages/set/test/index.test.ts b/packages/set/test/index.test.ts
index 8d97062e5..412fba986 100644
--- a/packages/set/test/index.test.ts
+++ b/packages/set/test/index.test.ts
@@ -100,11 +100,8 @@ describe("ReactiveSet", () => {
       const dispose = createRoot(dispose => {
         createEffect(() => {
           const run: number[] = [];
-          let i = 0;
           for (const key of fn(set)) {
             run.push(key);
-            if (i === 2) break; // don't iterate over all keys
-            i += 1;
           }
           captured.push(run);
         });
@@ -112,21 +109,30 @@ describe("ReactiveSet", () => {
       });
 
       expect(captured).toHaveLength(1);
-      expect(captured[0]).toEqual([1, 2, 3]);
+      expect(captured[0]).toEqual([1, 2, 3, 4]);
 
       set.delete(4);
-      expect(captured, "deleted unseen key").toHaveLength(1);
+      expect(captured).toHaveLength(2);
+      expect(captured[1]).toEqual([1, 2, 3]);
 
       set.delete(1);
-      expect(captured, "deleted seen").toHaveLength(2);
-      expect(captured[1]).toEqual([2, 3]);
+      expect(captured).toHaveLength(3);
+      expect(captured[2]).toEqual([2, 3]);
 
       set.add(4);
-      expect(captured, "added key in reach").toHaveLength(3);
-      expect(captured[2]).toEqual([2, 3, 4]);
+      expect(captured).toHaveLength(4);
+      expect(captured[3]).toEqual([2, 3, 4]);
 
       set.add(5);
-      expect(captured, "added key out of reach").toHaveLength(3);
+      expect(captured).toHaveLength(5);
+      expect(captured[4]).toEqual([2, 3, 4, 5]);
+
+      set.add(5);
+      expect(captured).toHaveLength(5);
+
+      set.clear();
+      expect(captured).toHaveLength(6);
+      expect(captured[5]).toEqual([]);
 
       dispose();
     });