Trie.ts

export class Node<Value> {
  children: {
    [key: string]: Node<Value>
  } = {}

  value?: Value
  hasValue: boolean = false

  setValue (value: Value) {
    this.hasValue = true
    this.value = value
  }

  getChild (key: string, force?: true): this
  getChild (key: string, force: false): null | this
  getChild (key: string, force = true) {
    const child = this.children[key]
    if (child) {
      return child
    } else if (!force) {
      return null
    } else {
      // @ts-ignore
      const newChild = new this.constructor()
      this.children[key] = newChild
      return newChild
    }
  }
}

type Reducer<T> = (value?: T) => T

export class Trie<Value> extends Node<Value> {
  constructor (public defaultValue?: () => Value) {
    super()
  }

  private reduceValue (reducer: Reducer<Value>, node: Trie<Value>) {
    if (node.hasValue) {
      return reducer(node.value)
    } else if (this.defaultValue) {
      return reducer(this.defaultValue())
    } else {
      return reducer()
    }
  }

  keys (): Generator<string[]>
  keys (...prefix: string[]): Generator<string[]>
  * keys (...prefix: string[]) {
    if (this.hasValue) {
      yield prefix
    }
    for (const key of Object.keys(this.children)) {
      yield * this.getChild(key).keys(...prefix, key)
    }
  }

  entries (): Generator<[string[], Value]>
  entries (...prefix: string[]): Generator<[string[], Value]>
  * entries (...prefix: string[]) {
    if (this.hasValue) {
      yield [prefix, this.value]
    }
    for (const key of Object.keys(this.children)) {
      yield * this.getChild(key).entries(...prefix, key)
    }
  }

  reduce (path: string[], reducer: Reducer<Value>) {
    let node = this
    for (const step of path) {
      node = node.getChild(step)
    }
    node.setValue(this.reduceValue(reducer, node))
  }

  find (path: string[]) {
    let node = this
    for (const step of path) {
      node = node.getChild(step, false)
      if (!node) {
        return null
      }
    }
    return node.value
  }
}