1package cow23//----------------------------------------4// Node56// 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}1516// 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}2526// Size returns the size of the subtree rooted at the node.27func (node *Node) Size() int {28 if node == nil {29 return 030 }31 return node.size32}3334// IsLeaf checks if the node is a leaf node (has no children).35func (node *Node) IsLeaf() bool {36 return node.height == 037}3839// Key returns the key of the node.40func (node *Node) Key() string {41 return node.key42}4344// Value returns the value of the node.45func (node *Node) Value() any {46 return node.value47}4849// _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}6263// 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 false67 }68 if node.key == key {69 return true70 }71 if node.height == 0 {72 return false73 }74 if key < node.key {75 return node.getLeftNode().Has(key)76 }77 return node.getRightNode().Has(key)78}7980// Get searches for a node with the given key in the subtree rooted at the node81// 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, false85 }8687 if node.height == 0 {88 if node.key == key {89 return 0, node.value, true90 }91 if node.key < key {92 return 1, nil, false93 }94 return 0, nil, false95 }9697 if key < node.key {98 return node.getLeftNode().Get(key)99 }100101 rightNode := node.getRightNode()102 index, value, exists = rightNode.Get(key)103 index += node.size - rightNode.size104 return index, value, exists105}106107// GetByIndex retrieves the key-value pair of the node at the given index108// 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.value113 }114 panic("GetByIndex asked for invalid index")115 }116 // TODO: could improve this by storing the sizes117 leftNode := node.getLeftNode()118 if index < leftNode.size {119 return leftNode.GetByIndex(index)120 }121 return node.getRightNode().GetByIndex(index - leftNode.size)122}123124// 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), false131 }132133 // Always create a new node for leaf nodes134 if node.height == 0 {135 return node.setLeaf(key, value)136 }137138 // Copy the node before modifying139 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 }145146 if !updated {147 newNode.calcHeightAndSize()148 return newNode.balance(), updated149 }150151 return newNode, updated152}153154// 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), true159 }160161 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 }, false169 }170171 return &Node{172 key: key,173 height: 1,174 size: 2,175 leftNode: node,176 rightNode: NewNode(key, value),177 }, false178}179180// 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, false188 }189 if node.height == 0 {190 if key == node.key {191 return nil, "", node.value, true192 }193 return node, "", nil, false194 }195 if key < node.key {196 var newLeftNode *Node197 newLeftNode, newKey, value, removed = node.getLeftNode().Remove(key)198 if !removed {199 return node, "", value, false200 }201 if newLeftNode == nil { // left node held value, was removed202 return node.rightNode, node.key, value, true203 }204 node = node._copy()205 node.leftNode = newLeftNode206 node.calcHeightAndSize()207 node = node.balance()208 return node, newKey, value, true209 }210211 var newRightNode *Node212 newRightNode, newKey, value, removed = node.getRightNode().Remove(key)213 if !removed {214 return node, "", value, false215 }216 if newRightNode == nil { // right node held value, was removed217 return node.leftNode, "", value, true218 }219 node = node._copy()220 node.rightNode = newRightNode221 if newKey != "" {222 node.key = newKey223 }224 node.calcHeightAndSize()225 node = node.balance()226 return node, "", value, true227}228229// getLeftNode returns the left child of the node.230func (node *Node) getLeftNode() *Node {231 return node.leftNode232}233234// getRightNode returns the right child of the node.235func (node *Node) getRightNode() *Node {236 return node.rightNode237}238239// rotateRight performs a right rotation on the node and returns the new root.240// NOTE: overwrites node241// TODO: optimize balance & rotate242func (node *Node) rotateRight() *Node {243 node = node._copy()244 l := node.getLeftNode()245 _l := l._copy()246247 _lrCached := _l.rightNode248 _l.rightNode = node249 node.leftNode = _lrCached250251 node.calcHeightAndSize()252 _l.calcHeightAndSize()253254 return _l255}256257// rotateLeft performs a left rotation on the node and returns the new root.258// NOTE: overwrites node259// TODO: optimize balance & rotate260func (node *Node) rotateLeft() *Node {261 node = node._copy()262 r := node.getRightNode()263 _r := r._copy()264265 _rlCached := _r.leftNode266 _r.leftNode = node267 node.rightNode = _rlCached268269 node.calcHeightAndSize()270 _r.calcHeightAndSize()271272 return _r273}274275// calcHeightAndSize updates the height and size of the node based on its children.276// NOTE: mutates height and size277func (node *Node) calcHeightAndSize() {278 node.height = maxInt8(node.getLeftNode().height, node.getRightNode().height) + 1279 node.size = node.getLeftNode().size + node.getRightNode().size280}281282// calcBalance calculates the balance factor of the node.283func (node *Node) calcBalance() int {284 return int(node.getLeftNode().height) - int(node.getRightNode().height)285}286287// balance balances the subtree rooted at the node and returns the new root.288// NOTE: assumes that node can be modified289// TODO: optimize balance & rotate290func (node *Node) balance() (newSelf *Node) {291 balance := node.calcBalance()292 if balance >= -1 {293 return node294 }295 if balance > 1 {296 if node.getLeftNode().calcBalance() >= 0 {297 // Left Left Case298 return node.rotateRight()299 }300 // Left Right Case301 left := node.getLeftNode()302 node.leftNode = left.rotateLeft()303 return node.rotateRight()304 }305306 if node.getRightNode().calcBalance() <= 0 {307 // Right Right Case308 return node.rotateLeft()309 }310311 // Right Left Case312 right := node.getRightNode()313 node.rightNode = right.rotateRight()314 return node.rotateLeft()315}316317// 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}321322// 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}326327// 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 false337 }338 afterStart := (start == "" || start < node.key)339 startOrAfter := (start == "" || start <= node.key)340 beforeEnd := false341 if ascending {342 beforeEnd = (end == "" || node.key < end)343 } else {344 beforeEnd = (end == "" || node.key <= end)345 }346347 // Run callback per inner/leaf node.348 stop := false349 if (!node.IsLeaf() && !leavesOnly) ||350 (node.IsLeaf() && startOrAfter && beforeEnd) {351 stop = cb(node)352 if stop {353 return stop354 }355 }356 if node.IsLeaf() {357 return stop358 }359360 if ascending {361 // check lower nodes, then higher362 if afterStart {363 stop = node.getLeftNode().TraverseInRange(start, end, ascending, leavesOnly, cb)364 }365 if stop {366 return stop367 }368 if beforeEnd {369 stop = node.getRightNode().TraverseInRange(start, end, ascending, leavesOnly, cb)370 }371 } else {372 // check the higher nodes first373 if beforeEnd {374 stop = node.getRightNode().TraverseInRange(start, end, ascending, leavesOnly, cb)375 }376 if stop {377 return stop378 }379 if afterStart {380 stop = node.getLeftNode().TraverseInRange(start, end, ascending, leavesOnly, cb)381 }382 }383384 return stop385}386387// 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 false392 }393394 // fast paths. these happen only if TraverseByOffset is called directly on a leaf.395 if limit <= 0 || offset >= node.size {396 return false397 }398 if node.IsLeaf() {399 if offset > 0 {400 return false401 }402 return cb(node)403 }404405 // go to the actual recursive function.406 return node.traverseByOffset(offset, limit, descending, leavesOnly, cb)407}408409// 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 true417 }418 }419 first, second := node.getLeftNode(), node.getRightNode()420 if descending {421 first, second = second, first422 }423 if first.IsLeaf() {424 // either run or skip, based on offset425 if offset > 0 {426 offset--427 } else {428 cb(first)429 limit--430 if limit <= 0 {431 return false432 }433 }434 } else {435 // possible cases:436 // 1 the offset given skips the first node entirely437 // 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 // 1441 } else {442 if first.traverseByOffset(offset, limit, descending, leavesOnly, cb) {443 return true444 }445 // number of leaves which could actually be called from inside446 delta := first.size - offset447 offset = 0448 if delta >= limit {449 return true // 3450 }451 limit -= delta // 2452 }453 }454455 // 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 in457 // this second node that we include.458459 // => 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 this464 // function with the updated offset and limit465 return second.traverseByOffset(offset, limit, descending, leavesOnly, cb)466}467468// Only used in testing...469func (node *Node) lmd() *Node {470 if node.height == 0 {471 return node472 }473 return node.getLeftNode().lmd()474}475476// Only used in testing...477func (node *Node) rmd() *Node {478 if node.height == 0 {479 return node480 }481 return node.getRightNode().rmd()482}483484func maxInt8(a, b int8) int8 {485 if a > b {486 return a487 }488 return b489}490491// 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 cases496 if node == nil || other == nil {497 return node == other498 }499500 // Compare node properties501 if node.key != other.key ||502 node.value != other.value ||503 node.height != other.height ||504 node.size != other.size {505 return false506 }507508 // Compare children509 leftEqual := (node.leftNode == nil && other.leftNode == nil) ||510 (node.leftNode != nil && other.leftNode != nil && node.leftNode.Equal(other.leftNode))511 if !leftEqual {512 return false513 }514515 rightEqual := (node.rightNode == nil && other.rightNode == nil) ||516 (node.rightNode != nil && other.rightNode != nil && node.rightNode.Equal(other.rightNode))517 return rightEqual518}519Signatures reconstructed verbatim from vm/qfuncs — interface params keep their inline definitions.