chromium/ash/webui/recorder_app_ui/resources/core/utils/iterable_weak_set.ts

// Copyright 2024 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

/**
 * An iterable version of WeakSet.
 *
 * TODO(pihsun): Unit tests.
 */
export class IterableWeakSet<T extends object> {
  private readonly set = new Set<WeakRef<T>>();

  private readonly weakRefMap = new WeakMap<T, WeakRef<T>>();

  private readonly finalizationGroup = new FinalizationRegistry(
    (ref: WeakRef<T>) => {
      this.set.delete(ref);
    },
  );

  add(it: T): void {
    if (this.weakRefMap.has(it)) {
      return;
    }
    const weakRef = new WeakRef(it);
    this.weakRefMap.set(it, weakRef);
    this.set.add(weakRef);
    this.finalizationGroup.register(it, weakRef, weakRef);
  }

  delete(it: T): void {
    const weakRef = this.weakRefMap.get(it);
    if (weakRef !== undefined) {
      this.weakRefMap.delete(it);
      this.set.delete(weakRef);
      this.finalizationGroup.unregister(weakRef);
    }
  }

  /**
   * Gets the estimated size.
   *
   * Note that this is just an estimate, and relies on garbage collection on
   * whether the finalization has been run or not, so this should only be used
   * in testing.
   */
  sizeForTesting(): number {
    return this.set.size;
  }

  * [Symbol.iterator](): Iterator<T> {
    for (const ref of this.set) {
      const it = ref.deref();
      if (it !== undefined) {
        yield it;
      }
    }
  }

  has(it: T): boolean {
    return this.weakRefMap.get(it)?.deref() !== undefined;
  }
}