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 PrimaryKey, typename SecondaryKey, typename Value, 142 typename PrimaryKeyCompare, typename SecondaryKeyCompare, 143 typename GetPrimaryKey, typename GetSecondaryKey> 144 class TwoKeyAVLTreeNodeStrategy { 145 public: 146 typedef TwoKeyAVLTreeKey<PrimaryKey, SecondaryKey> Key; 147 148 TwoKeyAVLTreeNodeStrategy( 149 const PrimaryKeyCompare& primaryKeyCompare = PrimaryKeyCompare(), 150 const SecondaryKeyCompare& secondaryKeyCompare = SecondaryKeyCompare(), 151 const GetPrimaryKey& getPrimaryKey = GetPrimaryKey(), 152 const GetSecondaryKey& getSecondaryKey = GetSecondaryKey()) 153 : 154 fPrimaryKeyCompare(primaryKeyCompare), 155 fSecondaryKeyCompare(secondaryKeyCompare), 156 fGetPrimaryKey(getPrimaryKey), 157 fGetSecondaryKey(getSecondaryKey) 158 { 159 fObjectCache = create_object_cache("packagefs TwoKeyAVLTreeNodes", sizeof(Node), 8, 160 NULL, NULL, NULL); 161 fObjectCacheRefs = new int32(1); 162 } 163 TwoKeyAVLTreeNodeStrategy(const TwoKeyAVLTreeNodeStrategy& other) 164 : 165 fPrimaryKeyCompare(other.fPrimaryKeyCompare), 166 fSecondaryKeyCompare(other.fSecondaryKeyCompare), 167 fGetPrimaryKey(other.fGetPrimaryKey), 168 fGetSecondaryKey(other.fGetSecondaryKey), 169 fObjectCache(other.fObjectCache), 170 fObjectCacheRefs(other.fObjectCacheRefs) 171 { 172 atomic_add(fObjectCacheRefs, 1); 173 } 174 ~TwoKeyAVLTreeNodeStrategy() 175 { 176 atomic_add(fObjectCacheRefs, -1); 177 if (atomic_get(fObjectCacheRefs) == 0) { 178 delete_object_cache(fObjectCache); 179 delete fObjectCacheRefs; 180 } 181 } 182 183 struct Node : AVLTreeNode { 184 static void* operator new(size_t size, object_cache* cache) { 185 if (size != sizeof(Node) || !cache) 186 panic("unexpected size passed to operator new!"); 187 return object_cache_alloc(cache, 0); 188 } 189 190 Node(const Value& value) 191 : 192 AVLTreeNode(), 193 value(value) 194 { 195 } 196 197 Value value; 198 }; 199 200 inline Node* Allocate(const Key& key, const Value& value) const 201 { 202 return new(fObjectCache) Node(value); 203 } 204 205 inline void Free(Node* node) const 206 { 207 if (node == NULL) 208 return; 209 210 // There is no way to overload operator delete with extra parameters. 211 node->~Node(); 212 object_cache_free(fObjectCache, node, 0); 213 } 214 215 // internal use (not part of the strategy) 216 inline Key GetValueKey(const Value& value) const 217 { 218 return Key(fGetPrimaryKey(value), fGetSecondaryKey(value)); 219 } 220 221 inline Key GetKey(Node* node) const 222 { 223 return GetValueKey(node->value); 224 } 225 226 inline Value& GetValue(Node* node) const 227 { 228 return node->value; 229 } 230 231 inline AVLTreeNode* GetAVLTreeNode(Node* node) const 232 { 233 return node; 234 } 235 236 inline Node* GetNode(AVLTreeNode* node) const 237 { 238 return static_cast<Node*>(node); 239 } 240 241 inline int CompareKeyNode(const Key& a, const Node* b) const 242 { 243 return _CompareKeys(a, GetKey(const_cast<Node*>(b))); 244 } 245 246 inline int CompareNodes(const Node* a, const Node* b) const 247 { 248 return _CompareKeys(GetKey(const_cast<Node*>(a)), 249 GetKey(const_cast<Node*>(b))); 250 } 251 252 private: 253 inline int _CompareKeys(const Key& a, const Key& b) const 254 { 255 int result = fPrimaryKeyCompare(a.primary, b.primary); 256 if (result == 0 && a.use_secondary && b.use_secondary) 257 result = fSecondaryKeyCompare(a.secondary, b.secondary); 258 return result; 259 } 260 261 PrimaryKeyCompare fPrimaryKeyCompare; 262 SecondaryKeyCompare fSecondaryKeyCompare; 263 GetPrimaryKey fGetPrimaryKey; 264 GetSecondaryKey fGetSecondaryKey; 265 266 object_cache* fObjectCache; 267 int32* fObjectCacheRefs; 268 }; 269 270 271 // for convenience 272 #define TWO_KEY_AVL_TREE_TEMPLATE_LIST template<typename Value, \ 273 typename PrimaryKey, typename PrimaryKeyCompare, \ 274 typename GetPrimaryKey, typename SecondaryKey, \ 275 typename SecondaryKeyCompare, typename GetSecondaryKey> 276 #define TWO_KEY_AVL_TREE_CLASS_NAME TwoKeyAVLTree<Value, PrimaryKey, \ 277 PrimaryKeyCompare, GetPrimaryKey, SecondaryKey, \ 278 SecondaryKeyCompare, GetSecondaryKey> 279 280 281 // #pragma mark - TwoKeyAVLTree 282 283 284 template<typename Value, typename PrimaryKey, 285 typename PrimaryKeyCompare, typename GetPrimaryKey, 286 typename SecondaryKey = Value, 287 typename SecondaryKeyCompare = TwoKeyAVLTreeStandardCompare<SecondaryKey>, 288 typename GetSecondaryKey 289 = TwoKeyAVLTreeStandardGetKey<Value, SecondaryKey> > 290 class TwoKeyAVLTree { 291 public: 292 typedef TwoKeyAVLTreeKey<PrimaryKey, SecondaryKey> Key; 293 typedef TwoKeyAVLTreeNodeStrategy<PrimaryKey, SecondaryKey, Value, 294 PrimaryKeyCompare, SecondaryKeyCompare, GetPrimaryKey, 295 GetSecondaryKey> NodeStrategy; 296 typedef AVLTreeMap<Key, Value, NodeStrategy> TreeMap; 297 298 typedef typename TreeMap::Iterator TreeMapIterator; 299 typedef typename NodeStrategy::Node Node; 300 301 class Iterator; 302 303 public: 304 TwoKeyAVLTree(); 305 TwoKeyAVLTree( 306 const PrimaryKeyCompare& primaryCompare, 307 const GetPrimaryKey& getPrimary, 308 const SecondaryKeyCompare& secondaryCompare, 309 const GetSecondaryKey& getSecondary); 310 ~TwoKeyAVLTree(); 311 312 inline int CountItems() const { return fTreeMap.Count(); } 313 314 Node* Previous(Node* node) const; 315 Node* Next(Node* node) const; 316 317 Value* FindFirst(const PrimaryKey& key, 318 Iterator* iterator = NULL); 319 Value* FindFirstClosest(const PrimaryKey& key, 320 bool less, Iterator* iterator = NULL); 321 Value* FindLast(const PrimaryKey& key, 322 Iterator* iterator = NULL); 323 inline Value* Find(const PrimaryKey& primaryKey, 324 const SecondaryKey& secondaryKey, 325 Iterator* iterator = NULL); 326 327 inline void GetIterator(Iterator* iterator); 328 inline void GetIterator(Node* node, Iterator* iterator); 329 330 inline status_t Insert(const Value& value, 331 Iterator* iterator); 332 inline status_t Insert(const Value& value, 333 Node** _node = NULL); 334 inline status_t Remove(const PrimaryKey& primaryKey, 335 const SecondaryKey& secondaryKey); 336 inline status_t Remove(Node* node); 337 338 private: 339 Node* _FindFirst(const PrimaryKey& key, 340 Node** _parent) const; 341 342 private: 343 TreeMap fTreeMap; 344 PrimaryKeyCompare fPrimaryKeyCompare; 345 GetPrimaryKey fGetPrimaryKey; 346 }; 347 348 349 // #pragma mark - Iterator 350 351 352 TWO_KEY_AVL_TREE_TEMPLATE_LIST 353 class TWO_KEY_AVL_TREE_CLASS_NAME::Iterator { 354 public: 355 typedef typename TWO_KEY_AVL_TREE_CLASS_NAME::TreeMapIterator 356 TreeMapIterator; 357 358 inline Iterator() 359 : 360 fIterator() 361 { 362 } 363 364 inline ~Iterator() 365 { 366 } 367 368 inline Value* Current() 369 { 370 return fIterator.CurrentValuePointer(); 371 } 372 373 inline Node* CurrentNode() 374 { 375 return fIterator.CurrentNode(); 376 } 377 378 inline Value* Next() 379 { 380 fIterator.Next(); 381 return Current(); 382 } 383 384 inline Value* Previous() 385 { 386 fIterator.Previous(); 387 return Current(); 388 } 389 390 inline void Remove() 391 { 392 fIterator.Remove(); 393 } 394 395 private: 396 friend class TWO_KEY_AVL_TREE_CLASS_NAME; 397 398 Iterator(const Iterator& other) 399 : 400 fIterator(other.fIterator) 401 { 402 } 403 404 Iterator& operator=(const Iterator& other) 405 { 406 fIterator = other.fIterator; 407 } 408 409 Iterator(const TreeMapIterator& iterator) 410 : 411 fIterator() 412 { 413 } 414 415 inline void _SetTo(const TreeMapIterator& iterator) 416 { 417 fIterator = iterator; 418 } 419 420 private: 421 TWO_KEY_AVL_TREE_CLASS_NAME::TreeMapIterator fIterator; 422 }; 423 424 425 TWO_KEY_AVL_TREE_TEMPLATE_LIST 426 TWO_KEY_AVL_TREE_CLASS_NAME::TwoKeyAVLTree() 427 : 428 fTreeMap(NodeStrategy(PrimaryKeyCompare(), SecondaryKeyCompare(), 429 GetPrimaryKey(), GetSecondaryKey())), 430 fPrimaryKeyCompare(PrimaryKeyCompare()), 431 fGetPrimaryKey(GetPrimaryKey()) 432 { 433 } 434 435 436 TWO_KEY_AVL_TREE_TEMPLATE_LIST 437 TWO_KEY_AVL_TREE_CLASS_NAME::TwoKeyAVLTree( 438 const PrimaryKeyCompare& primaryCompare, const GetPrimaryKey& getPrimary, 439 const SecondaryKeyCompare& secondaryCompare, 440 const GetSecondaryKey& getSecondary) 441 : 442 fTreeMap(NodeStrategy(primaryCompare, secondaryCompare, getPrimary, 443 getSecondary)), 444 fPrimaryKeyCompare(primaryCompare), 445 fGetPrimaryKey(getPrimary) 446 { 447 } 448 449 450 TWO_KEY_AVL_TREE_TEMPLATE_LIST 451 TWO_KEY_AVL_TREE_CLASS_NAME::~TwoKeyAVLTree() 452 { 453 } 454 455 456 TWO_KEY_AVL_TREE_TEMPLATE_LIST 457 typename TWO_KEY_AVL_TREE_CLASS_NAME::Node* 458 TWO_KEY_AVL_TREE_CLASS_NAME::Previous(Node* node) const 459 { 460 return fTreeMap.Previous(node); 461 } 462 463 464 TWO_KEY_AVL_TREE_TEMPLATE_LIST 465 typename TWO_KEY_AVL_TREE_CLASS_NAME::Node* 466 TWO_KEY_AVL_TREE_CLASS_NAME::Next(Node* node) const 467 { 468 return fTreeMap.Next(node); 469 } 470 471 472 TWO_KEY_AVL_TREE_TEMPLATE_LIST 473 Value* 474 TWO_KEY_AVL_TREE_CLASS_NAME::FindFirst(const PrimaryKey& key, 475 Iterator* iterator) 476 { 477 const NodeStrategy& strategy = fTreeMap.GetNodeStrategy(); 478 479 Node* node = _FindFirst(key, NULL); 480 if (node == NULL) 481 return NULL; 482 483 if (iterator != NULL) 484 iterator->_SetTo(fTreeMap.GetIterator(node)); 485 486 return &strategy.GetValue(node); 487 } 488 489 490 TWO_KEY_AVL_TREE_TEMPLATE_LIST 491 Value* 492 TWO_KEY_AVL_TREE_CLASS_NAME::FindFirstClosest(const PrimaryKey& key, bool less, 493 Iterator* iterator) 494 { 495 const NodeStrategy& strategy = fTreeMap.GetNodeStrategy(); 496 497 Node* parent = NULL; 498 Node* node = _FindFirst(key, &parent); 499 if (node == NULL) { 500 // not found -- try to get the closest node 501 if (parent == NULL) 502 return NULL; 503 504 node = parent; 505 int expectedCmp = less ? 1 : -1; 506 int cmp = fPrimaryKeyCompare(key, 507 fGetPrimaryKey(strategy.GetValue(strategy.GetNode(node)))); 508 509 if (cmp != expectedCmp) { 510 // The node's value is less although we were asked for a greater 511 // value, or the other way around. We need to iterate to the next 512 // node in the respective direction. If there is no node, we fail. 513 node = less ? Previous(node) : Next(node); 514 if (node == NULL) 515 return NULL; 516 } 517 } 518 519 if (iterator != NULL) 520 iterator->_SetTo(fTreeMap.GetIterator(node)); 521 522 return &strategy.GetValue(node); 523 } 524 525 526 TWO_KEY_AVL_TREE_TEMPLATE_LIST 527 Value* 528 TWO_KEY_AVL_TREE_CLASS_NAME::FindLast(const PrimaryKey& key, 529 Iterator* iterator) 530 { 531 const NodeStrategy& strategy = fTreeMap.GetNodeStrategy(); 532 Node* node = fTreeMap.RootNode(); 533 534 while (node) { 535 int cmp = fPrimaryKeyCompare(key, fGetPrimaryKey( 536 strategy.GetValue(node))); 537 if (cmp == 0) { 538 // found a matching node, now get the right-most node with that key 539 while (node->right && fPrimaryKeyCompare(key, 540 fGetPrimaryKey(strategy.GetValue( 541 strategy.GetNode(node->right)))) == 0) { 542 node = strategy.GetNode(node->right); 543 } 544 if (iterator) 545 iterator->_SetTo(fTreeMap.GetIterator(node)); 546 return &strategy.GetValue(node); 547 } 548 549 if (cmp < 0) 550 node = strategy.GetNode(node->left); 551 else 552 node = strategy.GetNode(node->right); 553 } 554 return NULL; 555 } 556 557 558 TWO_KEY_AVL_TREE_TEMPLATE_LIST 559 Value* 560 TWO_KEY_AVL_TREE_CLASS_NAME::Find(const PrimaryKey& primaryKey, 561 const SecondaryKey& secondaryKey, Iterator* iterator) 562 { 563 TreeMapIterator it = fTreeMap.Find(Key(primaryKey, secondaryKey)); 564 if (iterator) 565 iterator->_SetTo(it); 566 return it.CurrentValuePointer(); 567 } 568 569 570 TWO_KEY_AVL_TREE_TEMPLATE_LIST 571 void 572 TWO_KEY_AVL_TREE_CLASS_NAME::GetIterator(Iterator* iterator) 573 { 574 TreeMapIterator it = fTreeMap.GetIterator(); 575 it.Next(); 576 // Our iterator needs to point to the first entry already. 577 iterator->_SetTo(it); 578 } 579 580 581 TWO_KEY_AVL_TREE_TEMPLATE_LIST 582 void 583 TWO_KEY_AVL_TREE_CLASS_NAME::GetIterator(Node* node, Iterator* iterator) 584 { 585 iterator->_SetTo(fTreeMap.GetIterator(node)); 586 } 587 588 589 TWO_KEY_AVL_TREE_TEMPLATE_LIST 590 status_t 591 TWO_KEY_AVL_TREE_CLASS_NAME::Insert(const Value& value, Iterator* iterator) 592 { 593 NodeStrategy& strategy 594 = const_cast<NodeStrategy&>(fTreeMap.GetNodeStrategy()); 595 596 TreeMapIterator it; 597 status_t status = fTreeMap.Insert(strategy.GetValueKey(value), value, &it); 598 if (status != B_OK || !iterator) 599 return status; 600 601 iterator->_SetTo(it); 602 return B_OK; 603 } 604 605 606 TWO_KEY_AVL_TREE_TEMPLATE_LIST 607 status_t 608 TWO_KEY_AVL_TREE_CLASS_NAME::Insert(const Value& value, Node** _node) 609 { 610 NodeStrategy& strategy 611 = const_cast<NodeStrategy&>(fTreeMap.GetNodeStrategy()); 612 613 return fTreeMap.Insert(strategy.GetValueKey(value), value, _node); 614 } 615 616 617 TWO_KEY_AVL_TREE_TEMPLATE_LIST 618 status_t 619 TWO_KEY_AVL_TREE_CLASS_NAME::Remove(const PrimaryKey& primaryKey, 620 const SecondaryKey& secondaryKey) 621 { 622 return fTreeMap.Remove(Key(primaryKey, secondaryKey)); 623 } 624 625 626 TWO_KEY_AVL_TREE_TEMPLATE_LIST 627 status_t 628 TWO_KEY_AVL_TREE_CLASS_NAME::Remove(Node* node) 629 { 630 return fTreeMap.Remove(node); 631 } 632 633 634 TWO_KEY_AVL_TREE_TEMPLATE_LIST 635 typename TWO_KEY_AVL_TREE_CLASS_NAME::Node* 636 TWO_KEY_AVL_TREE_CLASS_NAME::_FindFirst(const PrimaryKey& key, 637 Node** _parent) const 638 { 639 const NodeStrategy& strategy = fTreeMap.GetNodeStrategy(); 640 Node* node = fTreeMap.RootNode(); 641 Node* parent = NULL; 642 643 while (node) { 644 int cmp = fPrimaryKeyCompare(key, fGetPrimaryKey( 645 strategy.GetValue(node))); 646 if (cmp == 0) { 647 // found a matching node, now get the left-most node with that key 648 while (node->left && fPrimaryKeyCompare(key, 649 fGetPrimaryKey(strategy.GetValue( 650 strategy.GetNode(node->left)))) == 0) { 651 node = strategy.GetNode(node->left); 652 } 653 654 return node; 655 } 656 657 parent = node; 658 659 if (cmp < 0) 660 node = strategy.GetNode(node->left); 661 else 662 node = strategy.GetNode(node->right); 663 } 664 665 if (_parent != NULL) 666 *_parent = parent; 667 668 return NULL; 669 } 670 671 672 #endif // TWO_KEY_AVL_TREE_H 673