Fix and optimize ReactiveMap and ReactiveWeakMap (#704)

solidjs-community · Jan 2, 2025 · 37f0616 · 37f0616
2 parents 1daf96b + 46e4379
commit 37f0616
Show file tree

Hide file tree

Showing 4 changed files with 120 additions and 106 deletions.
diff --git a/.changeset/perfect-jars-visit.md b/.changeset/perfect-jars-visit.md
@@ -0,0 +1,5 @@
+---
+"@solid-primitives/map": minor
+---
+
+Fix and optimize ReactiveMap and ReactiveWeakMap (https://github.com/solidjs-community/solid-primitives/pull/704)
diff --git a/packages/map/src/index.ts b/packages/map/src/index.ts
@@ -1,7 +1,7 @@
 import { Accessor, batch } from "solid-js";
 import { TriggerCache } from "@solid-primitives/trigger";
 
-const $KEYS = Symbol("track-keys");
+const $OBJECT = Symbol("track-object");
 
 /**
  * A reactive version of `Map` data structure. All the reads (like `get` or `has`) are signals, and all the writes (`delete` or `set`) will cause updates to appropriate signals.
@@ -21,100 +21,114 @@ const $KEYS = Symbol("track-keys");
  * userPoints.set(user1, { foo: "bar" });
  */
 export class ReactiveMap<K, V> extends Map<K, V> {
-  #keyTriggers = new TriggerCache<K | typeof $KEYS>();
-  #valueTriggers = new TriggerCache<K>();
+  #keyTriggers = new TriggerCache<K | typeof $OBJECT>();
+  #valueTriggers = new TriggerCache<K | typeof $OBJECT>();
 
-  constructor(initial?: Iterable<readonly [K, V]> | null) {
-    super();
-    if (initial) for (const v of initial) super.set(v[0], v[1]);
+  [Symbol.iterator](): MapIterator<[K, V]> {
+    return this.entries();
   }
 
-  // reads
-  has(key: K): boolean {
-    this.#keyTriggers.track(key);
-    return super.has(key);
-  }
-  get(key: K): V | undefined {
-    this.#valueTriggers.track(key);
-    return super.get(key);
+  constructor(entries?: Iterable<readonly [K, V]> | null) {
+    super();
+    if (entries) for (const entry of entries) super.set(...entry);
   }
+
   get size(): number {
-    this.#keyTriggers.track($KEYS);
+    this.#keyTriggers.track($OBJECT);
     return super.size;
   }
 
   *keys(): MapIterator<K> {
+    this.#keyTriggers.track($OBJECT);
+
     for (const key of super.keys()) {
-      this.#keyTriggers.track(key);
       yield key;
     }
-    this.#keyTriggers.track($KEYS);
   }
+
   *values(): MapIterator<V> {
-    for (const [key, v] of super.entries()) {
-      this.#valueTriggers.track(key);
-      yield v;
+    this.#valueTriggers.track($OBJECT);
+
+    for (const value of super.values()) {
+      yield value;
     }
-    this.#keyTriggers.track($KEYS);
   }
+
   *entries(): MapIterator<[K, V]> {
+    this.#keyTriggers.track($OBJECT);
+    this.#valueTriggers.track($OBJECT);
+
     for (const entry of super.entries()) {
-      this.#valueTriggers.track(entry[0]);
       yield entry;
     }
-    this.#keyTriggers.track($KEYS);
   }
 
-  // writes
+  forEach(callbackfn: (value: V, key: K, map: Map<K, V>) => void, thisArg?: any): void {
+    this.#keyTriggers.track($OBJECT);
+    this.#valueTriggers.track($OBJECT);
+    super.forEach(callbackfn, thisArg);
+  }
+
+  has(key: K): boolean {
+    this.#keyTriggers.track(key);
+    return super.has(key);
+  }
+
+  get(key: K): V | undefined {
+    this.#valueTriggers.track(key);
+    return super.get(key);
+  }
+
   set(key: K, value: V): this {
-    batch(() => {
-      if (super.has(key)) {
-        if (super.get(key)! === value) return;
-      } else {
-        this.#keyTriggers.dirty(key);
-        this.#keyTriggers.dirty($KEYS);
-      }
-      this.#valueTriggers.dirty(key);
-      super.set(key, value);
-    });
-    return this;
+    const hadNoKey = !super.has(key);
+    const hasChanged = super.get(key) !== value;
+    const result = super.set(key, value);
+
+    if (hasChanged || hadNoKey) {
+      batch(() => {
+        if (hadNoKey) {
+          this.#keyTriggers.dirty($OBJECT);
+          this.#keyTriggers.dirty(key);
+        }
+        if (hasChanged) {
+          this.#valueTriggers.dirty($OBJECT);
+          this.#valueTriggers.dirty(key);
+        }
+      });
+    }
+
+    return result;
   }
+
   delete(key: K): boolean {
-    const r = super.delete(key);
-    if (r) {
+    const isDefined = super.get(key) !== undefined;
+    const result = super.delete(key);
+
+    if (result) {
       batch(() => {
+        this.#keyTriggers.dirty($OBJECT);
+        this.#valueTriggers.dirty($OBJECT);
         this.#keyTriggers.dirty(key);
-        this.#keyTriggers.dirty($KEYS);
-        this.#valueTriggers.dirty(key);
+
+        if (isDefined) {
+          this.#valueTriggers.dirty(key);
+        }
       });
     }
-    return r;
+
+    return result;
   }
+
   clear(): void {
     if (super.size) {
+      super.clear();
+
       batch(() => {
-        for (const v of super.keys()) {
-          this.#keyTriggers.dirty(v);
-          this.#valueTriggers.dirty(v);
-        }
-        super.clear();
-        this.#keyTriggers.dirty($KEYS);
+        this.#keyTriggers.dirtyAll();
+        this.#valueTriggers.dirtyAll();
       });
     }
   }
-
-  // callback
-  forEach(callbackfn: (value: V, key: K, map: this) => void) {
-    this.#keyTriggers.track($KEYS);
-    for (const [key, v] of super.entries()) {
-      this.#valueTriggers.track(key);
-      callbackfn(v, key, this);
-    }
-  }
-
-  [Symbol.iterator](): MapIterator<[K, V]> {
-    return this.entries();
-  }
 }
 
 /**
@@ -137,9 +151,9 @@ export class ReactiveWeakMap<K extends object, V> extends WeakMap<K, V> {
   #keyTriggers = new TriggerCache<K>(WeakMap);
   #valueTriggers = new TriggerCache<K>(WeakMap);
 
-  constructor(initial?: Iterable<readonly [K, V]> | null) {
+  constructor(entries?: Iterable<readonly [K, V]> | null) {
     super();
-    if (initial) for (const v of initial) super.set(v[0], v[1]);
+    if (entries) for (const entry of entries) super.set(...entry);
   }
 
   has(key: K): boolean {
@@ -151,24 +165,31 @@ export class ReactiveWeakMap<K extends object, V> extends WeakMap<K, V> {
     return super.get(key);
   }
   set(key: K, value: V): this {
-    batch(() => {
-      if (super.has(key)) {
-        if (super.get(key)! === value) return;
-      } else this.#keyTriggers.dirty(key);
-      this.#valueTriggers.dirty(key);
-      super.set(key, value);
-    });
-    return this;
+    const hadNoKey = !super.has(key);
+    const hasChanged = super.get(key) !== value;
+    const result = super.set(key, value);
+
+    if (hasChanged || hadNoKey) {
+      batch(() => {
+        if (hadNoKey) this.#keyTriggers.dirty(key);
+        if (hasChanged) this.#valueTriggers.dirty(key);
+      });
+    }
+
+    return result;
   }
   delete(key: K): boolean {
-    const r = super.delete(key);
-    if (r) {
+    const isDefined = super.get(key) !== undefined;
+    const result = super.delete(key);
+
+    if (result) {
       batch(() => {
         this.#keyTriggers.dirty(key);
-        this.#valueTriggers.dirty(key);
+        if (isDefined) this.#valueTriggers.dirty(key);
       });
     }
-    return r;
+
+    return result;
   }
 }
 

diff --git a/packages/map/test/index.test.ts b/packages/map/test/index.test.ts
@@ -210,7 +210,6 @@ describe("ReactiveMap", () => {
       [1, "a"],
       [2, "b"],
       [3, "c"],
-      [4, "d"],
     ]);
 
     const captured: unknown[][] = [];
@@ -220,7 +219,6 @@ describe("ReactiveMap", () => {
         const run: unknown[] = [];
         for (const key of map.keys()) {
           run.push(key);
-          if (key === 3) break; // don't iterate over all keys
         }
         captured.push(run);
       });
@@ -233,12 +231,12 @@ describe("ReactiveMap", () => {
     map.set(1, "e");
     expect(captured, "value change").toHaveLength(1);
 
-    map.set(5, "f");
-    expect(captured, "not seen key change").toHaveLength(1);
+    map.set(4, "f");
+    expect(captured, "new key added").toHaveLength(2);
 
     map.delete(1);
-    expect(captured, "seen key change").toHaveLength(2);
-    expect(captured[1]).toEqual([2, 3]);
+    expect(captured, "seen key change").toHaveLength(3);
+    expect(captured[2]).toEqual([2, 3, 4]);
 
     dispose();
   });
@@ -248,19 +246,15 @@ describe("ReactiveMap", () => {
       [1, "a"],
       [2, "b"],
       [3, "c"],
-      [4, "d"],
     ]);
 
     const captured: unknown[][] = [];
 
     const dispose = createRoot(dispose => {
       createEffect(() => {
         const run: unknown[] = [];
-        let i = 0;
         for (const v of map.values()) {
           run.push(v);
-          if (i === 2) break; // don't iterate over all keys
-          i += 1;
         }
         captured.push(run);
       });
@@ -275,14 +269,12 @@ describe("ReactiveMap", () => {
     expect(captured[1]).toEqual(["e", "b", "c"]);
 
     map.set(4, "f");
-    expect(captured, "not seen value change").toHaveLength(2);
+    expect(captured, "new key added").toHaveLength(3);
+    expect(captured[2]).toEqual(["e", "b", "c", "f"]);
 
     map.delete(4);
-    expect(captured, "not seen key change").toHaveLength(2);
-
-    map.delete(1);
-    expect(captured, "seen key change").toHaveLength(3);
-    expect(captured[2]).toEqual(["b", "c"]);
+    expect(captured, "key removed").toHaveLength(4);
+    expect(captured[3]).toEqual(["e", "b", "c"]);
 
     dispose();
   });
@@ -292,19 +284,15 @@ describe("ReactiveMap", () => {
       [1, "a"],
       [2, "b"],
       [3, "c"],
-      [4, "d"],
     ]);
 
     const captured: unknown[][] = [];
 
     const dispose = createRoot(dispose => {
       createEffect(() => {
         const run: unknown[] = [];
-        let i = 0;
         for (const e of map.entries()) {
           run.push(e);
-          if (i === 2) break; // don't iterate over all keys
-          i += 1;
         }
         captured.push(run);
       });
@@ -327,14 +315,18 @@ describe("ReactiveMap", () => {
     ]);
 
     map.set(4, "f");
-    expect(captured, "not seen value change").toHaveLength(2);
+    expect(captured, "new key added").toHaveLength(3);
+    expect(captured[2]).toEqual([
+      [1, "e"],
+      [2, "b"],
+      [3, "c"],
+      [4, "f"],
+    ]);
 
     map.delete(4);
-    expect(captured, "not seen key change").toHaveLength(2);
-
-    map.delete(1);
-    expect(captured, "seen key change").toHaveLength(3);
-    expect(captured[2]).toEqual([
+    expect(captured, "key removed").toHaveLength(4);
+    expect(captured[3]).toEqual([
+      [1, "e"],
       [2, "b"],
       [3, "c"],
     ]);
@@ -347,7 +339,6 @@ describe("ReactiveMap", () => {
       [1, "a"],
       [2, "b"],
       [3, "c"],
-      [4, "d"],
     ]);
 
     const captured: unknown[][] = [];
@@ -368,7 +359,6 @@ describe("ReactiveMap", () => {
       [1, "a"],
       [2, "b"],
       [3, "c"],
-      [4, "d"],
     ]);
 
     map.set(1, "e");
@@ -377,15 +367,13 @@ describe("ReactiveMap", () => {
       [1, "e"],
       [2, "b"],
       [3, "c"],
-      [4, "d"],
     ]);
 
-    map.delete(4);
+    map.delete(3);
     expect(captured).toHaveLength(3);
     expect(captured[2]).toEqual([
       [1, "e"],
       [2, "b"],
-      [3, "c"],
     ]);
 
     dispose();