Data structure for (Key, Value) pairs that keeps keys and values available in contiguous arrays

SparseSet.ts

/**
 * A class to store key value pairs while keeping
 * both available in contiguous arrays. Basically
 * always prefer a Map over this, unless you regularly
 * need to create arrays from the values or keys
 * (i.e. inside an update loop). Is strict about
 * duplicate or non-existing keys.
 *
 * - Does not keep insertion order!
 * - Throws error when using .set() with existing keys (use .mutate() instead)
 * - Throws error when using .get() with non-existing keys (use .has() to guard)
 * - The iterator will be a lot slower than iterating directly over
     keys/values or using SparseSet.forEach(). Mainly useful for destructuring,
     use with caution.
 * - **Don't** modify the arrays directly. They're public for readonly access.
 */

export class SparseSet<K extends any = any, V extends any = any> {
  public readonly values = new Array<V>();
  public readonly keys = new Array<K>();
  private _idxByKey = new Map<K, number>();

  /**
   * Constructs a new SparseSet.
   * @param {Iterable<[K, V]>} [iterable] An iterable to populate the set.
   */
  constructor(iterable?: Iterable<[K, V]>) {
    let i = 0;
    if (iterable) {
      for (const [key, value] of iterable) {
        this.values[i] = value;
        this.keys[i] = key;
        this._idxByKey.set(key, i++);
      }
    }
  }

  // ----------------------------------------------------------
  // mutators

  /**
   * Sets a key-value pair in the set.
   * @param {K} key The key to set.
   * @param {V} value The value to set.
   */
  set(key: K, value: V) {
    if (this._idxByKey.has(key)) {
      throw new Error(`Key ${key} already exists in SparseSet.`);
    } else {
      const len = this.values.push(value);
      this.keys.push(key);
      this._idxByKey.set(key, len - 1);
    }
  }

  /**
   * Mutates a value in the set.
   * @param key The key to mutate the value for.
   * @param mutatorFn The function to mutate the value with.
   */
  mutate(key: K, mutatorFn: (prevValue: V) => V) {
    if (this._idxByKey.has(key)) {
      const idx = this._idxByKey.get(key)!;
      const prev = this.values[idx];
      this.values[idx] = mutatorFn(prev);
    } else {
      throw new Error(`Key ${key} does not exist in SparseSet.`);
    }
  }

  /**
   * Deletes a key-value pair from the set.
   * @param {K} key The key to delete.
   * @returns {boolean} True if the key was deleted, false otherwise.
   * @throws {Error} If the key does not exist in the set.
   */
  delete(key: K) {
    const {values, keys} = this;
    const idx = this._idxByKey.get(key);
    if (idx == undefined) {
      throw new Error(`Key ${key} does not exist in SparseSet.`);
    }
    const lastElement = values.pop()!;
    const lastKey = keys.pop()!;

    if (idx !== values.length) {
      values[idx] = lastElement;
      keys[idx] = lastKey;
      this._idxByKey.set(lastKey, idx);
    }
    this._idxByKey.delete(key);

    return true;
  }

  /**
   * Deletes a key-value pair from the set.
   * @param {K} key The key to delete.
   * @returns {boolean} True if the key was deleted, false otherwise.
   * @throws {Error} If the key does not exist in the set.
   */
  clear() {
    this.values.length = 0;
    this.keys.length = 0;
    this._idxByKey.clear();
  }

  // ----------------------------------------------------------
  // accessors

  /**
   * Gets the value associated with a key.
   * @param {K} key The key to get.
   * @returns {V} The value associated with the key.
   * @throws {Error} If the key does not exist in the set.
   */
  get(key: K): V {
    const idx = this._idxByKey.get(key);
    if (idx == undefined) {
      throw new Error(`Key ${key} does not exist in SparseSet.`);
    }
    return this.values[idx];
  }

  /**
   * Checks if a key exists in the set.
   * @param {K} key The key to check.
   * @returns {boolean} True if the key exists in the set, false otherwise.
   */
  has(key: K): boolean {
    return this._idxByKey.has(key);
  }

  /**
   * Gets the size of the set.
   * @returns {number} The size of the set.
   */
  get size() {
    return this.values.length;
  }

  // ----------------------------------------------------------
  // iterator and iteration based methods

  /**
   * Returns an iterator for the set.
   * @returns {Iterator} An iterator for the set.
   */
  [Symbol.iterator]() {
    const {values, keys} = this;
    let i = 0;
    return {
      next() {
        if (i < values.length) {
          return {value: [keys[i], values[i++]], done: false};
        } else {
          return {done: true};
        }
      }
    }
  }

  /**
   * Executes a provided function once for each key-value pair in the set.
   * @param {function(V, K, SparseSet<K, V>): void} callback The function to execute for each key-value pair.
   */
  forEach(callback: (value: V, key: K, set: SparseSet<K, V>) => void) {
    for (let i = 0; i < this.size; i++) {
      const key = this.keys[i];
      const value = this.values[i];
      callback(value, key, this);
    }
  }

  /**
   * Creates a new set with all key-value pairs that pass the test implemented by the provided function.
   * @param {function(V, K, SparseSet<K, V>): boolean} callback The function to test each key-value pair.
   * @returns {SparseSet<K, V>} A new set with all key-value pairs that pass the test.
   */
  filter(callback: (value: V, key: K, set: SparseSet<K, V>) => boolean): SparseSet<K, V> {
    const result = new SparseSet<K, V>();
    for (let i = 0; i < this.size; i++) {
      const key = this.keys[i];
      const value = this.values[i];
      if (callback(value, key, this)) {
        result.set(key, value);
      }
    }
    return result;
  }

  /**
   * Deletes all key-value pairs that pass the test implemented by the provided function.
   * @param {function(V, K, SparseSet<K, V>): boolean} callback The function to test each key-value pair.
   * @returns {boolean} True if any key-value pairs were deleted, false otherwise.
   */
  deleteIf(callback: (value: V, key: K, set: SparseSet<K, V>) => unknown): boolean {
    let deleted = false;

    for (let i = 0; i < this.size;) {
      const key = this.keys[i];
      const value = this.values[i];
      if (callback(value, key, this)) {
        this.delete(key);
        deleted = true;
      } else {
        i++;
      }
    }

    return deleted;
  }

  // ----------------------------------------------------------
  // copy and clone

  /**
   * Copies all key-value pairs from a set into a new set.
   * @param {SparseSet<K, V>} other The set to copy from.
   */
  static copyFrom<K, V>(other: SparseSet<K, V>) {
    type Mutable<T> = { -readonly [P in keyof T]: T[P] };
    const copy = new SparseSet<K, V>();
    (copy.keys as Mutable<K[]>) = Array.from(other.keys);
    (copy.values as Mutable<V[]>) = Array.from(other.values);
    copy._idxByKey = new Map(other._idxByKey);
    return copy;
  }

  /**
   * Creates a new set that is a shallow clone of this set.
   * @returns {SparseSet<K, V>} A new set that is a shallow clone of this set.
   */
  clone() {
    return SparseSet.copyFrom(this);
  }

}