1 /* 2 * Copyright 2003-2007, Ingo Weinhold <bonefish@cs.tu-berlin.de>. 3 * Distributed under the terms of the MIT License. 4 */ 5 #ifndef _AVL_TREE_MAP_H 6 #define _AVL_TREE_MAP_H 7 8 9 #include <OS.h> 10 #include <util/kernel_cpp.h> 11 12 #include <util/MallocFreeAllocator.h> 13 14 15 // maximal height of a tree 16 static const int kMaxAVLTreeHeight = 32; 17 18 class AVLTreeIterator; 19 20 21 // AVLTreeNode 22 struct AVLTreeNode { 23 AVLTreeNode* parent; 24 AVLTreeNode* left; 25 AVLTreeNode* right; 26 int balance_factor; 27 }; 28 29 30 // AVLTreeCompare 31 class AVLTreeCompare { 32 public: 33 virtual ~AVLTreeCompare(); 34 35 virtual int CompareKeyNode(const void* key, 36 const AVLTreeNode* node) = 0; 37 virtual int CompareNodes(const AVLTreeNode* node1, 38 const AVLTreeNode* node2) = 0; 39 }; 40 41 42 // AVLTree 43 class AVLTree { 44 public: 45 AVLTree(AVLTreeCompare* compare); 46 ~AVLTree(); 47 48 inline int Count() const { return fNodeCount; } 49 inline bool IsEmpty() const { return (fNodeCount == 0); } 50 void MakeEmpty(); 51 52 inline AVLTreeNode* Root() const { return fRoot; } 53 54 AVLTreeNode* LeftMost(AVLTreeNode* node) const; 55 AVLTreeNode* RightMost(AVLTreeNode* node) const; 56 57 AVLTreeNode* Previous(AVLTreeNode* node) const; 58 AVLTreeNode* Next(AVLTreeNode* node) const; 59 60 inline AVLTreeIterator GetIterator() const; 61 inline AVLTreeIterator GetIterator(AVLTreeNode* node) const; 62 63 AVLTreeNode* Find(const void* key); 64 AVLTreeNode* FindClose(const void* key, bool less); 65 66 status_t Insert(AVLTreeNode* element); 67 AVLTreeNode* Remove(const void* key); 68 bool Remove(AVLTreeNode* element); 69 70 private: 71 enum { 72 NOT_FOUND = -3, 73 DUPLICATE = -2, 74 NO_MEMORY = -1, 75 OK = 0, 76 HEIGHT_CHANGED = 1, 77 78 LEFT = -1, 79 BALANCED = 0, 80 RIGHT = 1, 81 }; 82 83 // rotations 84 void _RotateRight(AVLTreeNode** nodeP); 85 void _RotateLeft(AVLTreeNode** nodeP); 86 87 // insert 88 int _BalanceInsertLeft(AVLTreeNode** node); 89 int _BalanceInsertRight(AVLTreeNode** node); 90 int _Insert(AVLTreeNode* nodeToInsert); 91 92 // remove 93 int _BalanceRemoveLeft(AVLTreeNode** node); 94 int _BalanceRemoveRight(AVLTreeNode** node); 95 int _RemoveRightMostChild(AVLTreeNode** node, 96 AVLTreeNode** foundNode); 97 int _Remove(AVLTreeNode* node); 98 99 AVLTreeNode* fRoot; 100 int fNodeCount; 101 AVLTreeCompare* fCompare; 102 }; 103 104 105 // AVLTreeIterator 106 class AVLTreeIterator { 107 public: 108 inline AVLTreeIterator() 109 : fParent(NULL), 110 fCurrent(NULL), 111 fNext(NULL) 112 { 113 } 114 115 inline AVLTreeIterator(const AVLTreeIterator& other) 116 : fParent(other.fParent), 117 fCurrent(other.fCurrent), 118 fNext(other.fNext) 119 { 120 } 121 122 inline AVLTreeNode* Current() const 123 { 124 return fCurrent; 125 } 126 127 inline bool HasNext() const 128 { 129 return fNext; 130 } 131 132 inline AVLTreeNode* Next() 133 { 134 fCurrent = fNext; 135 136 if (fNext) 137 fNext = fParent->Next(fNext); 138 139 return fCurrent; 140 } 141 142 inline AVLTreeNode* Previous() 143 { 144 if (fCurrent) { 145 fNext = fCurrent; 146 fCurrent = fParent->Previous(fCurrent); 147 } else if (fNext) 148 fCurrent = fParent->Previous(fNext); 149 150 return fCurrent; 151 } 152 153 inline AVLTreeNode* Remove() 154 { 155 if (!fCurrent) 156 return NULL; 157 158 AVLTreeNode* node = fCurrent; 159 fCurrent = NULL; 160 161 return (const_cast<AVLTree*>(fParent)->Remove(node) ? node : NULL); 162 } 163 164 inline AVLTreeIterator& operator=(const AVLTreeIterator& other) 165 { 166 fParent = other.fParent; 167 fCurrent = other.fCurrent; 168 fNext = other.fNext; 169 return *this; 170 } 171 172 private: 173 inline AVLTreeIterator(const AVLTree* parent, AVLTreeNode* current, 174 AVLTreeNode* next) 175 : fParent(parent), 176 fCurrent(current), 177 fNext(next) 178 { 179 } 180 181 protected: 182 friend class AVLTree; 183 184 const AVLTree* fParent; 185 AVLTreeNode* fCurrent; 186 AVLTreeNode* fNext; 187 }; 188 189 190 // GetIterator 191 inline AVLTreeIterator 192 AVLTree::GetIterator() const 193 { 194 return AVLTreeIterator(this, NULL, LeftMost(fRoot)); 195 } 196 197 198 // GetIterator 199 inline AVLTreeIterator 200 AVLTree::GetIterator(AVLTreeNode* node) const 201 { 202 return AVLTreeIterator(this, node, Next(node)); 203 } 204 205 206 // #pragma mark - AVLTreeMap and friends 207 208 209 // strategies 210 namespace AVLTreeMapStrategy { 211 // key orders 212 template<typename Value> class Ascending; 213 template<typename Value> class Descending; 214 215 //! Automatic node strategy (works like STL containers do) 216 template <typename Key, typename Value, 217 typename KeyOrder = Ascending<Key>, 218 template <typename> class Allocator = MallocFreeAllocator> 219 class Auto; 220 221 /*! NodeStrategy must implement this interface: 222 typename Node; 223 inline Node* Allocate(const Key& key, const Value& value) 224 inline void Free(Node* node) 225 inline Key GetKey(Node* node) const 226 inline Value& GetValue(Node* node) const 227 inline AVLTreeNode* GetAVLTreeNode(Node* node) const 228 inline Node* GetNode(AVLTreeNode* node) const 229 inline int CompareKeyNode(const Key& a, const Node* b) 230 inline int CompareNodes(const Node* a, const Node* b) 231 */ 232 } 233 234 235 // for convenience 236 #define _AVL_TREE_MAP_TEMPLATE_LIST template<typename Key, typename Value, \ 237 typename NodeStrategy> 238 #define _AVL_TREE_MAP_CLASS_NAME AVLTreeMap<Key, Value, NodeStrategy> 239 240 241 // AVLTreeMap 242 template<typename Key, typename Value, 243 typename NodeStrategy = AVLTreeMapStrategy::Auto<Key, Value> > 244 class AVLTreeMap : protected AVLTreeCompare { 245 private: 246 typedef typename NodeStrategy::Node Node; 247 typedef _AVL_TREE_MAP_CLASS_NAME Class; 248 249 public: 250 class Iterator; 251 class ConstIterator; 252 253 public: 254 AVLTreeMap(const NodeStrategy& strategy 255 = NodeStrategy()); 256 virtual ~AVLTreeMap(); 257 258 inline int Count() const { return fTree.Count(); } 259 inline bool IsEmpty() const { return fTree.IsEmpty(); } 260 inline void MakeEmpty(); 261 262 Node* RootNode() const; 263 264 inline Iterator GetIterator(); 265 inline ConstIterator GetIterator() const; 266 267 inline Iterator GetIterator(Node* node); 268 inline ConstIterator GetIterator(Node* node) const; 269 270 Iterator Find(const Key& key); 271 Iterator FindClose(const Key& key, bool less); 272 273 status_t Insert(const Key& key, const Value& value, 274 Iterator* iterator); 275 status_t Remove(const Key& key); 276 277 const NodeStrategy& GetNodeStrategy() const { return fStrategy; } 278 279 protected: 280 // AVLTreeCompare 281 virtual int CompareKeyNode(const void* key, 282 const AVLTreeNode* node); 283 virtual int CompareNodes(const AVLTreeNode* node1, 284 const AVLTreeNode* node2); 285 286 void _FreeTree(AVLTreeNode* node); 287 288 // strategy shortcuts 289 inline Node* _Allocate(const Key& key, const Value& value); 290 inline void _Free(Node* node); 291 inline Key& _GetKey(Node* node) const; 292 inline Value& _GetValue(Node* node) const; 293 inline AVLTreeNode* _GetAVLTreeNode(const Node* node) const; 294 inline Node* _GetNode(const AVLTreeNode* node) const; 295 inline int _CompareKeyNode(const Key& a, const Node* b); 296 inline int _CompareNodes(const Node* a, const Node* b); 297 298 protected: 299 friend class Iterator; 300 friend class ConstIterator; 301 302 AVLTree fTree; 303 NodeStrategy fStrategy; 304 305 public: 306 // Iterator 307 // (need to implement it here, otherwise gcc 2.95.3 chokes) 308 class Iterator : public ConstIterator { 309 public: 310 inline Iterator() 311 : ConstIterator() 312 { 313 } 314 315 inline Iterator(const Iterator& other) 316 : ConstIterator(other) 317 { 318 } 319 320 inline void Remove() 321 { 322 if (AVLTreeNode* node = ConstIterator::fTreeIterator.Remove()) { 323 _AVL_TREE_MAP_CLASS_NAME* parent 324 = const_cast<_AVL_TREE_MAP_CLASS_NAME*>( 325 ConstIterator::fParent); 326 parent->_Free(parent->_GetNode(node)); 327 } 328 } 329 330 private: 331 inline Iterator(_AVL_TREE_MAP_CLASS_NAME* parent, 332 const AVLTreeIterator& treeIterator) 333 : ConstIterator(parent, treeIterator) 334 { 335 } 336 337 friend class _AVL_TREE_MAP_CLASS_NAME; 338 }; 339 }; 340 341 342 // ConstIterator 343 _AVL_TREE_MAP_TEMPLATE_LIST 344 class _AVL_TREE_MAP_CLASS_NAME::ConstIterator { 345 public: 346 inline ConstIterator() 347 : fParent(NULL), 348 fTreeIterator() 349 { 350 } 351 352 inline ConstIterator(const ConstIterator& other) 353 : fParent(other.fParent), 354 fTreeIterator(other.fTreeIterator) 355 { 356 } 357 358 inline bool HasCurrent() const 359 { 360 return fTreeIterator.Current(); 361 } 362 363 inline Key CurrentKey() 364 { 365 if (AVLTreeNode* node = fTreeIterator.Current()) 366 return fParent->_GetKey(fParent->_GetNode(node)); 367 return Key(); 368 } 369 370 inline Value Current() 371 { 372 if (AVLTreeNode* node = fTreeIterator.Current()) 373 return fParent->_GetValue(fParent->_GetNode(node)); 374 return Value(); 375 } 376 377 inline Value* CurrentValuePointer() 378 { 379 if (AVLTreeNode* node = fTreeIterator.Current()) 380 return &fParent->_GetValue(fParent->_GetNode(node)); 381 return NULL; 382 } 383 384 inline bool HasNext() const 385 { 386 return fTreeIterator.HasNext(); 387 } 388 389 inline Value Next() 390 { 391 if (AVLTreeNode* node = fTreeIterator.Next()) 392 return fParent->_GetValue(fParent->_GetNode(node)); 393 return Value(); 394 } 395 396 inline Value Previous() 397 { 398 if (AVLTreeNode* node = fTreeIterator.Previous()) 399 return fParent->_GetValue(fParent->_GetNode(node)); 400 return Value(); 401 } 402 403 inline ConstIterator& operator=(const ConstIterator& other) 404 { 405 fParent = other.fParent; 406 fTreeIterator = other.fTreeIterator; 407 return *this; 408 } 409 410 protected: 411 inline ConstIterator(const _AVL_TREE_MAP_CLASS_NAME* parent, 412 const AVLTreeIterator& treeIterator) 413 { 414 fParent = parent; 415 fTreeIterator = treeIterator; 416 } 417 418 friend class _AVL_TREE_MAP_CLASS_NAME; 419 420 const _AVL_TREE_MAP_CLASS_NAME* fParent; 421 AVLTreeIterator fTreeIterator; 422 }; 423 424 425 // constructor 426 _AVL_TREE_MAP_TEMPLATE_LIST 427 _AVL_TREE_MAP_CLASS_NAME::AVLTreeMap(const NodeStrategy& strategy) 428 : fTree(this), 429 fStrategy(strategy) 430 { 431 } 432 433 434 // destructor 435 _AVL_TREE_MAP_TEMPLATE_LIST 436 _AVL_TREE_MAP_CLASS_NAME::~AVLTreeMap() 437 { 438 } 439 440 441 // MakeEmpty 442 _AVL_TREE_MAP_TEMPLATE_LIST 443 inline void 444 _AVL_TREE_MAP_CLASS_NAME::MakeEmpty() 445 { 446 AVLTreeNode* root = fTree.Root(); 447 _FreeTree(root); 448 } 449 450 451 // RootNode 452 _AVL_TREE_MAP_TEMPLATE_LIST 453 inline typename _AVL_TREE_MAP_CLASS_NAME::Node* 454 _AVL_TREE_MAP_CLASS_NAME::RootNode() const 455 { 456 if (AVLTreeNode* root = fTree.Root()) 457 return _GetNode(root); 458 return NULL; 459 } 460 461 462 // GetIterator 463 _AVL_TREE_MAP_TEMPLATE_LIST 464 inline typename _AVL_TREE_MAP_CLASS_NAME::Iterator 465 _AVL_TREE_MAP_CLASS_NAME::GetIterator() 466 { 467 return Iterator(this, fTree.GetIterator()); 468 } 469 470 471 // GetIterator 472 _AVL_TREE_MAP_TEMPLATE_LIST 473 inline typename _AVL_TREE_MAP_CLASS_NAME::ConstIterator 474 _AVL_TREE_MAP_CLASS_NAME::GetIterator() const 475 { 476 return ConstIterator(this, fTree.GetIterator()); 477 } 478 479 480 // GetIterator 481 _AVL_TREE_MAP_TEMPLATE_LIST 482 inline typename _AVL_TREE_MAP_CLASS_NAME::Iterator 483 _AVL_TREE_MAP_CLASS_NAME::GetIterator(Node* node) 484 { 485 return Iterator(this, fTree.GetIterator(_GetAVLTreeNode(node))); 486 } 487 488 489 // GetIterator 490 _AVL_TREE_MAP_TEMPLATE_LIST 491 inline typename _AVL_TREE_MAP_CLASS_NAME::ConstIterator 492 _AVL_TREE_MAP_CLASS_NAME::GetIterator(Node* node) const 493 { 494 return ConstIterator(this, fTree.GetIterator(_GetAVLTreeNode(node))); 495 } 496 497 498 // Find 499 _AVL_TREE_MAP_TEMPLATE_LIST 500 typename _AVL_TREE_MAP_CLASS_NAME::Iterator 501 _AVL_TREE_MAP_CLASS_NAME::Find(const Key& key) 502 { 503 if (AVLTreeNode* node = fTree.Find(&key)) 504 return Iterator(this, fTree.GetIterator(node)); 505 return Iterator(); 506 } 507 508 509 // FindClose 510 _AVL_TREE_MAP_TEMPLATE_LIST 511 typename _AVL_TREE_MAP_CLASS_NAME::Iterator 512 _AVL_TREE_MAP_CLASS_NAME::FindClose(const Key& key, bool less) 513 { 514 if (AVLTreeNode* node = fTree.FindClose(&key, less)) 515 return Iterator(this, fTree.GetIterator(node)); 516 return Iterator(); 517 } 518 519 520 // Insert 521 _AVL_TREE_MAP_TEMPLATE_LIST 522 status_t 523 _AVL_TREE_MAP_CLASS_NAME::Insert(const Key& key, const Value& value, 524 Iterator* iterator) 525 { 526 // allocate a node 527 Node* userNode = _Allocate(key, value); 528 if (!userNode) 529 return B_NO_MEMORY; 530 531 // insert node 532 AVLTreeNode* node = _GetAVLTreeNode(userNode); 533 status_t error = fTree.Insert(node); 534 if (error != B_OK) { 535 _Free(userNode); 536 return error; 537 } 538 539 if (iterator) 540 *iterator = Iterator(this, fTree.GetIterator(node)); 541 542 return B_OK; 543 } 544 545 546 // Remove 547 _AVL_TREE_MAP_TEMPLATE_LIST 548 status_t 549 _AVL_TREE_MAP_CLASS_NAME::Remove(const Key& key) 550 { 551 AVLTreeNode* node = fTree.Remove(&key); 552 if (!node) 553 return B_ENTRY_NOT_FOUND; 554 555 _Free(_GetNode(node)); 556 return B_OK; 557 } 558 559 560 // CompareKeyNode 561 _AVL_TREE_MAP_TEMPLATE_LIST 562 int 563 _AVL_TREE_MAP_CLASS_NAME::CompareKeyNode(const void* key, 564 const AVLTreeNode* node) 565 { 566 return _CompareKeyNode(*(const Key*)key, _GetNode(node)); 567 } 568 569 570 // CompareNodes 571 _AVL_TREE_MAP_TEMPLATE_LIST 572 int 573 _AVL_TREE_MAP_CLASS_NAME::CompareNodes(const AVLTreeNode* node1, 574 const AVLTreeNode* node2) 575 { 576 return _CompareNodes(_GetNode(node1), _GetNode(node2)); 577 } 578 579 580 // _Allocate 581 _AVL_TREE_MAP_TEMPLATE_LIST 582 inline typename _AVL_TREE_MAP_CLASS_NAME::Node* 583 _AVL_TREE_MAP_CLASS_NAME::_Allocate(const Key& key, const Value& value) 584 { 585 return fStrategy.Allocate(key, value); 586 } 587 588 589 // _Free 590 _AVL_TREE_MAP_TEMPLATE_LIST 591 inline void 592 _AVL_TREE_MAP_CLASS_NAME::_Free(Node* node) 593 { 594 fStrategy.Free(node); 595 } 596 597 598 // _GetKey 599 _AVL_TREE_MAP_TEMPLATE_LIST 600 inline Key& 601 _AVL_TREE_MAP_CLASS_NAME::_GetKey(Node* node) const 602 { 603 return fStrategy.GetKey(node); 604 } 605 606 607 // _GetValue 608 _AVL_TREE_MAP_TEMPLATE_LIST 609 inline Value& 610 _AVL_TREE_MAP_CLASS_NAME::_GetValue(Node* node) const 611 { 612 return fStrategy.GetValue(node); 613 } 614 615 616 // _GetAVLTreeNode 617 _AVL_TREE_MAP_TEMPLATE_LIST 618 inline AVLTreeNode* 619 _AVL_TREE_MAP_CLASS_NAME::_GetAVLTreeNode(const Node* node) const 620 { 621 return fStrategy.GetAVLTreeNode(const_cast<Node*>(node)); 622 } 623 624 625 // _GetNode 626 _AVL_TREE_MAP_TEMPLATE_LIST 627 inline typename _AVL_TREE_MAP_CLASS_NAME::Node* 628 _AVL_TREE_MAP_CLASS_NAME::_GetNode(const AVLTreeNode* node) const 629 { 630 return fStrategy.GetNode(const_cast<AVLTreeNode*>(node)); 631 } 632 633 634 // _CompareKeyNode 635 _AVL_TREE_MAP_TEMPLATE_LIST 636 inline int 637 _AVL_TREE_MAP_CLASS_NAME::_CompareKeyNode(const Key& a, const Node* b) 638 { 639 return fStrategy.CompareKeyNode(a, b); 640 } 641 642 643 // _CompareNodes 644 _AVL_TREE_MAP_TEMPLATE_LIST 645 inline int 646 _AVL_TREE_MAP_CLASS_NAME::_CompareNodes(const Node* a, const Node* b) 647 { 648 return fStrategy.CompareNodes(a, b); 649 } 650 651 652 // _FreeTree 653 _AVL_TREE_MAP_TEMPLATE_LIST 654 void 655 _AVL_TREE_MAP_CLASS_NAME::_FreeTree(AVLTreeNode* node) 656 { 657 if (node) { 658 _FreeTree(node->left); 659 _FreeTree(node->right); 660 _Free(_GetNode(node)); 661 } 662 } 663 664 665 // #pragma mark - strategy parameters 666 667 // Ascending 668 namespace AVLTreeMapStrategy { 669 template<typename Value> 670 class Ascending { 671 public: 672 inline int operator()(const Value &a, const Value &b) const 673 { 674 if (a < b) 675 return -1; 676 else if (a > b) 677 return 1; 678 return 0; 679 } 680 }; 681 } 682 683 684 // Descending 685 namespace AVLTreeMapStrategy { 686 template<typename Value> 687 class Descending { 688 public: 689 inline int operator()(const Value &a, const Value &b) const 690 { 691 if (a < b) 692 return -1; 693 else if (a > b) 694 return 1; 695 return 0; 696 } 697 }; 698 } 699 700 701 // #pragma mark - strategies 702 703 704 // Auto 705 namespace AVLTreeMapStrategy { 706 template <typename Key, typename Value, typename KeyOrder, 707 template <typename> class Allocator> 708 class Auto { 709 public: 710 struct Node : AVLTreeNode { 711 Node(const Key &key, const Value &value) 712 : AVLTreeNode(), 713 key(key), 714 value(value) 715 { 716 } 717 718 Key key; 719 Value value; 720 }; 721 722 inline Node* Allocate(const Key& key, const Value& value) 723 { 724 Node* result = fAllocator.Allocate(); 725 if (result) 726 fAllocator.Construct(result, key, value); 727 return result; 728 } 729 730 inline void Free(Node* node) 731 { 732 fAllocator.Destruct(node); 733 fAllocator.Deallocate(node); 734 } 735 736 inline Key& GetKey(Node* node) const 737 { 738 return node->key; 739 } 740 741 inline Value& GetValue(Node* node) const 742 { 743 return node->value; 744 } 745 746 inline AVLTreeNode* GetAVLTreeNode(Node* node) const 747 { 748 return node; 749 } 750 751 inline Node* GetNode(AVLTreeNode* node) const 752 { 753 return static_cast<Node*>(node); 754 } 755 756 inline int CompareKeyNode(const Key& a, const Node* b) const 757 { 758 return fCompare(a, _GetKey(b)); 759 } 760 761 inline int CompareNodes(const Node* a, const Node* b) const 762 { 763 return fCompare(_GetKey(a), _GetKey(b)); 764 } 765 766 private: 767 KeyOrder fCompare; 768 Allocator<Node> fAllocator; 769 }; 770 } 771 772 #endif // _AVL_TREE_MAP_H 773