github.com/wangyougui/gf/v2@v2.6.5/container/gtree/gtree_btree.go (about) 1 // Copyright GoFrame Author(https://goframe.org). All Rights Reserved. 2 // 3 // This Source Code Form is subject to the terms of the MIT License. 4 // If a copy of the MIT was not distributed with this file, 5 // You can obtain one at https://github.com/wangyougui/gf. 6 7 package gtree 8 9 import ( 10 "bytes" 11 "context" 12 "fmt" 13 "strings" 14 15 "github.com/wangyougui/gf/v2/container/gvar" 16 "github.com/wangyougui/gf/v2/internal/intlog" 17 "github.com/wangyougui/gf/v2/internal/json" 18 "github.com/wangyougui/gf/v2/internal/rwmutex" 19 "github.com/wangyougui/gf/v2/util/gconv" 20 ) 21 22 // BTree holds elements of the B-tree. 23 type BTree struct { 24 mu rwmutex.RWMutex 25 root *BTreeNode 26 comparator func(v1, v2 interface{}) int 27 size int // Total number of keys in the tree 28 m int // order (maximum number of children) 29 } 30 31 // BTreeNode is a single element within the tree. 32 type BTreeNode struct { 33 Parent *BTreeNode 34 Entries []*BTreeEntry // Contained keys in node 35 Children []*BTreeNode // Children nodes 36 } 37 38 // BTreeEntry represents the key-value pair contained within nodes. 39 type BTreeEntry struct { 40 Key interface{} 41 Value interface{} 42 } 43 44 // NewBTree instantiates a B-tree with `m` (maximum number of children) and a custom key comparator. 45 // The parameter `safe` is used to specify whether using tree in concurrent-safety, 46 // which is false in default. 47 // Note that the `m` must be greater or equal than 3, or else it panics. 48 func NewBTree(m int, comparator func(v1, v2 interface{}) int, safe ...bool) *BTree { 49 if m < 3 { 50 panic("Invalid order, should be at least 3") 51 } 52 return &BTree{ 53 comparator: comparator, 54 mu: rwmutex.Create(safe...), 55 m: m, 56 } 57 } 58 59 // NewBTreeFrom instantiates a B-tree with `m` (maximum number of children), a custom key comparator and data map. 60 // The parameter `safe` is used to specify whether using tree in concurrent-safety, 61 // which is false in default. 62 func NewBTreeFrom(m int, comparator func(v1, v2 interface{}) int, data map[interface{}]interface{}, safe ...bool) *BTree { 63 tree := NewBTree(m, comparator, safe...) 64 for k, v := range data { 65 tree.doSet(k, v) 66 } 67 return tree 68 } 69 70 // Clone returns a new tree with a copy of current tree. 71 func (tree *BTree) Clone() *BTree { 72 newTree := NewBTree(tree.m, tree.comparator, tree.mu.IsSafe()) 73 newTree.Sets(tree.Map()) 74 return newTree 75 } 76 77 // Set inserts key-value item into the tree. 78 func (tree *BTree) Set(key interface{}, value interface{}) { 79 tree.mu.Lock() 80 defer tree.mu.Unlock() 81 tree.doSet(key, value) 82 } 83 84 // doSet inserts key-value pair node into the tree. 85 // If key already exists, then its value is updated with the new value. 86 func (tree *BTree) doSet(key interface{}, value interface{}) { 87 entry := &BTreeEntry{Key: key, Value: value} 88 if tree.root == nil { 89 tree.root = &BTreeNode{Entries: []*BTreeEntry{entry}, Children: []*BTreeNode{}} 90 tree.size++ 91 return 92 } 93 94 if tree.insert(tree.root, entry) { 95 tree.size++ 96 } 97 } 98 99 // Sets batch sets key-values to the tree. 100 func (tree *BTree) Sets(data map[interface{}]interface{}) { 101 tree.mu.Lock() 102 defer tree.mu.Unlock() 103 for k, v := range data { 104 tree.doSet(k, v) 105 } 106 } 107 108 // Get searches the node in the tree by `key` and returns its value or nil if key is not found in tree. 109 func (tree *BTree) Get(key interface{}) (value interface{}) { 110 value, _ = tree.Search(key) 111 return 112 } 113 114 // doSetWithLockCheck checks whether value of the key exists with mutex.Lock, 115 // if not exists, set value to the map with given `key`, 116 // or else just return the existing value. 117 // 118 // When setting value, if `value` is type of <func() interface {}>, 119 // it will be executed with mutex.Lock of the hash map, 120 // and its return value will be set to the map with `key`. 121 // 122 // It returns value with given `key`. 123 func (tree *BTree) doSetWithLockCheck(key interface{}, value interface{}) interface{} { 124 tree.mu.Lock() 125 defer tree.mu.Unlock() 126 if entry := tree.doSearch(key); entry != nil { 127 return entry.Value 128 } 129 if f, ok := value.(func() interface{}); ok { 130 value = f() 131 } 132 if value != nil { 133 tree.doSet(key, value) 134 } 135 return value 136 } 137 138 // GetOrSet returns the value by key, 139 // or sets value with given `value` if it does not exist and then returns this value. 140 func (tree *BTree) GetOrSet(key interface{}, value interface{}) interface{} { 141 if v, ok := tree.Search(key); !ok { 142 return tree.doSetWithLockCheck(key, value) 143 } else { 144 return v 145 } 146 } 147 148 // GetOrSetFunc returns the value by key, 149 // or sets value with returned value of callback function `f` if it does not exist 150 // and then returns this value. 151 func (tree *BTree) GetOrSetFunc(key interface{}, f func() interface{}) interface{} { 152 if v, ok := tree.Search(key); !ok { 153 return tree.doSetWithLockCheck(key, f()) 154 } else { 155 return v 156 } 157 } 158 159 // GetOrSetFuncLock returns the value by key, 160 // or sets value with returned value of callback function `f` if it does not exist 161 // and then returns this value. 162 // 163 // GetOrSetFuncLock differs with GetOrSetFunc function is that it executes function `f` 164 // with mutex.Lock of the hash map. 165 func (tree *BTree) GetOrSetFuncLock(key interface{}, f func() interface{}) interface{} { 166 if v, ok := tree.Search(key); !ok { 167 return tree.doSetWithLockCheck(key, f) 168 } else { 169 return v 170 } 171 } 172 173 // GetVar returns a gvar.Var with the value by given `key`. 174 // The returned gvar.Var is un-concurrent safe. 175 func (tree *BTree) GetVar(key interface{}) *gvar.Var { 176 return gvar.New(tree.Get(key)) 177 } 178 179 // GetVarOrSet returns a gvar.Var with result from GetVarOrSet. 180 // The returned gvar.Var is un-concurrent safe. 181 func (tree *BTree) GetVarOrSet(key interface{}, value interface{}) *gvar.Var { 182 return gvar.New(tree.GetOrSet(key, value)) 183 } 184 185 // GetVarOrSetFunc returns a gvar.Var with result from GetOrSetFunc. 186 // The returned gvar.Var is un-concurrent safe. 187 func (tree *BTree) GetVarOrSetFunc(key interface{}, f func() interface{}) *gvar.Var { 188 return gvar.New(tree.GetOrSetFunc(key, f)) 189 } 190 191 // GetVarOrSetFuncLock returns a gvar.Var with result from GetOrSetFuncLock. 192 // The returned gvar.Var is un-concurrent safe. 193 func (tree *BTree) GetVarOrSetFuncLock(key interface{}, f func() interface{}) *gvar.Var { 194 return gvar.New(tree.GetOrSetFuncLock(key, f)) 195 } 196 197 // SetIfNotExist sets `value` to the map if the `key` does not exist, and then returns true. 198 // It returns false if `key` exists, and `value` would be ignored. 199 func (tree *BTree) SetIfNotExist(key interface{}, value interface{}) bool { 200 if !tree.Contains(key) { 201 tree.doSetWithLockCheck(key, value) 202 return true 203 } 204 return false 205 } 206 207 // SetIfNotExistFunc sets value with return value of callback function `f`, and then returns true. 208 // It returns false if `key` exists, and `value` would be ignored. 209 func (tree *BTree) SetIfNotExistFunc(key interface{}, f func() interface{}) bool { 210 if !tree.Contains(key) { 211 tree.doSetWithLockCheck(key, f()) 212 return true 213 } 214 return false 215 } 216 217 // SetIfNotExistFuncLock sets value with return value of callback function `f`, and then returns true. 218 // It returns false if `key` exists, and `value` would be ignored. 219 // 220 // SetIfNotExistFuncLock differs with SetIfNotExistFunc function is that 221 // it executes function `f` with mutex.Lock of the hash map. 222 func (tree *BTree) SetIfNotExistFuncLock(key interface{}, f func() interface{}) bool { 223 if !tree.Contains(key) { 224 tree.doSetWithLockCheck(key, f) 225 return true 226 } 227 return false 228 } 229 230 // Contains checks whether `key` exists in the tree. 231 func (tree *BTree) Contains(key interface{}) bool { 232 _, ok := tree.Search(key) 233 return ok 234 } 235 236 // doRemove removes the node from the tree by key. 237 // Key should adhere to the comparator's type assertion, otherwise method panics. 238 func (tree *BTree) doRemove(key interface{}) (value interface{}) { 239 node, index, found := tree.searchRecursively(tree.root, key) 240 if found { 241 value = node.Entries[index].Value 242 tree.delete(node, index) 243 tree.size-- 244 } 245 return 246 } 247 248 // Remove removes the node from the tree by `key`. 249 func (tree *BTree) Remove(key interface{}) (value interface{}) { 250 tree.mu.Lock() 251 defer tree.mu.Unlock() 252 return tree.doRemove(key) 253 } 254 255 // Removes batch deletes values of the tree by `keys`. 256 func (tree *BTree) Removes(keys []interface{}) { 257 tree.mu.Lock() 258 defer tree.mu.Unlock() 259 for _, key := range keys { 260 tree.doRemove(key) 261 } 262 } 263 264 // IsEmpty returns true if tree does not contain any nodes 265 func (tree *BTree) IsEmpty() bool { 266 return tree.Size() == 0 267 } 268 269 // Size returns number of nodes in the tree. 270 func (tree *BTree) Size() int { 271 tree.mu.RLock() 272 defer tree.mu.RUnlock() 273 return tree.size 274 } 275 276 // Keys returns all keys in asc order. 277 func (tree *BTree) Keys() []interface{} { 278 keys := make([]interface{}, tree.Size()) 279 index := 0 280 tree.IteratorAsc(func(key, value interface{}) bool { 281 keys[index] = key 282 index++ 283 return true 284 }) 285 return keys 286 } 287 288 // Values returns all values in asc order based on the key. 289 func (tree *BTree) Values() []interface{} { 290 values := make([]interface{}, tree.Size()) 291 index := 0 292 tree.IteratorAsc(func(key, value interface{}) bool { 293 values[index] = value 294 index++ 295 return true 296 }) 297 return values 298 } 299 300 // Map returns all key-value items as map. 301 func (tree *BTree) Map() map[interface{}]interface{} { 302 m := make(map[interface{}]interface{}, tree.Size()) 303 tree.IteratorAsc(func(key, value interface{}) bool { 304 m[key] = value 305 return true 306 }) 307 return m 308 } 309 310 // MapStrAny returns all key-value items as map[string]interface{}. 311 func (tree *BTree) MapStrAny() map[string]interface{} { 312 m := make(map[string]interface{}, tree.Size()) 313 tree.IteratorAsc(func(key, value interface{}) bool { 314 m[gconv.String(key)] = value 315 return true 316 }) 317 return m 318 } 319 320 // Clear removes all nodes from the tree. 321 func (tree *BTree) Clear() { 322 tree.mu.Lock() 323 defer tree.mu.Unlock() 324 tree.root = nil 325 tree.size = 0 326 } 327 328 // Replace the data of the tree with given `data`. 329 func (tree *BTree) Replace(data map[interface{}]interface{}) { 330 tree.mu.Lock() 331 defer tree.mu.Unlock() 332 tree.root = nil 333 tree.size = 0 334 for k, v := range data { 335 tree.doSet(k, v) 336 } 337 } 338 339 // Height returns the height of the tree. 340 func (tree *BTree) Height() int { 341 tree.mu.RLock() 342 defer tree.mu.RUnlock() 343 return tree.root.height() 344 } 345 346 // Left returns the left-most (min) entry or nil if tree is empty. 347 func (tree *BTree) Left() *BTreeEntry { 348 tree.mu.RLock() 349 defer tree.mu.RUnlock() 350 node := tree.left(tree.root) 351 if node != nil { 352 return node.Entries[0] 353 } 354 return nil 355 } 356 357 // Right returns the right-most (max) entry or nil if tree is empty. 358 func (tree *BTree) Right() *BTreeEntry { 359 tree.mu.RLock() 360 defer tree.mu.RUnlock() 361 node := tree.right(tree.root) 362 if node != nil { 363 return node.Entries[len(node.Entries)-1] 364 } 365 return nil 366 } 367 368 // String returns a string representation of container (for debugging purposes) 369 func (tree *BTree) String() string { 370 if tree == nil { 371 return "" 372 } 373 tree.mu.RLock() 374 defer tree.mu.RUnlock() 375 var buffer bytes.Buffer 376 if tree.size != 0 { 377 tree.output(&buffer, tree.root, 0, true) 378 } 379 return buffer.String() 380 } 381 382 // Search searches the tree with given `key`. 383 // Second return parameter `found` is true if key was found, otherwise false. 384 func (tree *BTree) Search(key interface{}) (value interface{}, found bool) { 385 tree.mu.RLock() 386 defer tree.mu.RUnlock() 387 node, index, found := tree.searchRecursively(tree.root, key) 388 if found { 389 return node.Entries[index].Value, true 390 } 391 return nil, false 392 } 393 394 // Search searches the tree with given `key` without mutex. 395 // It returns the entry if found or otherwise nil. 396 func (tree *BTree) doSearch(key interface{}) *BTreeEntry { 397 node, index, found := tree.searchRecursively(tree.root, key) 398 if found { 399 return node.Entries[index] 400 } 401 return nil 402 } 403 404 // Print prints the tree to stdout. 405 func (tree *BTree) Print() { 406 fmt.Println(tree.String()) 407 } 408 409 // Iterator is alias of IteratorAsc. 410 func (tree *BTree) Iterator(f func(key, value interface{}) bool) { 411 tree.IteratorAsc(f) 412 } 413 414 // IteratorFrom is alias of IteratorAscFrom. 415 func (tree *BTree) IteratorFrom(key interface{}, match bool, f func(key, value interface{}) bool) { 416 tree.IteratorAscFrom(key, match, f) 417 } 418 419 // IteratorAsc iterates the tree readonly in ascending order with given callback function `f`. 420 // If `f` returns true, then it continues iterating; or false to stop. 421 func (tree *BTree) IteratorAsc(f func(key, value interface{}) bool) { 422 tree.mu.RLock() 423 defer tree.mu.RUnlock() 424 node := tree.left(tree.root) 425 if node == nil { 426 return 427 } 428 tree.doIteratorAsc(node, node.Entries[0], 0, f) 429 } 430 431 // IteratorAscFrom iterates the tree readonly in ascending order with given callback function `f`. 432 // The parameter `key` specifies the start entry for iterating. The `match` specifies whether 433 // starting iterating if the `key` is fully matched, or else using index searching iterating. 434 // If `f` returns true, then it continues iterating; or false to stop. 435 func (tree *BTree) IteratorAscFrom(key interface{}, match bool, f func(key, value interface{}) bool) { 436 tree.mu.RLock() 437 defer tree.mu.RUnlock() 438 node, index, found := tree.searchRecursively(tree.root, key) 439 if match { 440 if found { 441 tree.doIteratorAsc(node, node.Entries[index], index, f) 442 } 443 } else { 444 if index >= 0 && index < len(node.Entries) { 445 tree.doIteratorAsc(node, node.Entries[index], index, f) 446 } 447 } 448 } 449 450 func (tree *BTree) doIteratorAsc(node *BTreeNode, entry *BTreeEntry, index int, f func(key, value interface{}) bool) { 451 first := true 452 loop: 453 if entry == nil { 454 return 455 } 456 if !f(entry.Key, entry.Value) { 457 return 458 } 459 // Find current entry position in current node 460 if !first { 461 index, _ = tree.search(node, entry.Key) 462 } else { 463 first = false 464 } 465 // Try to go down to the child right of the current entry 466 if index+1 < len(node.Children) { 467 node = node.Children[index+1] 468 // Try to go down to the child left of the current node 469 for len(node.Children) > 0 { 470 node = node.Children[0] 471 } 472 // Return the left-most entry 473 entry = node.Entries[0] 474 goto loop 475 } 476 // Above assures that we have reached a leaf node, so return the next entry in current node (if any) 477 if index+1 < len(node.Entries) { 478 entry = node.Entries[index+1] 479 goto loop 480 } 481 // Reached leaf node and there are no entries to the right of the current entry, so go up to the parent 482 for node.Parent != nil { 483 node = node.Parent 484 // Find next entry position in current node (note: search returns the first equal or bigger than entry) 485 index, _ = tree.search(node, entry.Key) 486 // Check that there is a next entry position in current node 487 if index < len(node.Entries) { 488 entry = node.Entries[index] 489 goto loop 490 } 491 } 492 } 493 494 // IteratorDesc iterates the tree readonly in descending order with given callback function `f`. 495 // If `f` returns true, then it continues iterating; or false to stop. 496 func (tree *BTree) IteratorDesc(f func(key, value interface{}) bool) { 497 tree.mu.RLock() 498 defer tree.mu.RUnlock() 499 node := tree.right(tree.root) 500 if node == nil { 501 return 502 } 503 index := len(node.Entries) - 1 504 entry := node.Entries[index] 505 tree.doIteratorDesc(node, entry, index, f) 506 } 507 508 // IteratorDescFrom iterates the tree readonly in descending order with given callback function `f`. 509 // The parameter `key` specifies the start entry for iterating. The `match` specifies whether 510 // starting iterating if the `key` is fully matched, or else using index searching iterating. 511 // If `f` returns true, then it continues iterating; or false to stop. 512 func (tree *BTree) IteratorDescFrom(key interface{}, match bool, f func(key, value interface{}) bool) { 513 tree.mu.RLock() 514 defer tree.mu.RUnlock() 515 node, index, found := tree.searchRecursively(tree.root, key) 516 if match { 517 if found { 518 tree.doIteratorDesc(node, node.Entries[index], index, f) 519 } 520 } else { 521 if index >= 0 && index < len(node.Entries) { 522 tree.doIteratorDesc(node, node.Entries[index], index, f) 523 } 524 } 525 } 526 527 // IteratorDesc iterates the tree readonly in descending order with given callback function `f`. 528 // If `f` returns true, then it continues iterating; or false to stop. 529 func (tree *BTree) doIteratorDesc(node *BTreeNode, entry *BTreeEntry, index int, f func(key, value interface{}) bool) { 530 first := true 531 loop: 532 if entry == nil { 533 return 534 } 535 if !f(entry.Key, entry.Value) { 536 return 537 } 538 // Find current entry position in current node 539 if !first { 540 index, _ = tree.search(node, entry.Key) 541 } else { 542 first = false 543 } 544 // Try to go down to the child left of the current entry 545 if index < len(node.Children) { 546 node = node.Children[index] 547 // Try to go down to the child right of the current node 548 for len(node.Children) > 0 { 549 node = node.Children[len(node.Children)-1] 550 } 551 // Return the right-most entry 552 entry = node.Entries[len(node.Entries)-1] 553 goto loop 554 } 555 // Above assures that we have reached a leaf node, so return the previous entry in current node (if any) 556 if index-1 >= 0 { 557 entry = node.Entries[index-1] 558 goto loop 559 } 560 561 // Reached leaf node and there are no entries to the left of the current entry, so go up to the parent 562 for node.Parent != nil { 563 node = node.Parent 564 // Find previous entry position in current node (note: search returns the first equal or bigger than entry) 565 index, _ = tree.search(node, entry.Key) 566 // Check that there is a previous entry position in current node 567 if index-1 >= 0 { 568 entry = node.Entries[index-1] 569 goto loop 570 } 571 } 572 } 573 574 func (tree *BTree) output(buffer *bytes.Buffer, node *BTreeNode, level int, isTail bool) { 575 for e := 0; e < len(node.Entries)+1; e++ { 576 if e < len(node.Children) { 577 tree.output(buffer, node.Children[e], level+1, true) 578 } 579 if e < len(node.Entries) { 580 if _, err := buffer.WriteString(strings.Repeat(" ", level)); err != nil { 581 intlog.Errorf(context.TODO(), `%+v`, err) 582 } 583 if _, err := buffer.WriteString(fmt.Sprintf("%v", node.Entries[e].Key) + "\n"); err != nil { 584 intlog.Errorf(context.TODO(), `%+v`, err) 585 } 586 } 587 } 588 } 589 590 func (node *BTreeNode) height() int { 591 h := 0 592 n := node 593 for ; n != nil; n = n.Children[0] { 594 h++ 595 if len(n.Children) == 0 { 596 break 597 } 598 } 599 return h 600 } 601 602 func (tree *BTree) isLeaf(node *BTreeNode) bool { 603 return len(node.Children) == 0 604 } 605 606 // func (tree *BTree) isFull(node *BTreeNode) bool { 607 // return len(node.Entries) == tree.maxEntries() 608 // } 609 610 func (tree *BTree) shouldSplit(node *BTreeNode) bool { 611 return len(node.Entries) > tree.maxEntries() 612 } 613 614 func (tree *BTree) maxChildren() int { 615 return tree.m 616 } 617 618 func (tree *BTree) minChildren() int { 619 return (tree.m + 1) / 2 // ceil(m/2) 620 } 621 622 func (tree *BTree) maxEntries() int { 623 return tree.maxChildren() - 1 624 } 625 626 func (tree *BTree) minEntries() int { 627 return tree.minChildren() - 1 628 } 629 630 func (tree *BTree) middle() int { 631 // "-1" to favor right nodes to have more keys when splitting 632 return (tree.m - 1) / 2 633 } 634 635 // search does search only within the single node among its entries 636 func (tree *BTree) search(node *BTreeNode, key interface{}) (index int, found bool) { 637 low, mid, high := 0, 0, len(node.Entries)-1 638 for low <= high { 639 mid = low + (high-low)/2 640 compare := tree.getComparator()(key, node.Entries[mid].Key) 641 switch { 642 case compare > 0: 643 low = mid + 1 644 case compare < 0: 645 high = mid - 1 646 case compare == 0: 647 return mid, true 648 } 649 } 650 return low, false 651 } 652 653 // searchRecursively searches recursively down the tree starting at the startNode 654 func (tree *BTree) searchRecursively(startNode *BTreeNode, key interface{}) (node *BTreeNode, index int, found bool) { 655 if tree.size == 0 { 656 return nil, -1, false 657 } 658 node = startNode 659 for { 660 index, found = tree.search(node, key) 661 if found { 662 return node, index, true 663 } 664 if tree.isLeaf(node) { 665 return node, index, false 666 } 667 node = node.Children[index] 668 } 669 } 670 671 func (tree *BTree) insert(node *BTreeNode, entry *BTreeEntry) (inserted bool) { 672 if tree.isLeaf(node) { 673 return tree.insertIntoLeaf(node, entry) 674 } 675 return tree.insertIntoInternal(node, entry) 676 } 677 678 func (tree *BTree) insertIntoLeaf(node *BTreeNode, entry *BTreeEntry) (inserted bool) { 679 insertPosition, found := tree.search(node, entry.Key) 680 if found { 681 node.Entries[insertPosition] = entry 682 return false 683 } 684 // Insert entry's key in the middle of the node 685 node.Entries = append(node.Entries, nil) 686 copy(node.Entries[insertPosition+1:], node.Entries[insertPosition:]) 687 node.Entries[insertPosition] = entry 688 tree.split(node) 689 return true 690 } 691 692 func (tree *BTree) insertIntoInternal(node *BTreeNode, entry *BTreeEntry) (inserted bool) { 693 insertPosition, found := tree.search(node, entry.Key) 694 if found { 695 node.Entries[insertPosition] = entry 696 return false 697 } 698 return tree.insert(node.Children[insertPosition], entry) 699 } 700 701 func (tree *BTree) split(node *BTreeNode) { 702 if !tree.shouldSplit(node) { 703 return 704 } 705 706 if node == tree.root { 707 tree.splitRoot() 708 return 709 } 710 711 tree.splitNonRoot(node) 712 } 713 714 func (tree *BTree) splitNonRoot(node *BTreeNode) { 715 middle := tree.middle() 716 parent := node.Parent 717 718 left := &BTreeNode{Entries: append([]*BTreeEntry(nil), node.Entries[:middle]...), Parent: parent} 719 right := &BTreeNode{Entries: append([]*BTreeEntry(nil), node.Entries[middle+1:]...), Parent: parent} 720 721 // Move children from the node to be split into left and right nodes 722 if !tree.isLeaf(node) { 723 left.Children = append([]*BTreeNode(nil), node.Children[:middle+1]...) 724 right.Children = append([]*BTreeNode(nil), node.Children[middle+1:]...) 725 setParent(left.Children, left) 726 setParent(right.Children, right) 727 } 728 729 insertPosition, _ := tree.search(parent, node.Entries[middle].Key) 730 731 // Insert middle key into parent 732 parent.Entries = append(parent.Entries, nil) 733 copy(parent.Entries[insertPosition+1:], parent.Entries[insertPosition:]) 734 parent.Entries[insertPosition] = node.Entries[middle] 735 736 // Set child left of inserted key in parent to the created left node 737 parent.Children[insertPosition] = left 738 739 // Set child right of inserted key in parent to the created right node 740 parent.Children = append(parent.Children, nil) 741 copy(parent.Children[insertPosition+2:], parent.Children[insertPosition+1:]) 742 parent.Children[insertPosition+1] = right 743 744 tree.split(parent) 745 } 746 747 func (tree *BTree) splitRoot() { 748 middle := tree.middle() 749 left := &BTreeNode{Entries: append([]*BTreeEntry(nil), tree.root.Entries[:middle]...)} 750 right := &BTreeNode{Entries: append([]*BTreeEntry(nil), tree.root.Entries[middle+1:]...)} 751 752 // Move children from the node to be split into left and right nodes 753 if !tree.isLeaf(tree.root) { 754 left.Children = append([]*BTreeNode(nil), tree.root.Children[:middle+1]...) 755 right.Children = append([]*BTreeNode(nil), tree.root.Children[middle+1:]...) 756 setParent(left.Children, left) 757 setParent(right.Children, right) 758 } 759 760 // Root is a node with one entry and two children (left and right) 761 newRoot := &BTreeNode{ 762 Entries: []*BTreeEntry{tree.root.Entries[middle]}, 763 Children: []*BTreeNode{left, right}, 764 } 765 766 left.Parent = newRoot 767 right.Parent = newRoot 768 tree.root = newRoot 769 } 770 771 func setParent(nodes []*BTreeNode, parent *BTreeNode) { 772 for _, node := range nodes { 773 node.Parent = parent 774 } 775 } 776 777 func (tree *BTree) left(node *BTreeNode) *BTreeNode { 778 if tree.size == 0 { 779 return nil 780 } 781 current := node 782 for { 783 if tree.isLeaf(current) { 784 return current 785 } 786 current = current.Children[0] 787 } 788 } 789 790 func (tree *BTree) right(node *BTreeNode) *BTreeNode { 791 if tree.size == 0 { 792 return nil 793 } 794 current := node 795 for { 796 if tree.isLeaf(current) { 797 return current 798 } 799 current = current.Children[len(current.Children)-1] 800 } 801 } 802 803 // leftSibling returns the node's left sibling and child index (in parent) if it exists, otherwise (nil,-1) 804 // key is any of keys in node (could even be deleted). 805 func (tree *BTree) leftSibling(node *BTreeNode, key interface{}) (*BTreeNode, int) { 806 if node.Parent != nil { 807 index, _ := tree.search(node.Parent, key) 808 index-- 809 if index >= 0 && index < len(node.Parent.Children) { 810 return node.Parent.Children[index], index 811 } 812 } 813 return nil, -1 814 } 815 816 // rightSibling returns the node's right sibling and child index (in parent) if it exists, otherwise (nil,-1) 817 // key is any of keys in node (could even be deleted). 818 func (tree *BTree) rightSibling(node *BTreeNode, key interface{}) (*BTreeNode, int) { 819 if node.Parent != nil { 820 index, _ := tree.search(node.Parent, key) 821 index++ 822 if index < len(node.Parent.Children) { 823 return node.Parent.Children[index], index 824 } 825 } 826 return nil, -1 827 } 828 829 // delete deletes an entry in node at entries' index 830 // ref.: https://en.wikipedia.org/wiki/B-tree#Deletion 831 func (tree *BTree) delete(node *BTreeNode, index int) { 832 // deleting from a leaf node 833 if tree.isLeaf(node) { 834 deletedKey := node.Entries[index].Key 835 tree.deleteEntry(node, index) 836 tree.reBalance(node, deletedKey) 837 if len(tree.root.Entries) == 0 { 838 tree.root = nil 839 } 840 return 841 } 842 843 // deleting from an internal node 844 leftLargestNode := tree.right(node.Children[index]) // largest node in the left sub-tree (assumed to exist) 845 leftLargestEntryIndex := len(leftLargestNode.Entries) - 1 846 node.Entries[index] = leftLargestNode.Entries[leftLargestEntryIndex] 847 deletedKey := leftLargestNode.Entries[leftLargestEntryIndex].Key 848 tree.deleteEntry(leftLargestNode, leftLargestEntryIndex) 849 tree.reBalance(leftLargestNode, deletedKey) 850 } 851 852 // reBalance reBalances the tree after deletion if necessary and returns true, otherwise false. 853 // Note that we first delete the entry and then call reBalance, thus the passed deleted key as reference. 854 func (tree *BTree) reBalance(node *BTreeNode, deletedKey interface{}) { 855 // check if re-balancing is needed 856 if node == nil || len(node.Entries) >= tree.minEntries() { 857 return 858 } 859 860 // try to borrow from left sibling 861 leftSibling, leftSiblingIndex := tree.leftSibling(node, deletedKey) 862 if leftSibling != nil && len(leftSibling.Entries) > tree.minEntries() { 863 // rotate right 864 node.Entries = append([]*BTreeEntry{node.Parent.Entries[leftSiblingIndex]}, node.Entries...) // prepend parent's separator entry to node's entries 865 node.Parent.Entries[leftSiblingIndex] = leftSibling.Entries[len(leftSibling.Entries)-1] 866 tree.deleteEntry(leftSibling, len(leftSibling.Entries)-1) 867 if !tree.isLeaf(leftSibling) { 868 leftSiblingRightMostChild := leftSibling.Children[len(leftSibling.Children)-1] 869 leftSiblingRightMostChild.Parent = node 870 node.Children = append([]*BTreeNode{leftSiblingRightMostChild}, node.Children...) 871 tree.deleteChild(leftSibling, len(leftSibling.Children)-1) 872 } 873 return 874 } 875 876 // try to borrow from right sibling 877 rightSibling, rightSiblingIndex := tree.rightSibling(node, deletedKey) 878 if rightSibling != nil && len(rightSibling.Entries) > tree.minEntries() { 879 // rotate left 880 node.Entries = append(node.Entries, node.Parent.Entries[rightSiblingIndex-1]) // append parent's separator entry to node's entries 881 node.Parent.Entries[rightSiblingIndex-1] = rightSibling.Entries[0] 882 tree.deleteEntry(rightSibling, 0) 883 if !tree.isLeaf(rightSibling) { 884 rightSiblingLeftMostChild := rightSibling.Children[0] 885 rightSiblingLeftMostChild.Parent = node 886 node.Children = append(node.Children, rightSiblingLeftMostChild) 887 tree.deleteChild(rightSibling, 0) 888 } 889 return 890 } 891 892 // merge with siblings 893 if rightSibling != nil { 894 // merge with right sibling 895 node.Entries = append(node.Entries, node.Parent.Entries[rightSiblingIndex-1]) 896 node.Entries = append(node.Entries, rightSibling.Entries...) 897 deletedKey = node.Parent.Entries[rightSiblingIndex-1].Key 898 tree.deleteEntry(node.Parent, rightSiblingIndex-1) 899 tree.appendChildren(node.Parent.Children[rightSiblingIndex], node) 900 tree.deleteChild(node.Parent, rightSiblingIndex) 901 } else if leftSibling != nil { 902 // merge with left sibling 903 entries := append([]*BTreeEntry(nil), leftSibling.Entries...) 904 entries = append(entries, node.Parent.Entries[leftSiblingIndex]) 905 node.Entries = append(entries, node.Entries...) 906 deletedKey = node.Parent.Entries[leftSiblingIndex].Key 907 tree.deleteEntry(node.Parent, leftSiblingIndex) 908 tree.prependChildren(node.Parent.Children[leftSiblingIndex], node) 909 tree.deleteChild(node.Parent, leftSiblingIndex) 910 } 911 912 // make the merged node the root if its parent was the root and the root is empty 913 if node.Parent == tree.root && len(tree.root.Entries) == 0 { 914 tree.root = node 915 node.Parent = nil 916 return 917 } 918 919 // parent might be underflow, so try to reBalance if necessary 920 tree.reBalance(node.Parent, deletedKey) 921 } 922 923 func (tree *BTree) prependChildren(fromNode *BTreeNode, toNode *BTreeNode) { 924 children := append([]*BTreeNode(nil), fromNode.Children...) 925 toNode.Children = append(children, toNode.Children...) 926 setParent(fromNode.Children, toNode) 927 } 928 929 func (tree *BTree) appendChildren(fromNode *BTreeNode, toNode *BTreeNode) { 930 toNode.Children = append(toNode.Children, fromNode.Children...) 931 setParent(fromNode.Children, toNode) 932 } 933 934 func (tree *BTree) deleteEntry(node *BTreeNode, index int) { 935 copy(node.Entries[index:], node.Entries[index+1:]) 936 node.Entries[len(node.Entries)-1] = nil 937 node.Entries = node.Entries[:len(node.Entries)-1] 938 } 939 940 func (tree *BTree) deleteChild(node *BTreeNode, index int) { 941 if index >= len(node.Children) { 942 return 943 } 944 copy(node.Children[index:], node.Children[index+1:]) 945 node.Children[len(node.Children)-1] = nil 946 node.Children = node.Children[:len(node.Children)-1] 947 } 948 949 // MarshalJSON implements the interface MarshalJSON for json.Marshal. 950 func (tree BTree) MarshalJSON() (jsonBytes []byte, err error) { 951 if tree.root == nil { 952 return []byte("null"), nil 953 } 954 buffer := bytes.NewBuffer(nil) 955 buffer.WriteByte('{') 956 tree.Iterator(func(key, value interface{}) bool { 957 valueBytes, valueJsonErr := json.Marshal(value) 958 if valueJsonErr != nil { 959 err = valueJsonErr 960 return false 961 } 962 if buffer.Len() > 1 { 963 buffer.WriteByte(',') 964 } 965 buffer.WriteString(fmt.Sprintf(`"%v":%s`, key, valueBytes)) 966 return true 967 }) 968 buffer.WriteByte('}') 969 return buffer.Bytes(), nil 970 } 971 972 // getComparator returns the comparator if it's previously set, 973 // or else it panics. 974 func (tree *BTree) getComparator() func(a, b interface{}) int { 975 if tree.comparator == nil { 976 panic("comparator is missing for tree") 977 } 978 return tree.comparator 979 }