tree.gno

  1package trie
  2
  3import "gno.land/p/nt/avl/v0"
  4
  5// Verify that Tree implements the ITree interface.
  6var _ avl.ITree = (*Tree)(nil)
  7
  8// Tree is a prefix tree (trie) keyed by string with values of any type.
  9// It implements the avl.ITree interface, so it can be used as a drop-in
 10// alternative to an AVL tree. Keys are ordered lexicographically by theiri
 11// raw bytes.
 12//
 13// The zero Tree is an empty, ready-to-use tree.
 14type Tree struct {
 15	root *node
 16	size int
 17}
 18
 19// NewTree creates a new empty trie tree.
 20func NewTree() *Tree {
 21	return &Tree{root: &node{}}
 22}
 23
 24// Size returns the number of key-value pairs stored in the tree.
 25func (t *Tree) Size() int {
 26	return t.size
 27}
 28
 29// Has reports whether key exists in the tree.
 30func (t *Tree) Has(key string) bool {
 31	if t.root == nil {
 32		return false
 33	}
 34
 35	_, exists := t.root.get(key)
 36	return exists
 37}
 38
 39// Get retrieves the value associated with key.
 40func (t *Tree) Get(key string) any {
 41	if t.root == nil {
 42		return nil
 43	}
 44
 45	v, _ := t.root.get(key)
 46	return v
 47}
 48
 49// GetByIndex retrieves the key-value pair at the given index
 50// in lexicographic order. It panics if index is out of range.
 51func (t *Tree) GetByIndex(index int) (key string, value any) {
 52	if index < 0 || index >= t.size {
 53		panic("trie: index out of range")
 54	}
 55
 56	i := 0
 57	t.root.inorder("", true, func(k string, v any) bool {
 58		if i == index {
 59			key, value = k, v
 60			return true
 61		}
 62
 63		i++
 64		return false
 65	})
 66	return key, value
 67}
 68
 69// Set inserts or updates a key-value pair.
 70// It returns true when the key already existed and its value
 71// was updated, and false when the key was newly inserted.
 72func (t *Tree) Set(key string, value any) (updated bool) {
 73	if t.root == nil {
 74		t.root = &node{}
 75	}
 76
 77	updated = t.root.set(key, value)
 78	if !updated {
 79		t.size++
 80	}
 81	return updated
 82}
 83
 84// Remove deletes the pair stored under key.
 85// It returns the removed value and whether the key was found and removed.
 86func (t *Tree) Remove(key string) (value any, removed bool) {
 87	if t.root == nil {
 88		return nil, false
 89	}
 90
 91	value, removed = t.root.remove(key)
 92	if removed {
 93		t.size--
 94	}
 95	return value, removed
 96}
 97
 98// Iterate performs an in-order (ascending) traversal over the keys in the
 99// range [start, end), start is inclusive and end is exclusive.
100// An empty start or end means unbounded on that side.
101// Iteration stops when the callback returns true, the method returns whether
102// it was stopped that way.
103func (t *Tree) Iterate(start, end string, cb avl.IterCbFn) bool {
104	if t.root == nil {
105		return false
106	}
107
108	return t.root.inorder("", true, func(key string, value any) bool {
109		if start != "" && key < start {
110			return false
111		}
112
113		if end != "" && key >= end {
114			return false
115		}
116		return cb(key, value)
117	})
118}
119
120// ReverseIterate performs a reverse in-order (descending) traversal over the
121// keys in the range [start, end], where both start and end are inclusive.
122// An empty start or end means unbounded on that side.
123// Iteration stops when the callback returns true; the method returns whether
124// it was stopped that way.
125func (t *Tree) ReverseIterate(start, end string, cb avl.IterCbFn) bool {
126	if t.root == nil {
127		return false
128	}
129
130	return t.root.inorder("", false, func(key string, value any) bool {
131		if start != "" && key < start {
132			return false
133		}
134
135		if end != "" && key > end {
136			return false
137		}
138		return cb(key, value)
139	})
140}
141
142// IterateByOffset performs an in-order (ascending) traversal that skips the
143// first offset keys and then visits up to count keys.
144// Iteration stops early when the callback returns true, the method returns
145// whether it was stopped that way.
146func (t *Tree) IterateByOffset(offset int, count int, cb avl.IterCbFn) bool {
147	return t.iterateByOffset(offset, count, true, cb)
148}
149
150// ReverseIterateByOffset performs a reverse in-order (descending) traversal
151// that skips the first offset keys and then visits up to count keys.
152// Iteration stops early when the callback returns true; the method returns whether
153// it was stopped that way.
154func (t *Tree) ReverseIterateByOffset(offset int, count int, cb avl.IterCbFn) bool {
155	return t.iterateByOffset(offset, count, false, cb)
156}
157
158func (t *Tree) iterateByOffset(offset, count int, ascending bool, cb avl.IterCbFn) bool {
159	if t.root == nil || count <= 0 {
160		return false
161	}
162
163	skipped, emitted := 0, 0
164	stopped := false
165	t.root.inorder("", ascending, func(key string, value any) bool {
166		if skipped < offset {
167			skipped++
168			return false
169		}
170
171		stopped = cb(key, value)
172		emitted++
173		return stopped || emitted >= count
174	})
175	return stopped
176}