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

Trie-based ES6-like Map data structures with prefix search/query support

class MultiTrie {
  constructor(pairs, opts) {
    this.opts = opts;
    pairs && this.into(pairs);
  }
  opts;
  next = /* @__PURE__ */ new Map();
  vals;
  get size() {
    let size = 0;
    const stack = [this];
    while (stack.length) {
      const node = stack.pop();
      size += node.vals?.size ?? 0;
      if (node.next) stack.push(...node.next.values());
    }
    return size;
  }
  [Symbol.iterator]() {
    return this.iterate(
      (key, node) => [...node.vals].map((v) => [key, v])
    );
  }
  keys(prefix = [], includePrefix = true) {
    return this.iterate((key) => [key], prefix, includePrefix);
  }
  values(prefix) {
    return this.iterate((_, node) => node.vals, prefix);
  }
  entries(prefix, includePrefix = true) {
    return this.iterate(
      (key, node) => [...node.vals].map((v) => [key, v]),
      prefix,
      includePrefix
    );
  }
  clear() {
    this.next.clear();
    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.get(key[i]);
      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];
      const next = node.next.get(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];
      const next = node.next.get(k);
      if (!next) {
        const newNode = new MultiTrie(null, this.opts);
        node.next.set(k, newNode);
        node = newNode;
      } else {
        node = next;
      }
    }
    if (!node.vals) {
      node.vals = this.opts?.values?.() ?? /* @__PURE__ */ new Set();
    }
    node.vals.add(val);
  }
  into(pairs) {
    for (const [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];
      key.push(k);
      path.push(node);
      node = node.next.get(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]) {
      node.next.delete(key[i]);
      if (node.next.size) break;
    }
    return true;
  }
  toJSON() {
    return {
      next: [...this.next].reduce(
        (acc, [k, v]) => (acc[k] = v.toJSON(), acc),
        {}
      ),
      vals: this.vals ? [...this.vals] : void 0
    };
  }
  *iterate(fn, prefix = [], includePrefix = true) {
    const root = this.find(prefix);
    if (!root) return;
    const stack = [
      [includePrefix ? prefix : [], root]
    ];
    while (stack.length) {
      const [key, node] = stack.pop();
      if (node.vals) yield* fn(key, node);
      for (let [k, v] of node.next) {
        stack.push([key.concat(k), v]);
      }
    }
  }
}
const defMultiTrie = (pairs, opts) => new MultiTrie(pairs, opts);
const defMultiTrieFromJSON = (src, opts) => {
  const res = defMultiTrie(null, opts);
  const queue = [[[], src]];
  while (queue.length) {
    const [key, node] = queue.pop();
    if (node.vals) {
      for (const v of node.vals) res.add(key, v);
    }
    for (const [k, child] of Object.entries(node.next)) {
      queue.push([[...key, k], child]);
    }
  }
  return res;
};
export {
  MultiTrie,
  defMultiTrie,
  defMultiTrieFromJSON
};