1 /* 2 * Copyright 2003-2007, Ingo Weinhold <bonefish@cs.tu-berlin.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 #include <util/AVLTreeMap.h> 9 10 11 // TwoKeyAVLTreeKey 12 template<typename PrimaryKey, typename SecondaryKey> 13 class TwoKeyAVLTreeKey { 14 public: 15 inline TwoKeyAVLTreeKey(const PrimaryKey &primary, 16 const SecondaryKey &secondary) 17 : primary(primary), 18 secondary(secondary), 19 use_secondary(true) 20 { 21 } 22 23 inline TwoKeyAVLTreeKey(const PrimaryKey *primary) 24 : primary(primary), 25 secondary(NULL), 26 use_secondary(false) 27 { 28 } 29 30 PrimaryKey primary; 31 SecondaryKey secondary; 32 bool use_secondary; 33 }; 34 35 // TwoKeyAVLTreeKeyCompare 36 template<typename PrimaryKey, typename SecondaryKey, 37 typename PrimaryKeyCompare, typename SecondaryKeyCompare> 38 class TwoKeyAVLTreeKeyCompare { 39 private: 40 typedef TwoKeyAVLTreeKey<PrimaryKey, SecondaryKey> Key; 41 42 public: 43 inline TwoKeyAVLTreeKeyCompare(const PrimaryKeyCompare &primary, 44 const SecondaryKeyCompare &secondary) 45 : fPrimaryKeyCompare(primary), fSecondaryKeyCompare(secondary) {} 46 47 inline int operator()(const Key &a, const Key &b) const 48 { 49 int result = fPrimaryKeyCompare(a.primary, b.primary); 50 if (result == 0 && a.use_secondary && b.use_secondary) 51 result = fSecondaryKeyCompare(a.secondary, b.secondary); 52 return result; 53 } 54 55 private: 56 PrimaryKeyCompare fPrimaryKeyCompare; 57 SecondaryKeyCompare fSecondaryKeyCompare; 58 }; 59 60 // TwoKeyAVLTreeGetKey 61 template<typename Value, typename PrimaryKey, typename SecondaryKey, 62 typename GetPrimaryKey, typename GetSecondaryKey> 63 class TwoKeyAVLTreeGetKey 64 { 65 private: 66 typedef TwoKeyAVLTreeKey<PrimaryKey, SecondaryKey> Key; 67 68 public: 69 TwoKeyAVLTreeGetKey(const GetPrimaryKey &getPrimary, 70 const GetSecondaryKey &getSecondary) 71 : fGetPrimaryKey(getPrimary), 72 fGetSecondaryKey(getSecondary) 73 { 74 } 75 76 inline Key operator()(const Value &a) const 77 { 78 return Key(fGetPrimaryKey(a), fGetSecondaryKey(a)); 79 } 80 81 private: 82 GetPrimaryKey fGetPrimaryKey; 83 GetSecondaryKey fGetSecondaryKey; 84 }; 85 86 87 // TwoKeyAVLTreeStandardCompare 88 template<typename Value> 89 class TwoKeyAVLTreeStandardCompare 90 { 91 public: 92 inline int operator()(const Value &a, const Value &b) const 93 { 94 if (a < b) 95 return -1; 96 else if (a > b) 97 return 1; 98 return 0; 99 } 100 }; 101 102 103 // TwoKeyAVLTreeStandardGetKey 104 template<typename Value, typename Key> 105 class TwoKeyAVLTreeStandardGetKey 106 { 107 public: 108 inline const Key &operator()(const Value &a) const 109 { 110 return a; 111 } 112 113 inline Key &operator()(Value &a) const 114 { 115 return a; 116 } 117 }; 118 119 120 // TwoKeyAVLTreeNodeStrategy 121 template <typename PrimaryKey, typename SecondaryKey, typename Value, 122 typename PrimaryKeyCompare, typename SecondaryKeyCompare, 123 typename GetPrimaryKey, typename GetSecondaryKey> 124 class TwoKeyAVLTreeNodeStrategy { 125 public: 126 typedef TwoKeyAVLTreeKey<PrimaryKey, SecondaryKey> Key; 127 128 TwoKeyAVLTreeNodeStrategy( 129 const PrimaryKeyCompare& primaryKeyCompare = PrimaryKeyCompare(), 130 const SecondaryKeyCompare& secondaryKeyCompare = SecondaryKeyCompare(), 131 const GetPrimaryKey& getPrimaryKey = GetPrimaryKey(), 132 const GetSecondaryKey& getSecondaryKey = GetSecondaryKey()) 133 : fPrimaryKeyCompare(primaryKeyCompare), 134 fSecondaryKeyCompare(secondaryKeyCompare), 135 fGetPrimaryKey(getPrimaryKey), 136 fGetSecondaryKey(getSecondaryKey) 137 { 138 } 139 140 struct Node : AVLTreeNode { 141 Node(const Value &value) 142 : AVLTreeNode(), 143 value(value) 144 { 145 } 146 147 Value value; 148 }; 149 150 inline Node* Allocate(const Key& key, const Value& value) const 151 { 152 return new(nothrow) Node(value); 153 } 154 155 inline void Free(Node* node) const 156 { 157 delete node; 158 } 159 160 // internal use (not part of the strategy) 161 inline Key GetValueKey(const Value& value) const 162 { 163 return Key(fGetPrimaryKey(value), fGetSecondaryKey(value)); 164 } 165 166 inline Key GetKey(Node* node) const 167 { 168 return GetValueKey(node->value); 169 } 170 171 inline Value& GetValue(Node* node) const 172 { 173 return node->value; 174 } 175 176 inline AVLTreeNode* GetAVLTreeNode(Node* node) const 177 { 178 return node; 179 } 180 181 inline Node* GetNode(AVLTreeNode* node) const 182 { 183 return static_cast<Node*>(node); 184 } 185 186 inline int CompareKeyNode(const Key& a, const Node* b) const 187 { 188 return _CompareKeys(a, GetKey(const_cast<Node*>(b))); 189 } 190 191 inline int CompareNodes(const Node* a, const Node* b) const 192 { 193 return _CompareKeys(GetKey(const_cast<Node*>(a)), 194 GetKey(const_cast<Node*>(b))); 195 } 196 197 private: 198 inline int _CompareKeys(const Key& a, const Key& b) const 199 { 200 int result = fPrimaryKeyCompare(a.primary, b.primary); 201 if (result == 0 && a.use_secondary && b.use_secondary) 202 result = fSecondaryKeyCompare(a.secondary, b.secondary); 203 return result; 204 } 205 206 PrimaryKeyCompare fPrimaryKeyCompare; 207 SecondaryKeyCompare fSecondaryKeyCompare; 208 GetPrimaryKey fGetPrimaryKey; 209 GetSecondaryKey fGetSecondaryKey; 210 }; 211 212 213 // for convenience 214 #define TWO_KEY_AVL_TREE_TEMPLATE_LIST template<typename Value, \ 215 typename PrimaryKey, typename PrimaryKeyCompare, \ 216 typename GetPrimaryKey, typename SecondaryKey, \ 217 typename SecondaryKeyCompare, typename GetSecondaryKey> 218 #define TWO_KEY_AVL_TREE_CLASS_NAME TwoKeyAVLTree<Value, PrimaryKey, \ 219 PrimaryKeyCompare, GetPrimaryKey, SecondaryKey, \ 220 SecondaryKeyCompare, GetSecondaryKey> 221 222 223 // TwoKeyAVLTree 224 template<typename Value, typename PrimaryKey, 225 typename PrimaryKeyCompare, typename GetPrimaryKey, 226 typename SecondaryKey = Value, 227 typename SecondaryKeyCompare 228 = TwoKeyAVLTreeStandardCompare<SecondaryKey>, 229 typename GetSecondaryKey 230 = TwoKeyAVLTreeStandardGetKey<Value, SecondaryKey> > 231 class TwoKeyAVLTree { 232 public: 233 typedef TwoKeyAVLTreeKey<PrimaryKey, SecondaryKey> Key; 234 typedef TwoKeyAVLTreeNodeStrategy<PrimaryKey, SecondaryKey, Value, 235 PrimaryKeyCompare, SecondaryKeyCompare, GetPrimaryKey, 236 GetSecondaryKey> NodeStrategy; 237 typedef AVLTreeMap<Key, Value, NodeStrategy> TreeMap; 238 239 typedef typename TreeMap::Iterator TreeMapIterator; 240 typedef typename NodeStrategy::Node Node; 241 242 class Iterator; 243 244 TwoKeyAVLTree(); 245 TwoKeyAVLTree(const PrimaryKeyCompare &primaryCompare, 246 const GetPrimaryKey &getPrimary, 247 const SecondaryKeyCompare &secondaryCompare, 248 const GetSecondaryKey &getSecondary); 249 ~TwoKeyAVLTree(); 250 251 inline int CountItems() const { return fTreeMap.Count(); } 252 253 Value *FindFirst(const PrimaryKey &key, Iterator *iterator = NULL); 254 Value *FindLast(const PrimaryKey &key, Iterator *iterator = NULL); 255 inline Value *Find(const PrimaryKey &primaryKey, 256 const SecondaryKey &secondaryKey, 257 Iterator *iterator = NULL); 258 259 inline void GetIterator(Iterator *iterator); 260 261 inline status_t Insert(const Value &value, Iterator *iterator = NULL); 262 inline status_t Remove(const PrimaryKey &primaryKey, 263 const SecondaryKey &secondaryKey); 264 265 private: 266 TreeMap fTreeMap; 267 PrimaryKeyCompare fPrimaryKeyCompare; 268 GetPrimaryKey fGetPrimaryKey; 269 }; 270 271 272 // Iterator 273 TWO_KEY_AVL_TREE_TEMPLATE_LIST 274 class TWO_KEY_AVL_TREE_CLASS_NAME::Iterator { 275 public: 276 typedef typename TWO_KEY_AVL_TREE_CLASS_NAME::TreeMapIterator 277 TreeMapIterator; 278 279 inline Iterator() 280 : fIterator() 281 { 282 } 283 284 inline ~Iterator() 285 { 286 } 287 288 inline Value *GetCurrent() 289 { 290 return fIterator.CurrentValuePointer(); 291 } 292 293 inline Value *GetNext() 294 { 295 fIterator.Next(); 296 return GetCurrent(); 297 } 298 299 inline Value *GetPrevious() 300 { 301 fIterator.Previous(); 302 return GetCurrent(); 303 } 304 305 inline void Remove() 306 { 307 fIterator.Remove(); 308 } 309 310 private: 311 friend class TWO_KEY_AVL_TREE_CLASS_NAME; 312 313 Iterator(const Iterator& other) 314 : fIterator(other.fIterator) 315 { 316 } 317 318 Iterator &operator=(const Iterator& other) 319 { 320 fIterator = other.fIterator; 321 } 322 323 Iterator(const TreeMapIterator &iterator) 324 : fIterator() 325 { 326 } 327 328 inline void _SetTo(const TreeMapIterator &iterator) 329 { 330 fIterator = iterator; 331 } 332 333 private: 334 TWO_KEY_AVL_TREE_CLASS_NAME::TreeMapIterator fIterator; 335 }; 336 337 338 // constructor 339 TWO_KEY_AVL_TREE_TEMPLATE_LIST 340 TWO_KEY_AVL_TREE_CLASS_NAME::TwoKeyAVLTree() 341 : fTreeMap(NodeStrategy(PrimaryKeyCompare(), SecondaryKeyCompare(), 342 GetPrimaryKey(), GetSecondaryKey())), 343 fPrimaryKeyCompare(PrimaryKeyCompare()), 344 fGetPrimaryKey(GetPrimaryKey()) 345 { 346 } 347 348 349 // constructor 350 TWO_KEY_AVL_TREE_TEMPLATE_LIST 351 TWO_KEY_AVL_TREE_CLASS_NAME::TwoKeyAVLTree( 352 const PrimaryKeyCompare &primaryCompare, const GetPrimaryKey &getPrimary, 353 const SecondaryKeyCompare &secondaryCompare, 354 const GetSecondaryKey &getSecondary) 355 : fTreeMap(NodeStrategy(primaryCompare, secondaryCompare, getPrimary, 356 getSecondary)), 357 fPrimaryKeyCompare(primaryCompare), 358 fGetPrimaryKey(getPrimary) 359 { 360 } 361 362 363 // destructor 364 TWO_KEY_AVL_TREE_TEMPLATE_LIST 365 TWO_KEY_AVL_TREE_CLASS_NAME::~TwoKeyAVLTree() 366 { 367 } 368 369 370 // FindFirst 371 TWO_KEY_AVL_TREE_TEMPLATE_LIST 372 Value * 373 TWO_KEY_AVL_TREE_CLASS_NAME::FindFirst(const PrimaryKey &key, 374 Iterator *iterator) 375 { 376 const NodeStrategy& strategy = fTreeMap.GetNodeStrategy(); 377 Node *node = fTreeMap.RootNode(); 378 379 while (node) { 380 int cmp = fPrimaryKeyCompare(key, fGetPrimaryKey( 381 strategy.GetValue(node))); 382 if (cmp == 0) { 383 // found a matching node, now get the left-most node with that key 384 while (node->left && fPrimaryKeyCompare(key, 385 fGetPrimaryKey(strategy.GetValue( 386 strategy.GetNode(node->left)))) == 0) { 387 node = strategy.GetNode(node->left); 388 } 389 if (iterator) 390 iterator->_SetTo(fTreeMap.GetIterator(node)); 391 return &strategy.GetValue(node); 392 } 393 394 if (cmp < 0) 395 node = strategy.GetNode(node->left); 396 else 397 node = strategy.GetNode(node->right); 398 } 399 return NULL; 400 } 401 402 // FindLast 403 TWO_KEY_AVL_TREE_TEMPLATE_LIST 404 Value * 405 TWO_KEY_AVL_TREE_CLASS_NAME::FindLast(const PrimaryKey &key, 406 Iterator *iterator) 407 { 408 const NodeStrategy& strategy = fTreeMap.GetNodeStrategy(); 409 Node *node = fTreeMap.RootNode(); 410 411 while (node) { 412 int cmp = fPrimaryKeyCompare(key, fGetPrimaryKey( 413 strategy.GetValue(node))); 414 if (cmp == 0) { 415 // found a matching node, now get the right-most node with that key 416 while (node->right && fPrimaryKeyCompare(key, 417 fGetPrimaryKey(strategy.GetValue( 418 strategy.GetNode(node->right)))) == 0) { 419 node = strategy.GetNode(node->right); 420 } 421 if (iterator) 422 iterator->_SetTo(fTreeMap.GetIterator(node)); 423 return &strategy.GetValue(node); 424 } 425 426 if (cmp < 0) 427 node = strategy.GetNode(node->left); 428 else 429 node = strategy.GetNode(node->right); 430 } 431 return NULL; 432 } 433 434 // Find 435 TWO_KEY_AVL_TREE_TEMPLATE_LIST 436 Value * 437 TWO_KEY_AVL_TREE_CLASS_NAME::Find(const PrimaryKey &primaryKey, 438 const SecondaryKey &secondaryKey, Iterator *iterator) 439 { 440 441 TreeMapIterator it = fTreeMap.Find(Key(primaryKey, secondaryKey)); 442 if (iterator) 443 iterator->_SetTo(it); 444 return it.CurrentValuePointer(); 445 } 446 447 // GetIterator 448 TWO_KEY_AVL_TREE_TEMPLATE_LIST 449 void 450 TWO_KEY_AVL_TREE_CLASS_NAME::GetIterator(Iterator *iterator) 451 { 452 TreeMapIterator it = fTreeMap.GetIterator(); 453 it.Next(); 454 // Our iterator needs to point to the first entry already. 455 iterator->_SetTo(it); 456 } 457 458 // Insert 459 TWO_KEY_AVL_TREE_TEMPLATE_LIST 460 status_t 461 TWO_KEY_AVL_TREE_CLASS_NAME::Insert(const Value &value, Iterator *iterator) 462 { 463 NodeStrategy& strategy 464 = const_cast<NodeStrategy&>(fTreeMap.GetNodeStrategy()); 465 466 TreeMapIterator it; 467 status_t status = fTreeMap.Insert(strategy.GetValueKey(value), value, &it); 468 if (status != B_OK || !iterator) 469 return status; 470 471 iterator->_SetTo(it); 472 return B_OK; 473 } 474 475 // Remove 476 TWO_KEY_AVL_TREE_TEMPLATE_LIST 477 status_t 478 TWO_KEY_AVL_TREE_CLASS_NAME::Remove(const PrimaryKey &primaryKey, 479 const SecondaryKey &secondaryKey) 480 { 481 return fTreeMap.Remove(Key(primaryKey, secondaryKey)); 482 } 483 484 #endif // TWO_KEY_AVL_TREE_H 485