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
16template<typename PrimaryKey, typename SecondaryKey>
17class TwoKeyAVLTreeKey {
18public:
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
45template<typename PrimaryKey, typename SecondaryKey,
46	typename PrimaryKeyCompare, typename SecondaryKeyCompare>
47class TwoKeyAVLTreeKeyCompare {
48private:
49	typedef TwoKeyAVLTreeKey<PrimaryKey, SecondaryKey> Key;
50
51public:
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
68private:
69	PrimaryKeyCompare	fPrimaryKeyCompare;
70	SecondaryKeyCompare	fSecondaryKeyCompare;
71};
72
73
74// #pragma mark - TwoKeyAVLTreeGetKey
75
76
77template<typename Value, typename PrimaryKey, typename SecondaryKey,
78	typename GetPrimaryKey, typename GetSecondaryKey>
79class TwoKeyAVLTreeGetKey {
80private:
81	typedef TwoKeyAVLTreeKey<PrimaryKey, SecondaryKey> Key;
82
83public:
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
97private:
98	GetPrimaryKey	fGetPrimaryKey;
99	GetSecondaryKey	fGetSecondaryKey;
100};
101
102
103// #pragma mark - TwoKeyAVLTreeStandardCompare
104
105
106template<typename Value>
107class TwoKeyAVLTreeStandardCompare {
108public:
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
123template<typename Value, typename Key>
124class TwoKeyAVLTreeStandardGetKey {
125public:
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
141template <typename PrimaryKey, typename SecondaryKey, typename Value,
142	typename PrimaryKeyCompare, typename SecondaryKeyCompare,
143	typename GetPrimaryKey, typename GetSecondaryKey>
144class TwoKeyAVLTreeNodeStrategy {
145public:
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
250private:
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
282template<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> >
288class TwoKeyAVLTree {
289public:
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
301public:
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
336private:
337			Node*				_FindFirst(const PrimaryKey& key,
338									Node** _parent) const;
339
340private:
341			TreeMap				fTreeMap;
342			PrimaryKeyCompare	fPrimaryKeyCompare;
343			GetPrimaryKey		fGetPrimaryKey;
344};
345
346
347// #pragma mark - Iterator
348
349
350TWO_KEY_AVL_TREE_TEMPLATE_LIST
351class TWO_KEY_AVL_TREE_CLASS_NAME::Iterator {
352public:
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
393private:
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
418private:
419	TWO_KEY_AVL_TREE_CLASS_NAME::TreeMapIterator fIterator;
420};
421
422
423TWO_KEY_AVL_TREE_TEMPLATE_LIST
424TWO_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
434TWO_KEY_AVL_TREE_TEMPLATE_LIST
435TWO_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
448TWO_KEY_AVL_TREE_TEMPLATE_LIST
449TWO_KEY_AVL_TREE_CLASS_NAME::~TwoKeyAVLTree()
450{
451}
452
453
454TWO_KEY_AVL_TREE_TEMPLATE_LIST
455typename TWO_KEY_AVL_TREE_CLASS_NAME::Node*
456TWO_KEY_AVL_TREE_CLASS_NAME::Previous(Node* node) const
457{
458	return fTreeMap.Previous(node);
459}
460
461
462TWO_KEY_AVL_TREE_TEMPLATE_LIST
463typename TWO_KEY_AVL_TREE_CLASS_NAME::Node*
464TWO_KEY_AVL_TREE_CLASS_NAME::Next(Node* node) const
465{
466	return fTreeMap.Next(node);
467}
468
469
470TWO_KEY_AVL_TREE_TEMPLATE_LIST
471Value*
472TWO_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
488TWO_KEY_AVL_TREE_TEMPLATE_LIST
489Value*
490TWO_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
524TWO_KEY_AVL_TREE_TEMPLATE_LIST
525Value*
526TWO_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
556TWO_KEY_AVL_TREE_TEMPLATE_LIST
557Value*
558TWO_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
568TWO_KEY_AVL_TREE_TEMPLATE_LIST
569void
570TWO_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
579TWO_KEY_AVL_TREE_TEMPLATE_LIST
580void
581TWO_KEY_AVL_TREE_CLASS_NAME::GetIterator(Node* node, Iterator* iterator)
582{
583	iterator->_SetTo(fTreeMap.GetIterator(node));
584}
585
586
587TWO_KEY_AVL_TREE_TEMPLATE_LIST
588status_t
589TWO_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
604TWO_KEY_AVL_TREE_TEMPLATE_LIST
605status_t
606TWO_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
615TWO_KEY_AVL_TREE_TEMPLATE_LIST
616status_t
617TWO_KEY_AVL_TREE_CLASS_NAME::Remove(const PrimaryKey& primaryKey,
618	const SecondaryKey& secondaryKey)
619{
620	return fTreeMap.Remove(Key(primaryKey, secondaryKey));
621}
622
623
624TWO_KEY_AVL_TREE_TEMPLATE_LIST
625status_t
626TWO_KEY_AVL_TREE_CLASS_NAME::Remove(Node* node)
627{
628	return fTreeMap.Remove(node);
629}
630
631
632TWO_KEY_AVL_TREE_TEMPLATE_LIST
633typename TWO_KEY_AVL_TREE_CLASS_NAME::Node*
634TWO_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