feat(ui): add SyncableMap

Can be used with useSyncExternal store to make a `Map` reactive.

invokeai/frontend/web/src/common/util/SyncableMap/SyncableMap.test.ts

@@ -0,0 +1,115 @@
+import { describe, expect, it, vi } from 'vitest';
+
+import { SyncableMap } from './SyncableMap';
+
+describe('SyncableMap', () => {
+  it('should initialize with entries', () => {
+    const initialEntries = [
+      ['key1', 'value1'],
+      ['key2', 'value2'],
+    ] as const;
+    const map = new SyncableMap(initialEntries);
+    expect(map.size).toBe(2);
+    expect(map.get('key1')).toBe('value1');
+    expect(map.get('key2')).toBe('value2');
+  });
+
+  it('should notify subscribers when a key is set', () => {
+    const map = new SyncableMap<string, string>();
+    const subscriber = vi.fn();
+    map.subscribe(subscriber);
+
+    map.set('key1', 'value1');
+
+    expect(subscriber).toHaveBeenCalledTimes(1);
+    expect(map.get('key1')).toBe('value1');
+  });
+
+  it('should notify subscribers when a key is deleted', () => {
+    const map = new SyncableMap<string, string>([['key1', 'value1']]);
+    const subscriber = vi.fn();
+    map.subscribe(subscriber);
+
+    map.delete('key1');
+
+    expect(subscriber).toHaveBeenCalledTimes(1);
+    expect(map.get('key1')).toBeUndefined();
+  });
+
+  it('should notify subscribers when the map is cleared', () => {
+    const map = new SyncableMap<string, string>([
+      ['key1', 'value1'],
+      ['key2', 'value2'],
+    ]);
+    const subscriber = vi.fn();
+    map.subscribe(subscriber);
+
+    map.clear();
+
+    expect(subscriber).toHaveBeenCalledTimes(1);
+    expect(map.size).toBe(0);
+  });
+
+  it('should not notify unsubscribed callbacks', () => {
+    const map = new SyncableMap<string, string>();
+    const subscriber = vi.fn();
+    const unsubscribe = map.subscribe(subscriber);
+
+    unsubscribe();
+
+    map.set('key1', 'value1');
+
+    expect(subscriber).not.toHaveBeenCalled();
+  });
+
+  it('should return a snapshot of the current state', () => {
+    const map = new SyncableMap<string, string>([['key1', 'value1']]);
+
+    const snapshot = map.getSnapshot();
+
+    expect(snapshot.size).toBe(1);
+    expect(snapshot.get('key1')).toBe('value1');
+  });
+
+  it('should return the same snapshot if there were no changes', () => {
+    const map = new SyncableMap<string, string>([['key1', 'value1']]);
+
+    const firstSnapshot = map.getSnapshot();
+    const secondSnapshot = map.getSnapshot();
+
+    expect(firstSnapshot).toBe(secondSnapshot);
+  });
+
+  it('should return a new snapshot if changes were made', () => {
+    const map = new SyncableMap<string, string>([['key1', 'value1']]);
+
+    const firstSnapshot = map.getSnapshot();
+    map.set('key2', 'value2');
+    const secondSnapshot = map.getSnapshot();
+
+    expect(firstSnapshot).not.toBe(secondSnapshot);
+    expect(secondSnapshot.size).toBe(2);
+  });
+
+  it('should consider different snapshots unequal', () => {
+    const map = new SyncableMap<string, string>([['key1', 'value1']]);
+
+    const firstSnapshot = map.getSnapshot();
+    map.set('key2', 'value2');
+    const secondSnapshot = map.getSnapshot();
+
+    expect(map['areSnapshotsEqual'](firstSnapshot, secondSnapshot)).toBe(false);
+  });
+
+  it('should consider identical snapshots equal', () => {
+    const map = new SyncableMap<string, string>([
+      ['key1', 'value1'],
+      ['key2', 'value2'],
+    ]);
+
+    const firstSnapshot = map.getSnapshot();
+    const secondSnapshot = map.getSnapshot();
+
+    expect(map['areSnapshotsEqual'](firstSnapshot, secondSnapshot)).toBe(true);
+  });
+});

invokeai/frontend/web/src/common/util/SyncableMap/SyncableMap.ts

@@ -0,0 +1,86 @@
+/**
+ * A Map that allows for subscribing to changes and getting a snapshot of the current state.
+ *
+ * It can be used with the `useSyncExternalStore` hook to sync the state of the map with a React component.
+ *
+ * Reactivity is shallow, so changes to nested objects will not trigger a re-render.
+ */
+export class SyncableMap<K, V> extends Map<K, V> {
+  private subscriptions = new Set<() => void>();
+  private lastSnapshot: Map<K, V> | null = null;
+
+  constructor(entries?: readonly (readonly [K, V])[] | null) {
+    super(entries);
+  }
+
+  set = (key: K, value: V): this => {
+    super.set(key, value);
+    this.notifySubscribers();
+    return this;
+  };
+
+  delete = (key: K): boolean => {
+    const result = super.delete(key);
+    this.notifySubscribers();
+    return result;
+  };
+
+  clear = (): void => {
+    super.clear();
+    this.notifySubscribers();
+  };
+
+  /**
+   * Notify all subscribers that the map has changed.
+   */
+  private notifySubscribers = () => {
+    for (const callback of this.subscriptions) {
+      callback();
+    }
+  };
+
+  /**
+   * Subscribe to changes to the map.
+   * @param callback A function to call when the map changes
+   * @returns A function to unsubscribe from changes
+   */
+  subscribe = (callback: () => void): (() => void) => {
+    this.subscriptions.add(callback);
+    return () => {
+      this.subscriptions.delete(callback);
+    };
+  };
+
+  /**
+   * Get a snapshot of the current state of the map.
+   * @returns A snapshot of the current state of the map
+   */
+  getSnapshot = (): Map<K, V> => {
+    const currentSnapshot = new Map(this);
+    if (!this.lastSnapshot || !this.areSnapshotsEqual(this.lastSnapshot, currentSnapshot)) {
+      this.lastSnapshot = currentSnapshot;
+    }
+
+    return this.lastSnapshot;
+  };
+
+  /**
+   * Compare two snapshots to determine if they are equal.
+   * @param snapshotA The first snapshot to compare
+   * @param snapshotB The second snapshot to compare
+   * @returns Whether the two snapshots are equal
+   */
+  private areSnapshotsEqual = (snapshotA: Map<K, V>, snapshotB: Map<K, V>): boolean => {
+    if (snapshotA.size !== snapshotB.size) {
+      return false;
+    }
+
+    for (const [key, value] of snapshotA) {
+      if (!Object.is(value, snapshotB.get(key))) {
+        return false;
+      }
+    }
+
+    return true;
+  };
+}