## Trie implementation
## 2023/02/08
import options

const
  AlphabetSize = 26
type
  NodeValue = char # array[4, uint8]
  Trie* = object
    root: TrieNode
    size: int # number of words on the trie
  TrieNode = ref object
    nodes: seq[TrieNode]
    value: NodeValue
    isTerminal: bool # true if this node is the end of a word


proc add(self: TrieNode, value: NodeValue): TrieNode =
  result = TrieNode(
    nodes: newSeqOfCap[TrieNode](AlphabetSize),
    value: value,
  )
  self.nodes.add(result)

proc getNode(self: TrieNode, value: NodeValue): Option[TrieNode] =
  assert self != nil
  for n in self.nodes:
    if n.value == value:
      return some(n)
  return none(TrieNode)

proc followPath(self: Trie, str: string): Option[TrieNode] =
  var node = self.root
  for c in str:
    # echo "path: ", c
    let tmp = node.getNode(c)
    if tmp.isNone:
      return none(TrieNode)
    node = tmp.get()
  return some(node)

func len*(self: Trie): int = self.size

proc insert*(self: var Trie, word: string): var Trie {.discardable.} =
  result = self
  var node = self.root
  for c in word:
    let nextNode = node.getNode(c)
    if nextNode.isNone:
      node = node.add(c)
    else:
      node = nextNode.get()
  node.isTerminal = true
  inc self.size

proc delete*(self: var Trie, word: string): var Trie {.discardable.} =
  result = self
  let terminalNode = self.followPath(word)
  if terminalNode.isNone:
    return
  terminalNode.get.isTerminal = false
  dec self.size

proc prefixExists*(self: Trie, prefix: string): bool =
  self.followPath(prefix).isSome

proc searchWord*(self: Trie, str: string): bool =
  let nodeOpt = self.followPath(str)
  if nodeOpt.isNone:
    return false
  return nodeOpt.get.isTerminal

proc getWordsByPrefix*(self: Trie, prefix: string): seq[TrieNode] =
  let nodeOpt = self.followPath(prefix)
  if nodeOpt.isNone:
    return @[]
  for child in nodeOpt.unsafeGet.nodes:
    result.add(child)

when isMainModule:
  var t = Trie(root: new(TrieNode))
  let words = [
    "gnome",
    "gnostic",
    "genetics",
    "genome",
    "geek",
    "greek",
    "general",
    "generalist",
    "generator",
    "thomas",
    "though",
    "through"
  ]
  
  for w in words:
    t.insert(w)
    assert t.searchWord(w)
  
  for w in ["general", "gno", "th", "agact"]:
    echo "'", w, "' ", "exists? ", t.prefixExists(w)
  
  for n in t.getWordsByPrefix("gen"):
    echo "child value: ", n.value