Merge pull request #21128 from megothss/fix/tracked-map-set-reactivity

[BUGFIX] Fix trackedMap and trackedWeakMap reactivity for existing keys

Author: NullVoxPopuli

packages/@glimmer-workspace/integration-tests/test/collections/map-test.ts

@@ -268,6 +268,82 @@ class TrackedMapTest extends RenderTest {
     );
   }
 
+  @test
+  'values: set existing value triggers reactivity'() {
+    this.assertReactivity(
+      class extends Component {
+        map = trackedMap([['foo', 123]]);
+
+        get value() {
+          let sum = 0;
+          for (const v of this.map.values()) sum += v;
+          return sum;
+        }
+
+        update() {
+          this.map.set('foo', 456);
+        }
+      }
+    );
+  }
+
+  @test
+  'entries: set existing value triggers reactivity'() {
+    this.assertReactivity(
+      class extends Component {
+        map = trackedMap([['foo', 123]]);
+
+        get value() {
+          let result = '';
+          for (const [k, v] of this.map.entries()) result += `${k}:${v}`;
+          return result;
+        }
+
+        update() {
+          this.map.set('foo', 456);
+        }
+      }
+    );
+  }
+
+  @test
+  'forEach: set existing value triggers reactivity'() {
+    this.assertReactivity(
+      class extends Component {
+        map = trackedMap([['foo', 123]]);
+
+        get value() {
+          let sum = 0;
+          this.map.forEach((v) => {
+            sum += v;
+          });
+          return sum;
+        }
+
+        update() {
+          this.map.set('foo', 456);
+        }
+      }
+    );
+  }
+
+  @test
+  'set existing falsy value triggers reactivity'() {
+    this.assertReactivity(
+      class extends Component {
+        map = trackedMap<string, number>([['foo', 0]]);
+
+        get value() {
+          return this.map.get('foo');
+        }
+
+        update() {
+          this.map.set('foo', 123);
+        }
+      }
+    );
+  }
+
   @test
   'each: set'() {
     this.assertEachReactivity(

packages/@glimmer-workspace/integration-tests/test/collections/weak-map-test.ts

@@ -187,6 +187,24 @@ class TrackedWeakMapTest extends RenderTest {
     );
   }
 
+  @test
+  'set existing falsy value triggers reactivity'() {
+    this.assertReactivity(
+      class extends Component {
+        obj = {};
+        map = trackedWeakMap<object, number>([[this.obj, 0]]);
+
+        get value() {
+          return this.map.get(this.obj);
+        }
+
+        update() {
+          this.map.set(this.obj, 123);
+        }
+      }
+    );
+  }
+
   @test 'delete unrelated value'() {
     this.assertReactivity(
       class extends Component {

packages/@glimmer/validator/lib/collections/map.ts

@@ -51,10 +51,8 @@ class TrackedMap<K = unknown, V = unknown> implements Map<K, V> {
   }
 
   // **** ALL GETTERS ****
-  entries() {
-    consumeTag(this.#collection);
-
-    return this.#vals.entries();
+  entries(): MapIterator<[K, V]> {
+    return this[Symbol.iterator]();
   }
 
   getOrInsert(key: K, defaultValue: V): V {
@@ -82,15 +80,23 @@ class TrackedMap<K = unknown, V = unknown> implements Map<K, V> {
   }
 
   values() {
-    consumeTag(this.#collection);
+    let iterator = this[Symbol.iterator]();
 
-    return this.#vals.values();
+    return {
+      next() {
+        let { value, done } = iterator.next();
+        return { value: done ? (undefined as V) : value?.[1], done };
+      },
+      [Symbol.iterator]() {
+        return this;
+      },
+    } as MapIterator<V>;
   }
 
   forEach(fn: (value: V, key: K, map: Map<K, V>) => void): void {
-    consumeTag(this.#collection);
-
-    this.#vals.forEach(fn);
+    for (let [key, value] of this) {
+      fn(value, key, this);
+    }
   }
 
   get size(): number {
@@ -101,8 +107,10 @@ class TrackedMap<K = unknown, V = unknown> implements Map<K, V> {
 
   /**
    * When iterating:
-   * - we entangle with the collection (as we iterate over the whole thing
+   * - we entangle with the collection (as we iterate over the whole thing)
+   *   via keys() → consumeTag(#collection)
    * - for each individual item, we entangle with the item as well
+   *   via get() → consumeTag(#storageFor(key))
    */
   [Symbol.iterator]() {
     let keys = this.keys();
@@ -121,6 +129,9 @@ class TrackedMap<K = unknown, V = unknown> implements Map<K, V> {
         // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
         return { value: [currentKey, self.get(currentKey!)], done: false };
       },
+      [Symbol.iterator]() {
+        return this;
+      },
     } as MapIterator<[K, V]>;
   }
 
@@ -129,10 +140,10 @@ class TrackedMap<K = unknown, V = unknown> implements Map<K, V> {
   }
 
   set(key: K, value: V): this {
-    let existing = this.#vals.get(key);
+    let hasExisting = this.#vals.has(key);
 
-    if (existing) {
-      let isUnchanged = this.#options.equals(existing, value);
+    if (hasExisting) {
+      let isUnchanged = this.#options.equals(this.#vals.get(key) as V, value);
 
       if (isUnchanged) {
         return this;
@@ -141,7 +152,7 @@ class TrackedMap<K = unknown, V = unknown> implements Map<K, V> {
 
     this.#dirtyStorageFor(key);
 
-    if (!existing) {
+    if (!hasExisting) {
       DIRTY_TAG(this.#collection);
     }
 

packages/@glimmer/validator/lib/collections/weak-map.ts

@@ -51,10 +51,10 @@ class TrackedWeakMap<K extends WeakKey = object, V = unknown> implements WeakMap
   }
 
   set(key: K, value: V): this {
-    let existing = this.#vals.get(key);
+    let hasExisting = this.#vals.has(key);
 
-    if (existing) {
-      let isUnchanged = this.#options.equals(existing, value);
+    if (hasExisting) {
+      let isUnchanged = this.#options.equals(this.#vals.get(key) as V, value);
 
       if (isUnchanged) {
         return this;