1 /* 2 * Copyright 2003-2011, Ingo Weinhold, ingo_weinhold@gmx.de. 3 * Distributed under the terms of the MIT License. 4 */ 5 #ifndef TWO_KEY_AVL_TREE_H 6 #define TWO_KEY_AVL_TREE_H 7 8 9 #include <slab/Slab.h> 10 #include <util/AVLTreeMap.h> 11 12 13 // #pragma mark - TwoKeyAVLTreeKey 14 15 16 template<typename PrimaryKey, typename SecondaryKey> 17 class TwoKeyAVLTreeKey { 18 public: 19 inline TwoKeyAVLTreeKey(const PrimaryKey& primary, 20 const SecondaryKey& secondary) 21 : 22 primary(primary), 23 secondary(secondary), 24 use_secondary(true) 25 { 26 } 27 28 inline TwoKeyAVLTreeKey(const PrimaryKey* primary) 29 : 30 primary(primary), 31 secondary(NULL), 32 use_secondary(false) 33 { 34 } 35 36 PrimaryKey primary; 37 SecondaryKey secondary; 38 bool use_secondary; 39 }; 40 41 42 // #pragma mark - TwoKeyAVLTreeKeyCompare 43 44 45 template<typename PrimaryKey, typename SecondaryKey, 46 typename PrimaryKeyCompare, typename SecondaryKeyCompare> 47 class TwoKeyAVLTreeKeyCompare { 48 private: 49 typedef TwoKeyAVLTreeKey<PrimaryKey, SecondaryKey> Key; 50 51 public: 52 inline TwoKeyAVLTreeKeyCompare(const PrimaryKeyCompare& primary, 53 const SecondaryKeyCompare& secondary) 54 : 55 fPrimaryKeyCompare(primary), 56 fSecondaryKeyCompare(secondary) 57 { 58 } 59 60 inline int operator()(const Key& a, const Key& b) const 61 { 62 int result = fPrimaryKeyCompare(a.primary, b.primary); 63 if (result == 0 && a.use_secondary && b.use_secondary) 64 result = fSecondaryKeyCompare(a.secondary, b.secondary); 65 return result; 66 } 67 68 private: 69 PrimaryKeyCompare fPrimaryKeyCompare; 70 SecondaryKeyCompare fSecondaryKeyCompare; 71 }; 72 73 74 // #pragma mark - TwoKeyAVLTreeGetKey 75 76 77 template<typename Value, typename PrimaryKey, typename SecondaryKey, 78 typename GetPrimaryKey, typename GetSecondaryKey> 79 class TwoKeyAVLTreeGetKey { 80 private: 81 typedef TwoKeyAVLTreeKey<PrimaryKey, SecondaryKey> Key; 82 83 public: 84 TwoKeyAVLTreeGetKey(const GetPrimaryKey& getPrimary, 85 const GetSecondaryKey& getSecondary) 86 : 87 fGetPrimaryKey(getPrimary), 88 fGetSecondaryKey(getSecondary) 89 { 90 } 91 92 inline Key operator()(const Value& a) const 93 { 94 return Key(fGetPrimaryKey(a), fGetSecondaryKey(a)); 95 } 96 97 private: 98 GetPrimaryKey fGetPrimaryKey; 99 GetSecondaryKey fGetSecondaryKey; 100 }; 101 102 103 // #pragma mark - TwoKeyAVLTreeStandardCompare 104 105 106 template<typename Value> 107 class TwoKeyAVLTreeStandardCompare { 108 public: 109 inline int operator()(const Value& a, const Value& b) const 110 { 111 if (a < b) 112 return -1; 113 else if (a > b) 114 return 1; 115 return 0; 116 } 117 }; 118 119 120 // #pragma mark - TwoKeyAVLTreeStandardGetKey 121 122 123 template<typename Value, typename Key> 124 class TwoKeyAVLTreeStandardGetKey { 125 public: 126 inline const Key& operator()(const Value& a) const 127 { 128 return a; 129 } 130 131 inline Key& operator()(Value& a) const 132 { 133 return a; 134 } 135 }; 136 137 138 // #pragma mark - TwoKeyAVLTreeNodeStrategy 139 140 141 template<typename Value> 142 struct TwoKeyAVLTreeNode : AVLTreeNode { 143 static object_cache* sNodeCache; 144 145 static void* operator new(size_t size) 146 { 147 object_cache* cache = TwoKeyAVLTreeNode<void*>::sNodeCache; 148 const size_t nodeSize = sizeof(TwoKeyAVLTreeNode<void*>); 149 if (size != nodeSize || cache == NULL) 150 panic("unexpected size passed to operator new!"); 151 return object_cache_alloc(cache, 0); 152 } 153 static void operator delete(void* object) 154 { 155 object_cache_free(TwoKeyAVLTreeNode<void*>::sNodeCache, object, 0); 156 } 157 158 public: 159 TwoKeyAVLTreeNode(const Value& value) 160 : 161 AVLTreeNode(), 162 value(value) 163 { 164 } 165 166 Value value; 167 }; 168 169 170 template <typename PrimaryKey, typename SecondaryKey, typename Value, 171 typename PrimaryKeyCompare, typename SecondaryKeyCompare, 172 typename GetPrimaryKey, typename GetSecondaryKey> 173 class TwoKeyAVLTreeNodeStrategy { 174 public: 175 typedef TwoKeyAVLTreeKey<PrimaryKey, SecondaryKey> Key; 176 typedef TwoKeyAVLTreeNode<Value> Node; 177 178 TwoKeyAVLTreeNodeStrategy( 179 const PrimaryKeyCompare& primaryKeyCompare = PrimaryKeyCompare(), 180 const SecondaryKeyCompare& secondaryKeyCompare = SecondaryKeyCompare(), 181 const GetPrimaryKey& getPrimaryKey = GetPrimaryKey(), 182 const GetSecondaryKey& getSecondaryKey = GetSecondaryKey()) 183 : 184 fPrimaryKeyCompare(primaryKeyCompare), 185 fSecondaryKeyCompare(secondaryKeyCompare), 186 fGetPrimaryKey(getPrimaryKey), 187 fGetSecondaryKey(getSecondaryKey) 188 { 189 } 190 ~TwoKeyAVLTreeNodeStrategy() 191 { 192 } 193 194 inline Node* Allocate(const Key& key, const Value& value) const 195 { 196 return new Node(value); 197 } 198 199 inline void Free(Node* node) const 200 { 201 delete node; 202 } 203 204 // internal use (not part of the strategy) 205 inline Key GetValueKey(const Value& value) const 206 { 207 return Key(fGetPrimaryKey(value), fGetSecondaryKey(value)); 208 } 209 210 inline Key GetKey(Node* node) const 211 { 212 return GetValueKey(node->value); 213 } 214 215 inline Value& GetValue(Node* node) const 216 { 217 return node->value; 218 } 219 220 inline AVLTreeNode* GetAVLTreeNode(Node* node) const 221 { 222 return node; 223 } 224 225 inline Node* GetNode(AVLTreeNode* node) const 226 { 227 return static_cast<Node*>(node); 228 } 229 230 inline int CompareKeyNode(const Key& a, const Node* b) const 231 { 232 return _CompareKeys(a, GetKey(const_cast<Node*>(b))); 233 } 234 235 inline int CompareNodes(const Node* a, const Node* b) const 236 { 237 return _CompareKeys(GetKey(const_cast<Node*>(a)), 238 GetKey(const_cast<Node*>(b))); 239 } 240 241 private: 242 inline int _CompareKeys(const Key& a, const Key& b) const 243 { 244 int result = fPrimaryKeyCompare(a.primary, b.primary); 245 if (result == 0 && a.use_secondary && b.use_secondary) 246 result = fSecondaryKeyCompare(a.secondary, b.secondary); 247 return result; 248 } 249 250 PrimaryKeyCompare fPrimaryKeyCompare; 251 SecondaryKeyCompare fSecondaryKeyCompare; 252 GetPrimaryKey fGetPrimaryKey; 253 GetSecondaryKey fGetSecondaryKey; 254 }; 255 256 257 // for convenience 258 #define TWO_KEY_AVL_TREE_TEMPLATE_LIST template<typename Value, \ 259 typename PrimaryKey, typename PrimaryKeyCompare, \ 260 typename GetPrimaryKey, typename SecondaryKey, \ 261 typename SecondaryKeyCompare, typename GetSecondaryKey> 262 #define TWO_KEY_AVL_TREE_CLASS_NAME TwoKeyAVLTree<Value, PrimaryKey, \ 263 PrimaryKeyCompare, GetPrimaryKey, SecondaryKey, \ 264 SecondaryKeyCompare, GetSecondaryKey> 265 266 267 // #pragma mark - TwoKeyAVLTree 268 269 270 template<typename Value, typename PrimaryKey, 271 typename PrimaryKeyCompare, typename GetPrimaryKey, 272 typename SecondaryKey = Value, 273 typename SecondaryKeyCompare = TwoKeyAVLTreeStandardCompare<SecondaryKey>, 274 typename GetSecondaryKey 275 = TwoKeyAVLTreeStandardGetKey<Value, SecondaryKey> > 276 class TwoKeyAVLTree { 277 public: 278 typedef TwoKeyAVLTreeKey<PrimaryKey, SecondaryKey> Key; 279 typedef TwoKeyAVLTreeNodeStrategy<PrimaryKey, SecondaryKey, Value, 280 PrimaryKeyCompare, SecondaryKeyCompare, GetPrimaryKey, 281 GetSecondaryKey> NodeStrategy; 282 typedef AVLTreeMap<Key, Value, NodeStrategy> TreeMap; 283 284 typedef typename TreeMap::Iterator TreeMapIterator; 285 typedef typename NodeStrategy::Node Node; 286 287 class Iterator; 288 289 public: 290 TwoKeyAVLTree(); 291 TwoKeyAVLTree( 292 const PrimaryKeyCompare& primaryCompare, 293 const GetPrimaryKey& getPrimary, 294 const SecondaryKeyCompare& secondaryCompare, 295 const GetSecondaryKey& getSecondary); 296 ~TwoKeyAVLTree(); 297 298 inline int CountItems() const { return fTreeMap.Count(); } 299 300 Node* Previous(Node* node) const; 301 Node* Next(Node* node) const; 302 303 Value* FindFirst(const PrimaryKey& key, 304 Iterator* iterator = NULL); 305 Value* FindFirstClosest(const PrimaryKey& key, 306 bool less, Iterator* iterator = NULL); 307 Value* FindLast(const PrimaryKey& key, 308 Iterator* iterator = NULL); 309 inline Value* Find(const PrimaryKey& primaryKey, 310 const SecondaryKey& secondaryKey, 311 Iterator* iterator = NULL); 312 313 inline void GetIterator(Iterator* iterator); 314 inline void GetIterator(Node* node, Iterator* iterator); 315 316 inline status_t Insert(const Value& value, 317 Iterator* iterator); 318 inline status_t Insert(const Value& value, 319 Node** _node = NULL); 320 inline status_t Remove(const PrimaryKey& primaryKey, 321 const SecondaryKey& secondaryKey); 322 inline status_t Remove(Node* node); 323 324 private: 325 Node* _FindFirst(const PrimaryKey& key, 326 Node** _parent) const; 327 328 private: 329 TreeMap fTreeMap; 330 PrimaryKeyCompare fPrimaryKeyCompare; 331 GetPrimaryKey fGetPrimaryKey; 332 }; 333 334 335 // #pragma mark - Iterator 336 337 338 TWO_KEY_AVL_TREE_TEMPLATE_LIST 339 class TWO_KEY_AVL_TREE_CLASS_NAME::Iterator { 340 public: 341 typedef typename TWO_KEY_AVL_TREE_CLASS_NAME::TreeMapIterator 342 TreeMapIterator; 343 344 inline Iterator() 345 : 346 fIterator() 347 { 348 } 349 350 inline ~Iterator() 351 { 352 } 353 354 inline Value* Current() 355 { 356 return fIterator.CurrentValuePointer(); 357 } 358 359 inline Node* CurrentNode() 360 { 361 return fIterator.CurrentNode(); 362 } 363 364 inline Value* Next() 365 { 366 fIterator.Next(); 367 return Current(); 368 } 369 370 inline Value* Previous() 371 { 372 fIterator.Previous(); 373 return Current(); 374 } 375 376 inline void Remove() 377 { 378 fIterator.Remove(); 379 } 380 381 private: 382 friend class TWO_KEY_AVL_TREE_CLASS_NAME; 383 384 Iterator(const Iterator& other) 385 : 386 fIterator(other.fIterator) 387 { 388 } 389 390 Iterator& operator=(const Iterator& other) 391 { 392 fIterator = other.fIterator; 393 } 394 395 Iterator(const TreeMapIterator& iterator) 396 : 397 fIterator() 398 { 399 } 400 401 inline void _SetTo(const TreeMapIterator& iterator) 402 { 403 fIterator = iterator; 404 } 405 406 private: 407 TWO_KEY_AVL_TREE_CLASS_NAME::TreeMapIterator fIterator; 408 }; 409 410 411 TWO_KEY_AVL_TREE_TEMPLATE_LIST 412 TWO_KEY_AVL_TREE_CLASS_NAME::TwoKeyAVLTree() 413 : 414 fTreeMap(NodeStrategy(PrimaryKeyCompare(), SecondaryKeyCompare(), 415 GetPrimaryKey(), GetSecondaryKey())), 416 fPrimaryKeyCompare(PrimaryKeyCompare()), 417 fGetPrimaryKey(GetPrimaryKey()) 418 { 419 } 420 421 422 TWO_KEY_AVL_TREE_TEMPLATE_LIST 423 TWO_KEY_AVL_TREE_CLASS_NAME::TwoKeyAVLTree( 424 const PrimaryKeyCompare& primaryCompare, const GetPrimaryKey& getPrimary, 425 const SecondaryKeyCompare& secondaryCompare, 426 const GetSecondaryKey& getSecondary) 427 : 428 fTreeMap(NodeStrategy(primaryCompare, secondaryCompare, getPrimary, 429 getSecondary)), 430 fPrimaryKeyCompare(primaryCompare), 431 fGetPrimaryKey(getPrimary) 432 { 433 } 434 435 436 TWO_KEY_AVL_TREE_TEMPLATE_LIST 437 TWO_KEY_AVL_TREE_CLASS_NAME::~TwoKeyAVLTree() 438 { 439 } 440 441 442 TWO_KEY_AVL_TREE_TEMPLATE_LIST 443 typename TWO_KEY_AVL_TREE_CLASS_NAME::Node* 444 TWO_KEY_AVL_TREE_CLASS_NAME::Previous(Node* node) const 445 { 446 return fTreeMap.Previous(node); 447 } 448 449 450 TWO_KEY_AVL_TREE_TEMPLATE_LIST 451 typename TWO_KEY_AVL_TREE_CLASS_NAME::Node* 452 TWO_KEY_AVL_TREE_CLASS_NAME::Next(Node* node) const 453 { 454 return fTreeMap.Next(node); 455 } 456 457 458 TWO_KEY_AVL_TREE_TEMPLATE_LIST 459 Value* 460 TWO_KEY_AVL_TREE_CLASS_NAME::FindFirst(const PrimaryKey& key, 461 Iterator* iterator) 462 { 463 const NodeStrategy& strategy = fTreeMap.GetNodeStrategy(); 464 465 Node* node = _FindFirst(key, NULL); 466 if (node == NULL) 467 return NULL; 468 469 if (iterator != NULL) 470 iterator->_SetTo(fTreeMap.GetIterator(node)); 471 472 return &strategy.GetValue(node); 473 } 474 475 476 TWO_KEY_AVL_TREE_TEMPLATE_LIST 477 Value* 478 TWO_KEY_AVL_TREE_CLASS_NAME::FindFirstClosest(const PrimaryKey& key, bool less, 479 Iterator* iterator) 480 { 481 const NodeStrategy& strategy = fTreeMap.GetNodeStrategy(); 482 483 Node* parent = NULL; 484 Node* node = _FindFirst(key, &parent); 485 if (node == NULL) { 486 // not found -- try to get the closest node 487 if (parent == NULL) 488 return NULL; 489 490 node = parent; 491 int expectedCmp = less ? 1 : -1; 492 int cmp = fPrimaryKeyCompare(key, 493 fGetPrimaryKey(strategy.GetValue(strategy.GetNode(node)))); 494 495 if (cmp != expectedCmp) { 496 // The node's value is less although we were asked for a greater 497 // value, or the other way around. We need to iterate to the next 498 // node in the respective direction. If there is no node, we fail. 499 node = less ? Previous(node) : Next(node); 500 if (node == NULL) 501 return NULL; 502 } 503 } 504 505 if (iterator != NULL) 506 iterator->_SetTo(fTreeMap.GetIterator(node)); 507 508 return &strategy.GetValue(node); 509 } 510 511 512 TWO_KEY_AVL_TREE_TEMPLATE_LIST 513 Value* 514 TWO_KEY_AVL_TREE_CLASS_NAME::FindLast(const PrimaryKey& key, 515 Iterator* iterator) 516 { 517 const NodeStrategy& strategy = fTreeMap.GetNodeStrategy(); 518 Node* node = fTreeMap.RootNode(); 519 520 while (node) { 521 int cmp = fPrimaryKeyCompare(key, fGetPrimaryKey( 522 strategy.GetValue(node))); 523 if (cmp == 0) { 524 // found a matching node, now get the right-most node with that key 525 while (node->right && fPrimaryKeyCompare(key, 526 fGetPrimaryKey(strategy.GetValue( 527 strategy.GetNode(node->right)))) == 0) { 528 node = strategy.GetNode(node->right); 529 } 530 if (iterator) 531 iterator->_SetTo(fTreeMap.GetIterator(node)); 532 return &strategy.GetValue(node); 533 } 534 535 if (cmp < 0) 536 node = strategy.GetNode(node->left); 537 else 538 node = strategy.GetNode(node->right); 539 } 540 return NULL; 541 } 542 543 544 TWO_KEY_AVL_TREE_TEMPLATE_LIST 545 Value* 546 TWO_KEY_AVL_TREE_CLASS_NAME::Find(const PrimaryKey& primaryKey, 547 const SecondaryKey& secondaryKey, Iterator* iterator) 548 { 549 TreeMapIterator it = fTreeMap.Find(Key(primaryKey, secondaryKey)); 550 if (iterator) 551 iterator->_SetTo(it); 552 return it.CurrentValuePointer(); 553 } 554 555 556 TWO_KEY_AVL_TREE_TEMPLATE_LIST 557 void 558 TWO_KEY_AVL_TREE_CLASS_NAME::GetIterator(Iterator* iterator) 559 { 560 TreeMapIterator it = fTreeMap.GetIterator(); 561 it.Next(); 562 // Our iterator needs to point to the first entry already. 563 iterator->_SetTo(it); 564 } 565 566 567 TWO_KEY_AVL_TREE_TEMPLATE_LIST 568 void 569 TWO_KEY_AVL_TREE_CLASS_NAME::GetIterator(Node* node, Iterator* iterator) 570 { 571 iterator->_SetTo(fTreeMap.GetIterator(node)); 572 } 573 574 575 TWO_KEY_AVL_TREE_TEMPLATE_LIST 576 status_t 577 TWO_KEY_AVL_TREE_CLASS_NAME::Insert(const Value& value, Iterator* iterator) 578 { 579 NodeStrategy& strategy 580 = const_cast<NodeStrategy&>(fTreeMap.GetNodeStrategy()); 581 582 TreeMapIterator it; 583 status_t status = fTreeMap.Insert(strategy.GetValueKey(value), value, &it); 584 if (status != B_OK || !iterator) 585 return status; 586 587 iterator->_SetTo(it); 588 return B_OK; 589 } 590 591 592 TWO_KEY_AVL_TREE_TEMPLATE_LIST 593 status_t 594 TWO_KEY_AVL_TREE_CLASS_NAME::Insert(const Value& value, Node** _node) 595 { 596 NodeStrategy& strategy 597 = const_cast<NodeStrategy&>(fTreeMap.GetNodeStrategy()); 598 599 return fTreeMap.Insert(strategy.GetValueKey(value), value, _node); 600 } 601 602 603 TWO_KEY_AVL_TREE_TEMPLATE_LIST 604 status_t 605 TWO_KEY_AVL_TREE_CLASS_NAME::Remove(const PrimaryKey& primaryKey, 606 const SecondaryKey& secondaryKey) 607 { 608 return fTreeMap.Remove(Key(primaryKey, secondaryKey)); 609 } 610 611 612 TWO_KEY_AVL_TREE_TEMPLATE_LIST 613 status_t 614 TWO_KEY_AVL_TREE_CLASS_NAME::Remove(Node* node) 615 { 616 return fTreeMap.Remove(node); 617 } 618 619 620 TWO_KEY_AVL_TREE_TEMPLATE_LIST 621 typename TWO_KEY_AVL_TREE_CLASS_NAME::Node* 622 TWO_KEY_AVL_TREE_CLASS_NAME::_FindFirst(const PrimaryKey& key, 623 Node** _parent) const 624 { 625 const NodeStrategy& strategy = fTreeMap.GetNodeStrategy(); 626 Node* node = fTreeMap.RootNode(); 627 Node* parent = NULL; 628 629 while (node) { 630 int cmp = fPrimaryKeyCompare(key, fGetPrimaryKey( 631 strategy.GetValue(node))); 632 if (cmp == 0) { 633 // found a matching node, now get the left-most node with that key 634 while (node->left && fPrimaryKeyCompare(key, 635 fGetPrimaryKey(strategy.GetValue( 636 strategy.GetNode(node->left)))) == 0) { 637 node = strategy.GetNode(node->left); 638 } 639 640 return node; 641 } 642 643 parent = node; 644 645 if (cmp < 0) 646 node = strategy.GetNode(node->left); 647 else 648 node = strategy.GetNode(node->right); 649 } 650 651 if (_parent != NULL) 652 *_parent = parent; 653 654 return NULL; 655 } 656 657 658 #endif // TWO_KEY_AVL_TREE_H 659