nn/src/data/trie.ts

export interface InternalTrieNode {
  children: Map<string, InternalTrieNode>;
  token?: string;
}

export interface TrieNode {
  readonly children: Map<string, TrieNode>;
  readonly token?: string;
}

export class Trie implements Iterable<string> {
  protected root: InternalTrieNode = { children: new Map() };

  constructor(values?: string[]) {
    for (const value of values ?? []) {
      this.insert(value);
    }
  }

  public insert(value: string) {
    let node = this.root;
    for (const char of value) {
      node = node.children.get(char) ?? node.children.set(char, { children: new Map() }).get(char)!;
    }
    node.token = value;
  }

  public match(value: string): string | undefined {
    let node: TrieNode = this.root;
    let match: string | undefined = undefined;
    for (let i = 0; i < value.length && node.children.has(value[i]); i++) {
      match = (node = node.children.get(value[i])!).token ?? match;
    }
    return match;
  }

  public *[Symbol.iterator](): IterableIterator<string> {
    const queue: InternalTrieNode[] = [];
    let item: InternalTrieNode | undefined = this.root;
    while (item) {
      if (item.token) {
        yield item.token;
      }
      queue.push(...item.children.values());
      item = queue.pop();
    }
  }

  public toJSON(node: InternalTrieNode = this.root): object {
    return [
      ...node.token ? [node.token] : [],
      ...[...node.children.entries()].map(([k, c]) => [k, this.toJSON(c)]),
    ];
  }

  protected static nodeFromJSON(element: object): InternalTrieNode {
    if (!Array.isArray(element) || !element.every(e => typeof e === 'string'
      || Array.isArray(e) && e.length === 2)) {
      throw new Error('Malformed Trie JSON.');
    }

    return {
      children: new Map((typeof element[0] === 'string' ? element.slice(1) : element).map(([k, v]) => [k, this.nodeFromJSON(v)])),
      token: typeof element[0] === 'string' ? element[0] : undefined,
    };
  }

  public static fromJSON(json: object): Trie {
    const result = new Trie();
    result.root = Trie.nodeFromJSON(json);
    return result;
  }
}