modules/indexes-search/hyper-p2p-trie-prefix/index.js

require('bare-process/global')
const EventEmitter = require('bare-events')
const { assertNonEmpty } = require('../../_shared/lib/errors.js')

const PROTOCOL = 'trie-prefix/v1'

class TrieNode {
  constructor () {
    this.children = new Map()
    this.terminal = false
    this.count = 0
  }
}

class HyperP2PTriePrefix extends EventEmitter {
  constructor (opts = {}) {
    super()
    this._root = new TrieNode()
    this._words = new Set()
    this._stats = { inserted: 0, queries: 0 }
  }

  insert (word) {
    assertNonEmpty(word, 'word')
    const w = String(word).toLowerCase()
    let node = this._root
    for (const ch of w) {
      if (!node.children.has(ch)) node.children.set(ch, new TrieNode())
      node = node.children.get(ch)
      node.count++
    }
    if (!node.terminal) {
      node.terminal = true
      this._words.add(w)
      this._stats.inserted++
      this.emit('insert', { word: w })
    }
    return true
  }

  prefixSearch (prefix) {
    this._stats.queries++
    const p = String(prefix || '').toLowerCase()
    let node = this._root
    for (const ch of p) {
      if (!node.children.has(ch)) return []
      node = node.children.get(ch)
    }
    const out = []
    const walk = (n, path) => {
      if (n.terminal) out.push(path)
      for (const [ch, child] of n.children) walk(child, path + ch)
    }
    walk(node, p)
    return out.sort()
  }

  has (word) {
    return this._words.has(String(word).toLowerCase())
  }

  remove (word) {
    const w = String(word).toLowerCase()
    if (!this._words.has(w)) return false
    let node = this._root
    const stack = []
    for (const ch of w) {
      stack.push([node, ch])
      node = node.children.get(ch)
      if (!node) return false
    }
    if (!node.terminal) return false
    node.terminal = false
    this._words.delete(w)
    for (let i = stack.length - 1; i >= 0; i--) {
      const [parent, ch] = stack[i]
      const child = parent.children.get(ch)
      child.count--
      if (child.count === 0 && !child.terminal) parent.children.delete(ch)
    }
    return true
  }

  list () {
    return [...this._words].sort()
  }

  countPrefix (prefix) {
    return this.prefixSearch(prefix).length
  }

  autocomplete (prefix, limit = 10) {
    const all = this.prefixSearch(prefix)
    return all.slice(0, Math.max(0, limit | 0))
  }

  getStats () {
    return {
      ...this._stats,
      words: this._words.size,
      protocol: PROTOCOL
    }
  }

  async ready () { return this }

  async close () {
    this._root = new TrieNode()
    this._words.clear()
    this.emit('closed')
  }
}

module.exports = { HyperP2PTriePrefix, PROTOCOL }