trie.ts

// Ideally, Map would have this as a method, but the proposal is still only on
// stage 2: https://github.com/tc39/proposal-upsert
function emplace<K, V>(
  map: Map<K, V>,
  key: K,
  supplier: (key: K) => V,
): V;
function emplace<K, V, This>(
  map: Map<K, V>,
  key: K,
  supplier: (this: This, key: K) => V,
  thisArg: This,
): V;
function emplace<K, V, This>(
  map: Map<K, V>,
  key: K,
  supplier: (this: This | undefined, key: K) => V,
  thisArg?: This,
): V {
  if (map.has(key)) {
    return map.get(key)!;
  }

  const value = supplier.call(thisArg, key);
  map.set(key, value);
  return value;
}

type TrieNode<K> = {
  occupied: false;
  value: null;
  children: Map<K, TrieNode<K>>;
} | {
  occupied: true;
  value: Iterable<K>;
  children: Map<K, TrieNode<K>> | null;
};

export class Trie<K = never> {
  // This assertion is safe, because `head` is only used inside the class and
  // we only ever access its `children` property.
  private readonly _head = { children: null } as TrieNode<K>;

  private _size = 0;

  constructor(values?: Iterable<Iterable<K>>) {
    if (values != null) {
      for (const value of values) {
        this.add(value);
      }
    }
  }

  get size(): number {
    return this._size;
  }

  /**
   * If multiple values are inserted that are made up of the same referentially
   * equal elements, then only the first value is stored in the trie.
   *
   * @example
   * const trie = new Trie();
   * const sentinel = {};
   * const a = [sentinel];
   * const b = [sentinel];
   * const [value, ...rest] = trie.add(a).add(b);
   * a === value; // true
   * rest.length === 0; // true
   */
  add(iterable: Iterable<K>): this {
    let node = this._head;

    for (const value of iterable) {
      const children = this._ensureChildren(node);
      node = emplace(children, value, this._newNode);
    }

    if (node !== this._head && !node.occupied) {
      ++this._size;
      this._occupyWith(node, iterable);
    }

    return this;
  }

  get(iterable: Iterable<K>): Iterable<K> | undefined {
    const node = this._find(iterable);

    return node != null && node.occupied ? node.value : void 0;
  }

  has(iterable: Iterable<K>): boolean {
    return this._isOccupied(this._find(iterable));
  }

  delete(iterable: Iterable<K>): boolean {
    const stack: [K, TrieNode<K>][] = [];
    const node = this._find(iterable, this._accumulator, stack)!;
    if (!this._isOccupied(node)) {
      // This guarantees that `node` isn't null for the rest of the method.
      return false;
    }

    if (this.size === 1) {
      this.clear();
      return true;
    }

    if (node.children != null) {
      node.occupied = false;
      node.value = null;
      --this._size;
      return true;
    }

    // The stack is guaranteed to contain at least one pair.
    for (let k = stack.length - 1, p = k - 1; ; --k, --p) {
      if (p === -1) {
        this._head.children!.delete(stack[k][0]);
        break;
      }

      const children = stack[p][1].children!;
      if (children.size !== 1) {
        children.delete(stack[k][0]);
        break;
      }
    }

    --this._size;
    return true;
  }

  /**
   * Empty iterables will not match anything.
   */
  isPrefix(iterable: Iterable<K>): boolean {
    return this._find(iterable) != null;
  }

  clear(): void {
    this._size = 0;
    this._head.children = null;
  }

  /**
   * Empty iterables will not match anything.
   */
  *find(iterable: Iterable<K>): IterableIterator<Iterable<K>> {
    const node = this._find(iterable);

    if (node != null) {
      if (node.occupied) {
        yield node.value;
      }

      yield* this._walker(node.children);
    }
  }

  values(): IterableIterator<Iterable<K>> {
    return this._walker(this._head.children);
  }

  toJSON(): string {
    return JSON.stringify([...this]);
  }

  private _isOccupied(node: TrieNode<K> | null): boolean {
    return node != null && node.occupied;
  }

  /**
   * This method cannot be inlined, because the code would be inside a block
   * with a type guard.
   */
  private _occupyWith(node: TrieNode<K>, value: Iterable<K>): void {
    node.occupied = true;
    node.value = value;
  }

  private _identity(node: TrieNode<K>): TrieNode<K> {
    return node;
  }

  private _accumulator(
    this: [K, TrieNode<K>][],
    node: TrieNode<K>,
    key: K,
  ): TrieNode<K> {
    this.push([key, node]);
    return node;
  }

  private _find(
    iterable: Iterable<K>,
    nodeMapper: (node: TrieNode<K>, key: K) => TrieNode<K> = this._identity,
    thisArg?: any,
  ): TrieNode<K> | null {
    let node = this._head;

    for (const value of iterable) {
      const { children } = node;
      if (children == null) {
        return null;
      }

      const leaf = children.get(value);
      if (leaf == null) {
        return null;
      }

      node = nodeMapper.call(thisArg, leaf, value);
    }

    return node === this._head ? null : node;
  }

  private _newNode(): TrieNode<K> {
    // This assertion is safe, because either `occupied` will be left false and
    // more children are appended to this node or this is the last node in the
    // chain and in Trie#add() `occupied` is set to true.
    const children = null as unknown as Map<K, TrieNode<K>>;
    return { occupied: false, value: null, children };
  }

  private _ensureChildren(node: TrieNode<K>): Map<K, TrieNode<K>> {
    if (node.children == null) {
      node.children = new Map();
    }

    return node.children;
  }

  private *_walker(
    map: Map<K, TrieNode<K>> | null,
  ): IterableIterator<Iterable<K>> {
    if (map == null) {
      return;
    }

    for (const node of map.values()) {
      if (node.occupied) {
        yield node.value;
      }

      yield* this._walker(node.children);
    }
  }
}

Object.defineProperty(Trie.prototype, Symbol.iterator, {
  value: Trie.prototype.values,
  writable: true,
  configurable: true,
});

export interface Trie<K> extends Iterable<Iterable<K>> {
  [Symbol.iterator](): IterableIterator<Iterable<K>>;
}