wesl-test/src/DeviceCache.ts

Write GPU shader tests as easily as regular unit tests. Test WGSL and WESL shaders with vitest or your favorite Node.js test framework.

/**
 * WeakMap cache for GPUDevice-based caching
 *
 * GPUDevices are weakly held to avoid memory leaks.
 */
export class DeviceCache<T> {
  private cache = new WeakMap<GPUDevice, Map<string, T>>();

  get(device: GPUDevice, key: string): T | undefined {
    return this.cache.get(device)?.get(key);
  }

  set(device: GPUDevice, key: string, value: T): void {
    let deviceCache = this.cache.get(device);
    if (!deviceCache) {
      deviceCache = new Map();
      this.cache.set(device, deviceCache);
    }
    deviceCache.set(key, value);
  }
}