trie.gno

  1// Package trie is a prefix tree (trie) for autocomplete, as a pure, reusable
  2// package: insert words, then ask for every word sharing a prefix.
  3//
  4// Everything is deterministic and allocation-friendly so it runs reproducibly
  5// on-chain: no maps in the hot path (map iteration order is unspecified, which
  6// would make Render output vary), no clocks, no chain imports. Children are
  7// kept in a slice sorted by rune, so completions always come out in
  8// lexicographic order — the same input always yields the same output.
  9//
 10// A live demo of this package (a gnoweb autocomplete box) is at
 11// [r/moul/x/daily/triedemo](/r/nym-moultest123/x/daily/triedemo/v0).
 12package trie
 13
 14import "sort"
 15
 16// MaxWordLen bounds a single word so insertion gas stays predictable.
 17const MaxWordLen = 64
 18
 19// node is one position in the tree. Children are ordered by Key so walks are
 20// deterministic; `terminal` marks the end of an inserted word (so "car" can be
 21// a word even when "carpet" is also stored).
 22type node struct {
 23	key      rune
 24	terminal bool
 25	children []*node
 26}
 27
 28// Trie is a prefix tree. The zero value is an empty, ready-to-use Trie.
 29type Trie struct {
 30	root  node
 31	count int
 32}
 33
 34// New returns an empty Trie.
 35func New() *Trie { return &Trie{} }
 36
 37// Len returns how many distinct words the Trie holds.
 38func (t *Trie) Len() int { return t.count }
 39
 40// child finds n's child for rune r, or nil. The slice is sorted, so this is a
 41// binary search.
 42func (n *node) child(r rune) *node {
 43	i := sort.Search(len(n.children), func(i int) bool { return n.children[i].key >= r })
 44	if i < len(n.children) && n.children[i].key == r {
 45		return n.children[i]
 46	}
 47	return nil
 48}
 49
 50// addChild inserts a child for rune r, keeping children sorted by key.
 51func (n *node) addChild(r rune) *node {
 52	i := sort.Search(len(n.children), func(i int) bool { return n.children[i].key >= r })
 53	if i < len(n.children) && n.children[i].key == r {
 54		return n.children[i]
 55	}
 56	c := &node{key: r}
 57	n.children = append(n.children, nil)
 58	copy(n.children[i+1:], n.children[i:])
 59	n.children[i] = c
 60	return c
 61}
 62
 63// Insert adds word to the Trie and reports whether it was newly added.
 64// The empty string and words longer than MaxWordLen are rejected (false).
 65// Inserting the same word twice is a no-op.
 66func (t *Trie) Insert(word string) bool {
 67	rs := []rune(word)
 68	if len(rs) == 0 || len(rs) > MaxWordLen {
 69		return false
 70	}
 71	n := &t.root
 72	for _, r := range rs {
 73		n = n.addChild(r)
 74	}
 75	if n.terminal {
 76		return false
 77	}
 78	n.terminal = true
 79	t.count++
 80	return true
 81}
 82
 83// Contains reports whether word was inserted as a complete word. A stored
 84// "carpet" does not make "car" Contains-true — only Insert does.
 85func (t *Trie) Contains(word string) bool {
 86	n := t.find(word)
 87	return n != nil && n.terminal
 88}
 89
 90// HasPrefix reports whether any stored word starts with prefix. The empty
 91// prefix matches whenever the Trie is non-empty.
 92func (t *Trie) HasPrefix(prefix string) bool {
 93	if prefix == "" {
 94		return t.count > 0
 95	}
 96	return t.find(prefix) != nil
 97}
 98
 99// find walks to the node for s, or returns nil when the path is absent.
100func (t *Trie) find(s string) *node {
101	n := &t.root
102	for _, r := range s {
103		n = n.child(r)
104		if n == nil {
105			return nil
106		}
107	}
108	return n
109}
110
111// Complete returns up to limit words starting with prefix, in lexicographic
112// order. A limit <= 0 means "no cap". An absent prefix yields an empty slice
113// (never nil), so callers can range over the result unconditionally.
114//
115// The empty prefix lists the whole Trie, which is what makes this usable as a
116// plain sorted listing too.
117func (t *Trie) Complete(prefix string, limit int) []string {
118	out := []string{}
119	start := t.find(prefix)
120	if start == nil {
121		return out
122	}
123	collect(start, []rune(prefix), &out, limit)
124	return out
125}
126
127// collect appends every word under n (depth-first, children already sorted)
128// to *out, stopping once limit is reached.
129func collect(n *node, path []rune, out *[]string, limit int) {
130	if limit > 0 && len(*out) >= limit {
131		return
132	}
133	if n.terminal {
134		*out = append(*out, string(path))
135	}
136	for _, c := range n.children {
137		if limit > 0 && len(*out) >= limit {
138			return
139		}
140		collect(c, append(path, c.key), out, limit)
141	}
142}
143
144// Words returns every stored word in lexicographic order.
145func (t *Trie) Words() []string { return t.Complete("", 0) }
146
147// FromWords builds a Trie from words, skipping any the Trie rejects.
148func FromWords(words []string) *Trie {
149	t := New()
150	for _, w := range words {
151		t.Insert(w)
152	}
153	return t
154}