node.gno

  1package trie
  2
  3// node is a single position in the trie tree
  4type node struct {
  5	char     byte    // node's char, unused in the root node
  6	value    any     // node's value, stored when isKey is true
  7	isKey    bool    // whether a key terminates at this node
  8	children []*node // child nodes, kept sorted ascending by char
  9}
 10
 11// childIndex returns the index of the child reached by byte c and whether such
 12// a child exists. When it does not exist, the returned index is the position at
 13// which a child with byte c must be inserted to keep children sorted.
 14func (n *node) childIndex(c byte) (int, bool) {
 15	lo := 0
 16	hi := len(n.children)
 17	for lo < hi {
 18		mid := (lo + hi) / 2
 19
 20		switch {
 21		case n.children[mid].char < c:
 22			lo = mid + 1
 23		case n.children[mid].char > c:
 24			hi = mid
 25		default:
 26			return mid, true
 27		}
 28	}
 29	return lo, false
 30}
 31
 32// get retrieves the value stored under key in the subtree rooted at n.
 33func (n *node) get(key string) (value any, exists bool) {
 34	cur := n
 35	for i := 0; i < len(key); i++ {
 36		idx, ok := cur.childIndex(key[i])
 37		if !ok {
 38			return nil, false
 39		}
 40
 41		cur = cur.children[idx]
 42	}
 43
 44	if !cur.isKey {
 45		return nil, false
 46	}
 47	return cur.value, true
 48}
 49
 50// set inserts or updates key with value. It returns true when the key already
 51// existed and its value updated, and false when it was newly inserted.
 52func (n *node) set(key string, value any) (updated bool) {
 53	cur := n
 54	for i := 0; i < len(key); i++ {
 55		c := key[i]
 56		idx, ok := cur.childIndex(c)
 57		if ok {
 58			cur = cur.children[idx]
 59			continue
 60		}
 61
 62		child := &node{char: c}
 63		cur.children = append(cur.children, nil)
 64		copy(cur.children[idx+1:], cur.children[idx:])
 65		cur.children[idx] = child
 66		cur = child
 67	}
 68
 69	if cur.isKey {
 70		cur.value = value
 71		return true
 72	}
 73
 74	cur.isKey = true
 75	cur.value = value
 76	return false
 77}
 78
 79// remove deletes key from the subtree rooted at n, pruning any nodes that no
 80// longer terminate a key and have no children. It returns the removed value and
 81// whether the key was found.
 82func (n *node) remove(key string) (value any, removed bool) {
 83	if key == "" {
 84		if !n.isKey {
 85			return nil, false
 86		}
 87
 88		value = n.value
 89		n.isKey = false
 90		n.value = nil
 91		return value, true
 92	}
 93
 94	idx, ok := n.childIndex(key[0])
 95	if !ok {
 96		return nil, false
 97	}
 98
 99	child := n.children[idx]
100	value, removed = child.remove(key[1:])
101	if removed && !child.isKey && len(child.children) == 0 {
102		n.children = append(n.children[:idx], n.children[idx+1:]...)
103	}
104	return value, removed
105}
106
107// inorder visits this node (when it terminates a key) and its whole subtree.
108// key is the full key spelled from the root to this node. When ascending is
109// true keys are produced in lexicographic order, otherwise in reverse. Because
110// a key always sorts before any key that extends it, the terminal value is
111// emitted before the children when ascending and after them when descending.
112// It returns true as soon as cb requests the walk to stop.
113func (n *node) inorder(key string, ascending bool, cb func(key string, value any) bool) (stop bool) {
114	if ascending {
115		if n.isKey && cb(key, n.value) {
116			return true
117		}
118
119		for _, c := range n.children {
120			if c.inorder(key+string([]byte{c.char}), ascending, cb) {
121				return true
122			}
123		}
124		return false
125	}
126
127	for i := len(n.children) - 1; i >= 0; i-- {
128		c := n.children[i]
129		if c.inorder(key+string([]byte{c.char}), ascending, cb) {
130			return true
131		}
132	}
133
134	if n.isKey && cb(key, n.value) {
135		return true
136	}
137	return false
138}