Fix SignaledSets

Exelord · Sep 9, 2024 · a4bbd16 · a4bbd16
1 parent 5ecd93b
commit a4bbd16
Show file tree

Hide file tree

Showing 3 changed files with 36 additions and 26 deletions.
diff --git a/src/primitives/set.ts b/src/primitives/set.ts
@@ -6,9 +6,9 @@ const OBJECT_KEYS = Symbol("objectKeys");
 export class SignaledSet<T> extends Set<T> {
   private readonly valuesCache = createCache();
 
-  constructor(values?: readonly T[] | null) {
+  constructor(values?: Iterable<T> | null) {
     super();
-    if (values) for (const v of values) super.add(v);
+    if (values) for (const value of values) super.add(value);
   }
 
   [Symbol.iterator](): IterableIterator<T> {
@@ -25,24 +25,27 @@ export class SignaledSet<T> extends Set<T> {
   }
 
   *values(): IterableIterator<T> {
-    for (const key of super.values()) {
-      track(key, this.valuesCache);
-      yield key;
-    }
     track(OBJECT_KEYS, this.valuesCache);
+
+    for (const value of super.values()) {
+      yield value;
+    }
   }
 
   *entries(): IterableIterator<[T, T]> {
-    for (const [key, value] of super.entries()) {
-      track(key, this.valuesCache);
-      yield [key, value];
-    }
     track(OBJECT_KEYS, this.valuesCache);
+
+    for (const entry of super.entries()) {
+      yield entry;
+    }
   }
 
-  forEach(fn: (value1: T, value2: T, set: Set<T>) => void): void {
+  forEach(
+    callbackfn: (value1: T, value2: T, set: Set<T>) => void,
+    thisArg?: any
+  ): void {
     track(OBJECT_KEYS, this.valuesCache);
-    super.forEach(fn);
+    super.forEach(callbackfn, thisArg);
   }
 
   has(value: T): boolean {
@@ -78,13 +81,14 @@ export class SignaledSet<T> extends Set<T> {
   clear(): void {
     if (super.size) {
       super.clear();
+
       batch(() => {
         dirtyAll(this.valuesCache);
       });
     }
   }
 }
 
-export function createSet<T>(values?: readonly T[] | null): SignaledSet<T> {
+export function createSet<T>(values?: Iterable<T> | null): SignaledSet<T> {
   return new SignaledSet<T>(values);
 }
diff --git a/src/primitives/weak-set.ts b/src/primitives/weak-set.ts
@@ -3,9 +3,9 @@ import { createWeakCache, track, dirty } from "../utils/cache";
 export class SignaledWeakSet<T extends object = object> extends WeakSet<T> {
   private readonly signalsCache = createWeakCache();
 
-  constructor(values?: readonly T[] | null) {
+  constructor(values?: Iterable<T> | null) {
     super();
-    if (values) for (const v of values) super.add(v);
+    if (values) for (const value of values) super.add(value);
   }
 
   has(value: T): boolean {
@@ -34,7 +34,7 @@ export class SignaledWeakSet<T extends object = object> extends WeakSet<T> {
 }
 
 export function createWeakSet<T extends object = object>(
-  values?: readonly T[] | null
+  values?: Iterable<T> | null
 ): SignaledWeakSet<T> {
   return new SignaledWeakSet(values);
 }
diff --git a/tests/src/primitives/set.test.ts b/tests/src/primitives/set.test.ts
@@ -105,33 +105,39 @@ describe("ReactiveSet - Solid Primitives", () => {
       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);
         });
         return dispose;
       });
 
       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();
     });