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:
TwoKeyAVLTreeKey(const PrimaryKey & primary,const SecondaryKey & secondary)15 inline TwoKeyAVLTreeKey(const PrimaryKey &primary,
16 const SecondaryKey &secondary)
17 : primary(primary),
18 secondary(secondary),
19 use_secondary(true)
20 {
21 }
22
TwoKeyAVLTreeKey(const PrimaryKey * primary)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:
TwoKeyAVLTreeKeyCompare(const PrimaryKeyCompare & primary,const SecondaryKeyCompare & secondary)43 inline TwoKeyAVLTreeKeyCompare(const PrimaryKeyCompare &primary,
44 const SecondaryKeyCompare &secondary)
45 : fPrimaryKeyCompare(primary), fSecondaryKeyCompare(secondary) {}
46
operator()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:
TwoKeyAVLTreeGetKey(const GetPrimaryKey & getPrimary,const GetSecondaryKey & getSecondary)69 TwoKeyAVLTreeGetKey(const GetPrimaryKey &getPrimary,
70 const GetSecondaryKey &getSecondary)
71 : fGetPrimaryKey(getPrimary),
72 fGetSecondaryKey(getSecondary)
73 {
74 }
75
operator()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:
operator()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:
operator()108 inline const Key &operator()(const Value &a) const
109 {
110 return a;
111 }
112
operator()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())
fPrimaryKeyCompare(primaryKeyCompare)133 : fPrimaryKeyCompare(primaryKeyCompare),
134 fSecondaryKeyCompare(secondaryKeyCompare),
135 fGetPrimaryKey(getPrimaryKey),
136 fGetSecondaryKey(getSecondaryKey)
137 {
138 }
139
140 struct Node : AVLTreeNode {
NodeNode141 Node(const Value &value)
142 : AVLTreeNode(),
143 value(value)
144 {
145 }
146
147 Value value;
148 };
149
Allocate(const Key & key,const Value & value)150 inline Node* Allocate(const Key& key, const Value& value) const
151 {
152 return new(nothrow) Node(value);
153 }
154
Free(Node * node)155 inline void Free(Node* node) const
156 {
157 delete node;
158 }
159
160 // internal use (not part of the strategy)
GetValueKey(const Value & value)161 inline Key GetValueKey(const Value& value) const
162 {
163 return Key(fGetPrimaryKey(value), fGetSecondaryKey(value));
164 }
165
GetKey(Node * node)166 inline Key GetKey(Node* node) const
167 {
168 return GetValueKey(node->value);
169 }
170
GetValue(Node * node)171 inline Value& GetValue(Node* node) const
172 {
173 return node->value;
174 }
175
GetAVLTreeNode(Node * node)176 inline AVLTreeNode* GetAVLTreeNode(Node* node) const
177 {
178 return node;
179 }
180
GetNode(AVLTreeNode * node)181 inline Node* GetNode(AVLTreeNode* node) const
182 {
183 return static_cast<Node*>(node);
184 }
185
CompareKeyNode(const Key & a,const Node * b)186 inline int CompareKeyNode(const Key& a, const Node* b) const
187 {
188 return _CompareKeys(a, GetKey(const_cast<Node*>(b)));
189 }
190
CompareNodes(const Node * a,const Node * b)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:
_CompareKeys(const Key & a,const Key & b)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
CountItems()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
Iterator()279 inline Iterator()
280 : fIterator()
281 {
282 }
283
~Iterator()284 inline ~Iterator()
285 {
286 }
287
GetCurrent()288 inline Value *GetCurrent()
289 {
290 return fIterator.CurrentValuePointer();
291 }
292
GetNext()293 inline Value *GetNext()
294 {
295 fIterator.Next();
296 return GetCurrent();
297 }
298
GetPrevious()299 inline Value *GetPrevious()
300 {
301 fIterator.Previous();
302 return GetCurrent();
303 }
304
Remove()305 inline void Remove()
306 {
307 fIterator.Remove();
308 }
309
310 private:
311 friend class TWO_KEY_AVL_TREE_CLASS_NAME;
312
Iterator(const Iterator & other)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
Iterator(const TreeMapIterator & iterator)323 Iterator(const TreeMapIterator &iterator)
324 : fIterator()
325 {
326 }
327
_SetTo(const TreeMapIterator & iterator)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
TwoKeyAVLTree()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
TwoKeyAVLTree(const PrimaryKeyCompare & primaryCompare,const GetPrimaryKey & getPrimary,const SecondaryKeyCompare & secondaryCompare,const GetSecondaryKey & getSecondary)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
~TwoKeyAVLTree()365 TWO_KEY_AVL_TREE_CLASS_NAME::~TwoKeyAVLTree()
366 {
367 }
368
369
370 // FindFirst
371 TWO_KEY_AVL_TREE_TEMPLATE_LIST
372 Value *
FindFirst(const PrimaryKey & key,Iterator * iterator)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 *
FindLast(const PrimaryKey & key,Iterator * iterator)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 *
Find(const PrimaryKey & primaryKey,const SecondaryKey & secondaryKey,Iterator * iterator)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
GetIterator(Iterator * iterator)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
Insert(const Value & value,Iterator * iterator)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
Remove(const PrimaryKey & primaryKey,const SecondaryKey & secondaryKey)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