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