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