node.gno
13.29 Kb · 518 lines
1package cow
2
3//----------------------------------------
4// Node
5
6// Node represents a node in an AVL tree.
7type Node struct {
8 key string // key is the unique identifier for the node.
9 value any // value is the data stored in the node.
10 height int8 // height is the height of the node in the tree.
11 size int // size is the number of nodes in the subtree rooted at this node.
12 leftNode *Node // leftNode is the left child of the node.
13 rightNode *Node // rightNode is the right child of the node.
14}
15
16// NewNode creates a new node with the given key and value.
17func NewNode(key string, value any) *Node {
18 return &Node{
19 key: key,
20 value: value,
21 height: 0,
22 size: 1,
23 }
24}
25
26// Size returns the size of the subtree rooted at the node.
27func (node *Node) Size() int {
28 if node == nil {
29 return 0
30 }
31 return node.size
32}
33
34// IsLeaf checks if the node is a leaf node (has no children).
35func (node *Node) IsLeaf() bool {
36 return node.height == 0
37}
38
39// Key returns the key of the node.
40func (node *Node) Key() string {
41 return node.key
42}
43
44// Value returns the value of the node.
45func (node *Node) Value() any {
46 return node.value
47}
48
49// _copy creates a copy of the node (excluding value).
50func (node *Node) _copy() *Node {
51 if node.height == 0 {
52 panic("Why are you copying a value node?")
53 }
54 return &Node{
55 key: node.key,
56 height: node.height,
57 size: node.size,
58 leftNode: node.leftNode,
59 rightNode: node.rightNode,
60 }
61}
62
63// Has checks if a node with the given key exists in the subtree rooted at the node.
64func (node *Node) Has(key string) (has bool) {
65 if node == nil {
66 return false
67 }
68 if node.key == key {
69 return true
70 }
71 if node.height == 0 {
72 return false
73 }
74 if key < node.key {
75 return node.getLeftNode().Has(key)
76 }
77 return node.getRightNode().Has(key)
78}
79
80// Get searches for a node with the given key in the subtree rooted at the node
81// and returns its index, value, and whether it exists.
82func (node *Node) Get(key string) (index int, value any, exists bool) {
83 if node == nil {
84 return 0, nil, false
85 }
86
87 if node.height == 0 {
88 if node.key == key {
89 return 0, node.value, true
90 }
91 if node.key < key {
92 return 1, nil, false
93 }
94 return 0, nil, false
95 }
96
97 if key < node.key {
98 return node.getLeftNode().Get(key)
99 }
100
101 rightNode := node.getRightNode()
102 index, value, exists = rightNode.Get(key)
103 index += node.size - rightNode.size
104 return index, value, exists
105}
106
107// GetByIndex retrieves the key-value pair of the node at the given index
108// in the subtree rooted at the node.
109func (node *Node) GetByIndex(index int) (key string, value any) {
110 if node.height == 0 {
111 if index == 0 {
112 return node.key, node.value
113 }
114 panic("GetByIndex asked for invalid index")
115 }
116 // TODO: could improve this by storing the sizes
117 leftNode := node.getLeftNode()
118 if index < leftNode.size {
119 return leftNode.GetByIndex(index)
120 }
121 return node.getRightNode().GetByIndex(index - leftNode.size)
122}
123
124// Set inserts a new node with the given key-value pair into the subtree rooted at the node,
125// and returns the new root of the subtree and whether an existing node was updated.
126//
127// XXX consider a better way to do this... perhaps split Node from Node.
128func (node *Node) Set(key string, value any) (newSelf *Node, updated bool) {
129 if node == nil {
130 return NewNode(key, value), false
131 }
132
133 // Always create a new node for leaf nodes
134 if node.height == 0 {
135 return node.setLeaf(key, value)
136 }
137
138 // Copy the node before modifying
139 newNode := node._copy()
140 if key < node.key {
141 newNode.leftNode, updated = node.getLeftNode().Set(key, value)
142 } else {
143 newNode.rightNode, updated = node.getRightNode().Set(key, value)
144 }
145
146 if !updated {
147 newNode.calcHeightAndSize()
148 return newNode.balance(), updated
149 }
150
151 return newNode, updated
152}
153
154// setLeaf inserts a new leaf node with the given key-value pair into the subtree rooted at the node,
155// and returns the new root of the subtree and whether an existing node was updated.
156func (node *Node) setLeaf(key string, value any) (newSelf *Node, updated bool) {
157 if key == node.key {
158 return NewNode(key, value), true
159 }
160
161 if key < node.key {
162 return &Node{
163 key: node.key,
164 height: 1,
165 size: 2,
166 leftNode: NewNode(key, value),
167 rightNode: node,
168 }, false
169 }
170
171 return &Node{
172 key: key,
173 height: 1,
174 size: 2,
175 leftNode: node,
176 rightNode: NewNode(key, value),
177 }, false
178}
179
180// Remove deletes the node with the given key from the subtree rooted at the node.
181// returns the new root of the subtree, the new leftmost leaf key (if changed),
182// the removed value and the removal was successful.
183func (node *Node) Remove(key string) (
184 newNode *Node, newKey string, value any, removed bool,
185) {
186 if node == nil {
187 return nil, "", nil, false
188 }
189 if node.height == 0 {
190 if key == node.key {
191 return nil, "", node.value, true
192 }
193 return node, "", nil, false
194 }
195 if key < node.key {
196 var newLeftNode *Node
197 newLeftNode, newKey, value, removed = node.getLeftNode().Remove(key)
198 if !removed {
199 return node, "", value, false
200 }
201 if newLeftNode == nil { // left node held value, was removed
202 return node.rightNode, node.key, value, true
203 }
204 node = node._copy()
205 node.leftNode = newLeftNode
206 node.calcHeightAndSize()
207 node = node.balance()
208 return node, newKey, value, true
209 }
210
211 var newRightNode *Node
212 newRightNode, newKey, value, removed = node.getRightNode().Remove(key)
213 if !removed {
214 return node, "", value, false
215 }
216 if newRightNode == nil { // right node held value, was removed
217 return node.leftNode, "", value, true
218 }
219 node = node._copy()
220 node.rightNode = newRightNode
221 if newKey != "" {
222 node.key = newKey
223 }
224 node.calcHeightAndSize()
225 node = node.balance()
226 return node, "", value, true
227}
228
229// getLeftNode returns the left child of the node.
230func (node *Node) getLeftNode() *Node {
231 return node.leftNode
232}
233
234// getRightNode returns the right child of the node.
235func (node *Node) getRightNode() *Node {
236 return node.rightNode
237}
238
239// rotateRight performs a right rotation on the node and returns the new root.
240// NOTE: overwrites node
241// TODO: optimize balance & rotate
242func (node *Node) rotateRight() *Node {
243 node = node._copy()
244 l := node.getLeftNode()
245 _l := l._copy()
246
247 _lrCached := _l.rightNode
248 _l.rightNode = node
249 node.leftNode = _lrCached
250
251 node.calcHeightAndSize()
252 _l.calcHeightAndSize()
253
254 return _l
255}
256
257// rotateLeft performs a left rotation on the node and returns the new root.
258// NOTE: overwrites node
259// TODO: optimize balance & rotate
260func (node *Node) rotateLeft() *Node {
261 node = node._copy()
262 r := node.getRightNode()
263 _r := r._copy()
264
265 _rlCached := _r.leftNode
266 _r.leftNode = node
267 node.rightNode = _rlCached
268
269 node.calcHeightAndSize()
270 _r.calcHeightAndSize()
271
272 return _r
273}
274
275// calcHeightAndSize updates the height and size of the node based on its children.
276// NOTE: mutates height and size
277func (node *Node) calcHeightAndSize() {
278 node.height = maxInt8(node.getLeftNode().height, node.getRightNode().height) + 1
279 node.size = node.getLeftNode().size + node.getRightNode().size
280}
281
282// calcBalance calculates the balance factor of the node.
283func (node *Node) calcBalance() int {
284 return int(node.getLeftNode().height) - int(node.getRightNode().height)
285}
286
287// balance balances the subtree rooted at the node and returns the new root.
288// NOTE: assumes that node can be modified
289// TODO: optimize balance & rotate
290func (node *Node) balance() (newSelf *Node) {
291 balance := node.calcBalance()
292 if balance >= -1 {
293 return node
294 }
295 if balance > 1 {
296 if node.getLeftNode().calcBalance() >= 0 {
297 // Left Left Case
298 return node.rotateRight()
299 }
300 // Left Right Case
301 left := node.getLeftNode()
302 node.leftNode = left.rotateLeft()
303 return node.rotateRight()
304 }
305
306 if node.getRightNode().calcBalance() <= 0 {
307 // Right Right Case
308 return node.rotateLeft()
309 }
310
311 // Right Left Case
312 right := node.getRightNode()
313 node.rightNode = right.rotateRight()
314 return node.rotateLeft()
315}
316
317// Shortcut for TraverseInRange.
318func (node *Node) Iterate(start, end string, cb func(*Node) bool) bool {
319 return node.TraverseInRange(start, end, true, true, cb)
320}
321
322// Shortcut for TraverseInRange.
323func (node *Node) ReverseIterate(start, end string, cb func(*Node) bool) bool {
324 return node.TraverseInRange(start, end, false, true, cb)
325}
326
327// TraverseInRange traverses all nodes, including inner nodes.
328// Start is inclusive and end is exclusive when ascending,
329// Start and end are inclusive when descending.
330// Empty start and empty end denote no start and no end.
331// If leavesOnly is true, only visit leaf nodes.
332// NOTE: To simulate an exclusive reverse traversal,
333// just append 0x00 to start.
334func (node *Node) TraverseInRange(start, end string, ascending bool, leavesOnly bool, cb func(*Node) bool) bool {
335 if node == nil {
336 return false
337 }
338 afterStart := (start == "" || start < node.key)
339 startOrAfter := (start == "" || start <= node.key)
340 beforeEnd := false
341 if ascending {
342 beforeEnd = (end == "" || node.key < end)
343 } else {
344 beforeEnd = (end == "" || node.key <= end)
345 }
346
347 // Run callback per inner/leaf node.
348 stop := false
349 if (!node.IsLeaf() && !leavesOnly) ||
350 (node.IsLeaf() && startOrAfter && beforeEnd) {
351 stop = cb(node)
352 if stop {
353 return stop
354 }
355 }
356 if node.IsLeaf() {
357 return stop
358 }
359
360 if ascending {
361 // check lower nodes, then higher
362 if afterStart {
363 stop = node.getLeftNode().TraverseInRange(start, end, ascending, leavesOnly, cb)
364 }
365 if stop {
366 return stop
367 }
368 if beforeEnd {
369 stop = node.getRightNode().TraverseInRange(start, end, ascending, leavesOnly, cb)
370 }
371 } else {
372 // check the higher nodes first
373 if beforeEnd {
374 stop = node.getRightNode().TraverseInRange(start, end, ascending, leavesOnly, cb)
375 }
376 if stop {
377 return stop
378 }
379 if afterStart {
380 stop = node.getLeftNode().TraverseInRange(start, end, ascending, leavesOnly, cb)
381 }
382 }
383
384 return stop
385}
386
387// TraverseByOffset traverses all nodes, including inner nodes.
388// A limit of math.MaxInt means no limit.
389func (node *Node) TraverseByOffset(offset, limit int, descending bool, leavesOnly bool, cb func(*Node) bool) bool {
390 if node == nil {
391 return false
392 }
393
394 // fast paths. these happen only if TraverseByOffset is called directly on a leaf.
395 if limit <= 0 || offset >= node.size {
396 return false
397 }
398 if node.IsLeaf() {
399 if offset > 0 {
400 return false
401 }
402 return cb(node)
403 }
404
405 // go to the actual recursive function.
406 return node.traverseByOffset(offset, limit, descending, leavesOnly, cb)
407}
408
409// TraverseByOffset traverses the subtree rooted at the node by offset and limit,
410// in either ascending or descending order, and applies the callback function to each traversed node.
411// If leavesOnly is true, only leaf nodes are visited.
412func (node *Node) traverseByOffset(offset, limit int, descending bool, leavesOnly bool, cb func(*Node) bool) bool {
413 // caller guarantees: offset < node.size; limit > 0.
414 if !leavesOnly {
415 if cb(node) {
416 return true
417 }
418 }
419 first, second := node.getLeftNode(), node.getRightNode()
420 if descending {
421 first, second = second, first
422 }
423 if first.IsLeaf() {
424 // either run or skip, based on offset
425 if offset > 0 {
426 offset--
427 } else {
428 cb(first)
429 limit--
430 if limit <= 0 {
431 return false
432 }
433 }
434 } else {
435 // possible cases:
436 // 1 the offset given skips the first node entirely
437 // 2 the offset skips none or part of the first node, but the limit requires some of the second node.
438 // 3 the offset skips none or part of the first node, and the limit stops our search on the first node.
439 if offset >= first.size {
440 offset -= first.size // 1
441 } else {
442 if first.traverseByOffset(offset, limit, descending, leavesOnly, cb) {
443 return true
444 }
445 // number of leaves which could actually be called from inside
446 delta := first.size - offset
447 offset = 0
448 if delta >= limit {
449 return true // 3
450 }
451 limit -= delta // 2
452 }
453 }
454
455 // because of the caller guarantees and the way we handle the first node,
456 // at this point we know that limit > 0 and there must be some values in
457 // this second node that we include.
458
459 // => if the second node is a leaf, it has to be included.
460 if second.IsLeaf() {
461 return cb(second)
462 }
463 // => if it is not a leaf, it will still be enough to recursively call this
464 // function with the updated offset and limit
465 return second.traverseByOffset(offset, limit, descending, leavesOnly, cb)
466}
467
468// Only used in testing...
469func (node *Node) lmd() *Node {
470 if node.height == 0 {
471 return node
472 }
473 return node.getLeftNode().lmd()
474}
475
476// Only used in testing...
477func (node *Node) rmd() *Node {
478 if node.height == 0 {
479 return node
480 }
481 return node.getRightNode().rmd()
482}
483
484func maxInt8(a, b int8) int8 {
485 if a > b {
486 return a
487 }
488 return b
489}
490
491// Equal compares two nodes for structural equality.
492// WARNING: This is an expensive operation that recursively traverses the entire tree structure.
493// It should only be used in tests or when absolutely necessary.
494func (node *Node) Equal(other *Node) bool {
495 // Handle nil cases
496 if node == nil || other == nil {
497 return node == other
498 }
499
500 // Compare node properties
501 if node.key != other.key ||
502 node.value != other.value ||
503 node.height != other.height ||
504 node.size != other.size {
505 return false
506 }
507
508 // Compare children
509 leftEqual := (node.leftNode == nil && other.leftNode == nil) ||
510 (node.leftNode != nil && other.leftNode != nil && node.leftNode.Equal(other.leftNode))
511 if !leftEqual {
512 return false
513 }
514
515 rightEqual := (node.rightNode == nil && other.rightNode == nil) ||
516 (node.rightNode != nil && other.rightNode != nil && node.rightNode.Equal(other.rightNode))
517 return rightEqual
518}
