Merge pull request #21234 from emberjs/copilot/improve-tracked-collections-implementation

Author: kategengler

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

@@ -1,16 +1,30 @@
-import type { ReactiveOptions } from './types';
+/* eslint-disable @typescript-eslint/no-explicit-any */
+// Using a Proxy-based approach so that any new methods added to the Map
+// interface (like getOrInsert, getOrInsertComputed, etc.) are automatically
+// supported without needing to manually re-implement each one.
 
 import { consumeTag } from '../tracking';
 import { createUpdatableTag, DIRTY_TAG } from '../validators';
 
-class TrackedMap<K = unknown, V = unknown> implements Map<K, V> {
-  #options: ReactiveOptions<V>;
-  #collection = createUpdatableTag();
-  #storages = new Map<K, ReturnType<typeof createUpdatableTag>>();
-  #vals: Map<K, V>;
+type Tag = ReturnType<typeof createUpdatableTag>;
 
-  #storageFor(key: K): ReturnType<typeof createUpdatableTag> {
-    const storages = this.#storages;
+export function trackedMap<Key = any, Value = any>(
+  data?:
+    | Map<Key, Value>
+    | Iterable<readonly [Key, Value]>
+    | readonly (readonly [Key, Value])[]
+    | null,
+  options?: { equals?: (a: Value, b: Value) => boolean; description?: string }
+): Map<Key, Value> {
+  const equals = options?.equals ?? Object.is;
+  // TypeScript doesn't correctly resolve the overloads for calling the `Map`
+  // constructor for the no-value constructor. This resolves that.
+  const target: Map<Key, Value> =
+    data instanceof Map ? new Map(data.entries()) : new Map(data ?? []);
+  const collection = createUpdatableTag();
+  const storages = new Map<Key, Tag>();
+
+  function storageFor(key: Key): Tag {
     let storage = storages.get(key);
 
     if (storage === undefined) {
@@ -20,181 +34,98 @@ class TrackedMap<K = unknown, V = unknown> implements Map<K, V> {
 
     return storage;
   }
-  #dirtyStorageFor(key: K): void {
-    const storage = this.#storages.get(key);
+
+  function dirtyStorageFor(key: Key): void {
+    const storage = storages.get(key);
 
     if (storage) {
       DIRTY_TAG(storage);
     }
   }
 
-  constructor(
-    existing: readonly (readonly [K, V])[] | Iterable<readonly [K, V]> | null | Map<K, V>,
-    options: ReactiveOptions<V>
-  ) {
-    // TypeScript doesn't correctly resolve the overloads for calling the `Map`
-    // constructor for the no-value constructor. This resolves that.
-    this.#vals = existing instanceof Map ? new Map(existing.entries()) : new Map(existing);
-    this.#options = options;
-  }
-
-  get(key: K): V | undefined {
-    consumeTag(this.#storageFor(key));
-
-    return this.#vals.get(key);
-  }
-
-  has(key: K): boolean {
-    consumeTag(this.#storageFor(key));
+  const proxy: Map<Key, Value> = new Proxy(target, {
+    get(target, prop, receiver) {
+      if (prop === 'set') {
+        return function (key: Key, value: Value): Map<Key, Value> {
+          const hasExisting = target.has(key);
 
-    return this.#vals.has(key);
-  }
-
-  // **** ALL GETTERS ****
-  entries(): MapIterator<[K, V]> {
-    return this[Symbol.iterator]();
-  }
+          if (hasExisting) {
+            const isUnchanged = equals(target.get(key) as Value, value);
 
-  getOrInsert(key: K, defaultValue: V): V {
-    consumeTag(this.#storageFor(key));
+            if (isUnchanged) return proxy;
+          }
 
-    // eslint-disable-next-line @typescript-eslint/ban-ts-comment
-    // @ts-ignore -- older versions of TS don't yet have this method
-    // eslint-disable-next-line @typescript-eslint/no-unsafe-return, @typescript-eslint/no-unsafe-call
-    return this.#vals.getOrInsert(key, defaultValue);
-  }
-
-  getOrInsertComputed(key: K, creator: (key: K) => V): V {
-    consumeTag(this.#storageFor(key));
-
-    // eslint-disable-next-line @typescript-eslint/ban-ts-comment
-    // @ts-ignore -- older versions of TS don't yet have this method
-    // eslint-disable-next-line @typescript-eslint/no-unsafe-return, @typescript-eslint/no-unsafe-call
-    return this.#vals.getOrInsertComputed(key, creator);
-  }
+          dirtyStorageFor(key);
+          DIRTY_TAG(collection);
 
-  keys() {
-    consumeTag(this.#collection);
-
-    return this.#vals.keys();
-  }
-
-  values() {
-    let iterator = this[Symbol.iterator]();
-
-    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 {
-    for (let [key, value] of this) {
-      fn(value, key, this);
-    }
-  }
+          target.set(key, value);
 
-  get size(): number {
-    consumeTag(this.#collection);
+          return proxy;
+        };
+      }
 
-    return this.#vals.size;
-  }
+      if (prop === 'delete') {
+        return function (key: Key): boolean {
+          if (!target.has(key)) return false;
 
-  /**
-   * When iterating:
-   * - 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();
-    // eslint-disable-next-line @typescript-eslint/no-this-alias
-    let self = this;
-
-    return {
-      next() {
-        let next = keys.next();
-        let currentKey = next.value;
-
-        if (next.done) {
-          return { value: [undefined, undefined], done: true };
-        }
-
-        // 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]>;
-  }
+          dirtyStorageFor(key);
+          DIRTY_TAG(collection);
 
-  get [Symbol.toStringTag](): string {
-    return this.#vals[Symbol.toStringTag];
-  }
+          storages.delete(key);
+          return target.delete(key);
+        };
+      }
 
-  set(key: K, value: V): this {
-    let hasExisting = this.#vals.has(key);
+      if (prop === 'clear') {
+        return function (): void {
+          if (target.size === 0) return;
 
-    if (hasExisting) {
-      let isUnchanged = this.#options.equals(this.#vals.get(key) as V, value);
+          storages.forEach((s) => DIRTY_TAG(s));
+          storages.clear();
 
-      if (isUnchanged) {
-        return this;
+          DIRTY_TAG(collection);
+          target.clear();
+        };
       }
-    }
 
-    this.#dirtyStorageFor(key);
+      if (prop === 'get') {
+        return function (key: Key): Value | undefined {
+          consumeTag(storageFor(key));
 
-    if (!hasExisting) {
-      DIRTY_TAG(this.#collection);
-    }
-
-    this.#vals.set(key, value);
-
-    return this;
-  }
+          return target.get(key);
+        };
+      }
 
-  delete(key: K): boolean {
-    if (!this.#vals.has(key)) return false;
+      if (prop === 'has') {
+        return function (key: Key): boolean {
+          consumeTag(storageFor(key));
 
-    this.#dirtyStorageFor(key);
-    DIRTY_TAG(this.#collection);
+          return target.has(key);
+        };
+      }
 
-    this.#storages.delete(key);
-    return this.#vals.delete(key);
-  }
+      if (prop === 'size') {
+        consumeTag(collection);
 
-  clear(): void {
-    if (this.#vals.size === 0) return;
+        return target.size;
+      }
 
-    this.#storages.forEach((s) => DIRTY_TAG(s));
-    this.#storages.clear();
+      // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
+      const value = Reflect.get(target, prop, receiver);
 
-    DIRTY_TAG(this.#collection);
-    this.#vals.clear();
-  }
-}
+      if (typeof value === 'function') {
+        return function (this: any, ...args: any[]) {
+          consumeTag(collection);
 
-// So instanceof works
-Object.setPrototypeOf(TrackedMap.prototype, Map.prototype);
+          // eslint-disable-next-line @typescript-eslint/no-unsafe-return, @typescript-eslint/no-unsafe-member-access, @typescript-eslint/no-unsafe-call
+          return value.apply(target, args);
+        };
+      }
 
-export function trackedMap<Key = unknown, Value = unknown>(
-  data?:
-    | Map<Key, Value>
-    | Iterable<readonly [Key, Value]>
-    | readonly (readonly [Key, Value])[]
-    | null,
-  options?: { equals?: (a: Value, b: Value) => boolean; description?: string }
-): Map<Key, Value> {
-  return new TrackedMap(data ?? [], {
-    equals: options?.equals ?? Object.is,
-    description: options?.description,
+      // eslint-disable-next-line @typescript-eslint/no-unsafe-return
+      return value;
+    },
   });
+
+  return proxy;
 }