xref: /haiku/src/add-ons/kernel/file_systems/packagefs/util/TwoKeyAVLTree.h (revision ed24eb5ff12640d052171c6a7feba37fab8a75d1)
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