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 12template<typename PrimaryKey, typename SecondaryKey> 13class TwoKeyAVLTreeKey { 14public: 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 36template<typename PrimaryKey, typename SecondaryKey, 37 typename PrimaryKeyCompare, typename SecondaryKeyCompare> 38class TwoKeyAVLTreeKeyCompare { 39private: 40 typedef TwoKeyAVLTreeKey<PrimaryKey, SecondaryKey> Key; 41 42public: 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 55private: 56 PrimaryKeyCompare fPrimaryKeyCompare; 57 SecondaryKeyCompare fSecondaryKeyCompare; 58}; 59 60// TwoKeyAVLTreeGetKey 61template<typename Value, typename PrimaryKey, typename SecondaryKey, 62 typename GetPrimaryKey, typename GetSecondaryKey> 63class TwoKeyAVLTreeGetKey 64{ 65private: 66 typedef TwoKeyAVLTreeKey<PrimaryKey, SecondaryKey> Key; 67 68public: 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 81private: 82 GetPrimaryKey fGetPrimaryKey; 83 GetSecondaryKey fGetSecondaryKey; 84}; 85 86 87// TwoKeyAVLTreeStandardCompare 88template<typename Value> 89class TwoKeyAVLTreeStandardCompare 90{ 91public: 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 104template<typename Value, typename Key> 105class TwoKeyAVLTreeStandardGetKey 106{ 107public: 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 121template <typename PrimaryKey, typename SecondaryKey, typename Value, 122 typename PrimaryKeyCompare, typename SecondaryKeyCompare, 123 typename GetPrimaryKey, typename GetSecondaryKey> 124class TwoKeyAVLTreeNodeStrategy { 125public: 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 197private: 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 224template<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> > 231class TwoKeyAVLTree { 232public: 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 265private: 266 TreeMap fTreeMap; 267 PrimaryKeyCompare fPrimaryKeyCompare; 268 GetPrimaryKey fGetPrimaryKey; 269}; 270 271 272// Iterator 273TWO_KEY_AVL_TREE_TEMPLATE_LIST 274class TWO_KEY_AVL_TREE_CLASS_NAME::Iterator { 275public: 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 310private: 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 333private: 334 TWO_KEY_AVL_TREE_CLASS_NAME::TreeMapIterator fIterator; 335}; 336 337 338// constructor 339TWO_KEY_AVL_TREE_TEMPLATE_LIST 340TWO_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 350TWO_KEY_AVL_TREE_TEMPLATE_LIST 351TWO_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 364TWO_KEY_AVL_TREE_TEMPLATE_LIST 365TWO_KEY_AVL_TREE_CLASS_NAME::~TwoKeyAVLTree() 366{ 367} 368 369 370// FindFirst 371TWO_KEY_AVL_TREE_TEMPLATE_LIST 372Value * 373TWO_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 403TWO_KEY_AVL_TREE_TEMPLATE_LIST 404Value * 405TWO_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 435TWO_KEY_AVL_TREE_TEMPLATE_LIST 436Value * 437TWO_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 448TWO_KEY_AVL_TREE_TEMPLATE_LIST 449void 450TWO_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 459TWO_KEY_AVL_TREE_TEMPLATE_LIST 460status_t 461TWO_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 476TWO_KEY_AVL_TREE_TEMPLATE_LIST 477status_t 478TWO_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