@thi.ng/trie/multi-trie.js

Trie-based map data structure with prefix search/query support

class MultiTrie {
  constructor(pairs, opts) {
    this.opts = opts;
    pairs && this.into(pairs);
  }
  next = {};
  vals;
  n = 0;
  *[Symbol.iterator]() {
    const queue = [["", this]];
    while (queue.length) {
      const [prefix, node] = queue.pop();
      if (node.vals) {
        for (let v of node.vals) yield [prefix, v];
      } else {
        node.queueChildren(queue, prefix);
      }
    }
  }
  *keys(sep = "", prefix = "") {
    const queue = [[prefix, this]];
    while (queue.length) {
      const [key, node] = queue.pop();
      if (node.vals) {
        yield key;
      } else {
        node.queueChildren(queue, key, sep);
      }
    }
  }
  *values() {
    const queue = [this];
    while (queue.length) {
      const node = queue.pop();
      if (node.vals) {
        yield* node.vals;
      } else {
        queue.push(...Object.values(node.next));
      }
    }
  }
  *suffixes(prefix, withPrefix = false, sep = "") {
    const node = this.find(prefix);
    if (node) {
      yield* node.keys(
        sep,
        withPrefix ? Array.isArray(prefix) ? prefix.join(sep) : prefix.toString() : ""
      );
    }
  }
  clear() {
    this.next = {};
    this.n = 0;
    this.vals = void 0;
  }
  has(key) {
    return !!this.get(key);
  }
  hasPrefix(prefix) {
    return !!this.find(prefix);
  }
  get(key) {
    const node = this.find(key);
    return node ? node.vals : void 0;
  }
  find(key) {
    let node = this;
    for (let i = 0, n = key.length; i < n; i++) {
      node = node.next[key[i].toString()];
      if (!node) return;
    }
    return node;
  }
  /**
   * Returns longest known prefix for `key` as array. If array is
   * empty, the given key has no partial matches.
   *
   * @param key -
   */
  knownPrefix(key) {
    let node = this;
    const prefix = [];
    for (let i = 0, n = key.length; i < n; i++) {
      const k = key[i].toString();
      const next = node.next[k];
      if (!next) break;
      prefix.push(k);
      node = next;
    }
    return prefix;
  }
  hasKnownPrefix(key) {
    return this.knownPrefix(key).length > 0;
  }
  add(key, val) {
    let node = this;
    for (let i = 0, n = key.length; i < n; i++) {
      const k = key[i].toString();
      const next = node.next[k];
      node = !next ? (node.n++, node.next[k] = new MultiTrie(null, this.opts)) : next;
    }
    if (!node.vals) {
      const ctor = this.opts?.vals;
      node.vals = ctor ? ctor() : /* @__PURE__ */ new Set();
    }
    node.vals.add(val);
  }
  into(pairs) {
    for (let [k, v] of pairs) {
      this.add(k, v);
    }
  }
  delete(prefix, val) {
    const n = prefix.length;
    if (n < 1) return false;
    const path = [];
    const key = [];
    let i = 0;
    let node = this;
    for (; i < n; i++) {
      const k = prefix[i].toString();
      key.push(k);
      path.push(node);
      node = node.next[k];
      if (!node) return false;
    }
    if (val !== void 0) {
      const vals = node.vals;
      if (vals?.has(val)) {
        vals.delete(val);
        if (vals.size > 0) return true;
      } else {
        return false;
      }
    }
    while (node = path[--i]) {
      delete node.next[key[i]];
      if (--node.n) break;
    }
    return true;
  }
  queueChildren(queue, prefix, sep = "") {
    prefix = prefix.length ? prefix + sep : prefix;
    queue.push(
      ...Object.keys(this.next).map(
        (k) => [prefix + k, this.next[k]]
      )
    );
  }
}
const defMultiTrie = (pairs, opts) => new MultiTrie(pairs, opts);
export {
  MultiTrie,
  defMultiTrie
};