diff --git a/invokeai/frontend/web/src/common/util/SyncableMap/SyncableMap.test.ts b/invokeai/frontend/web/src/common/util/SyncableMap/SyncableMap.test.ts new file mode 100644 index 0000000000..e65e8bf5a2 --- /dev/null +++ b/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(); + 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([['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([ + ['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(); + 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([['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([['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([['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([['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([ + ['key1', 'value1'], + ['key2', 'value2'], + ]); + + const firstSnapshot = map.getSnapshot(); + const secondSnapshot = map.getSnapshot(); + + expect(map['areSnapshotsEqual'](firstSnapshot, secondSnapshot)).toBe(true); + }); +}); diff --git a/invokeai/frontend/web/src/common/util/SyncableMap/SyncableMap.ts b/invokeai/frontend/web/src/common/util/SyncableMap/SyncableMap.ts new file mode 100644 index 0000000000..5743ce5ece --- /dev/null +++ b/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 extends Map { + private subscriptions = new Set<() => void>(); + private lastSnapshot: Map | 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 => { + 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, snapshotB: Map): 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; + }; +}