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 <util/AVLTreeMap.h>
10
11
12// #pragma mark - TwoKeyAVLTreeKey
13
14
15template<typename PrimaryKey, typename SecondaryKey>
16class TwoKeyAVLTreeKey {
17public:
18	inline TwoKeyAVLTreeKey(const PrimaryKey& primary,
19		const SecondaryKey& secondary)
20		:
21		primary(primary),
22		secondary(secondary),
23		use_secondary(true)
24	{
25	}
26
27	inline TwoKeyAVLTreeKey(const PrimaryKey* primary)
28		:
29		primary(primary),
30		secondary(NULL),
31		use_secondary(false)
32	{
33	}
34
35	PrimaryKey		primary;
36	SecondaryKey	secondary;
37	bool			use_secondary;
38};
39
40
41// #pragma mark - TwoKeyAVLTreeKeyCompare
42
43
44template<typename PrimaryKey, typename SecondaryKey,
45	typename PrimaryKeyCompare, typename SecondaryKeyCompare>
46class TwoKeyAVLTreeKeyCompare {
47private:
48	typedef TwoKeyAVLTreeKey<PrimaryKey, SecondaryKey> Key;
49
50public:
51	inline TwoKeyAVLTreeKeyCompare(const PrimaryKeyCompare& primary,
52								   const SecondaryKeyCompare& secondary)
53		:
54		fPrimaryKeyCompare(primary),
55		fSecondaryKeyCompare(secondary)
56	{
57	}
58
59	inline int operator()(const Key& a, const Key& b) const
60	{
61		int result = fPrimaryKeyCompare(a.primary, b.primary);
62		if (result == 0 && a.use_secondary && b.use_secondary)
63			result = fSecondaryKeyCompare(a.secondary, b.secondary);
64		return result;
65	}
66
67private:
68	PrimaryKeyCompare	fPrimaryKeyCompare;
69	SecondaryKeyCompare	fSecondaryKeyCompare;
70};
71
72
73// #pragma mark - TwoKeyAVLTreeGetKey
74
75
76template<typename Value, typename PrimaryKey, typename SecondaryKey,
77	typename GetPrimaryKey, typename GetSecondaryKey>
78class TwoKeyAVLTreeGetKey {
79private:
80	typedef TwoKeyAVLTreeKey<PrimaryKey, SecondaryKey> Key;
81
82public:
83	TwoKeyAVLTreeGetKey(const GetPrimaryKey& getPrimary,
84		const GetSecondaryKey& getSecondary)
85		:
86		fGetPrimaryKey(getPrimary),
87		fGetSecondaryKey(getSecondary)
88	{
89	}
90
91	inline Key operator()(const Value& a) const
92	{
93		return Key(fGetPrimaryKey(a), fGetSecondaryKey(a));
94	}
95
96private:
97	GetPrimaryKey	fGetPrimaryKey;
98	GetSecondaryKey	fGetSecondaryKey;
99};
100
101
102// #pragma mark - TwoKeyAVLTreeStandardCompare
103
104
105template<typename Value>
106class TwoKeyAVLTreeStandardCompare {
107public:
108	inline int operator()(const Value& a, const Value& b) const
109	{
110		if (a < b)
111			return -1;
112		else if (a > b)
113			return 1;
114		return 0;
115	}
116};
117
118
119// #pragma mark - TwoKeyAVLTreeStandardGetKey
120
121
122template<typename Value, typename Key>
123class TwoKeyAVLTreeStandardGetKey {
124public:
125	inline const Key& operator()(const Value& a) const
126	{
127		return a;
128	}
129
130	inline Key& operator()(Value& a) const
131	{
132		return a;
133	}
134};
135
136
137// #pragma mark - TwoKeyAVLTreeNodeStrategy
138
139
140template <typename PrimaryKey, typename SecondaryKey, typename Value,
141	typename PrimaryKeyCompare, typename SecondaryKeyCompare,
142	typename GetPrimaryKey, typename GetSecondaryKey>
143class TwoKeyAVLTreeNodeStrategy {
144public:
145	typedef TwoKeyAVLTreeKey<PrimaryKey, SecondaryKey> Key;
146
147	TwoKeyAVLTreeNodeStrategy(
148		const PrimaryKeyCompare& primaryKeyCompare = PrimaryKeyCompare(),
149		const SecondaryKeyCompare& secondaryKeyCompare = SecondaryKeyCompare(),
150		const GetPrimaryKey& getPrimaryKey = GetPrimaryKey(),
151		const GetSecondaryKey& getSecondaryKey = GetSecondaryKey())
152		:
153		fPrimaryKeyCompare(primaryKeyCompare),
154		fSecondaryKeyCompare(secondaryKeyCompare),
155		fGetPrimaryKey(getPrimaryKey),
156		fGetSecondaryKey(getSecondaryKey)
157	{
158	}
159
160	struct Node : AVLTreeNode {
161		Node(const Value& value)
162			:
163			AVLTreeNode(),
164			value(value)
165		{
166		}
167
168		Value	value;
169	};
170
171	inline Node* Allocate(const Key& key, const Value& value) const
172	{
173		return new(nothrow) Node(value);
174	}
175
176	inline void Free(Node* node) const
177	{
178		delete node;
179	}
180
181	// internal use (not part of the strategy)
182	inline Key GetValueKey(const Value& value) const
183	{
184		return Key(fGetPrimaryKey(value), fGetSecondaryKey(value));
185	}
186
187	inline Key GetKey(Node* node) const
188	{
189		return GetValueKey(node->value);
190	}
191
192	inline Value& GetValue(Node* node) const
193	{
194		return node->value;
195	}
196
197	inline AVLTreeNode* GetAVLTreeNode(Node* node) const
198	{
199		return node;
200	}
201
202	inline Node* GetNode(AVLTreeNode* node) const
203	{
204		return static_cast<Node*>(node);
205	}
206
207	inline int CompareKeyNode(const Key& a, const Node* b) const
208	{
209		return _CompareKeys(a, GetKey(const_cast<Node*>(b)));
210	}
211
212	inline int CompareNodes(const Node* a, const Node* b) const
213	{
214		return _CompareKeys(GetKey(const_cast<Node*>(a)),
215			GetKey(const_cast<Node*>(b)));
216	}
217
218private:
219	inline int _CompareKeys(const Key& a, const Key& b) const
220	{
221		int result = fPrimaryKeyCompare(a.primary, b.primary);
222		if (result == 0 && a.use_secondary && b.use_secondary)
223			result = fSecondaryKeyCompare(a.secondary, b.secondary);
224		return result;
225	}
226
227	PrimaryKeyCompare	fPrimaryKeyCompare;
228	SecondaryKeyCompare	fSecondaryKeyCompare;
229	GetPrimaryKey		fGetPrimaryKey;
230	GetSecondaryKey		fGetSecondaryKey;
231};
232
233
234// for convenience
235#define TWO_KEY_AVL_TREE_TEMPLATE_LIST template<typename Value, \
236	typename PrimaryKey, typename PrimaryKeyCompare, \
237	typename GetPrimaryKey, typename SecondaryKey, \
238	typename SecondaryKeyCompare, typename GetSecondaryKey>
239#define TWO_KEY_AVL_TREE_CLASS_NAME TwoKeyAVLTree<Value, PrimaryKey, \
240	PrimaryKeyCompare, GetPrimaryKey, SecondaryKey, \
241	SecondaryKeyCompare, GetSecondaryKey>
242
243
244// #pragma mark - TwoKeyAVLTree
245
246
247template<typename Value, typename PrimaryKey,
248	typename PrimaryKeyCompare, typename GetPrimaryKey,
249	typename SecondaryKey = Value,
250	typename SecondaryKeyCompare = TwoKeyAVLTreeStandardCompare<SecondaryKey>,
251	typename GetSecondaryKey
252		= TwoKeyAVLTreeStandardGetKey<Value, SecondaryKey> >
253class TwoKeyAVLTree {
254public:
255			typedef TwoKeyAVLTreeKey<PrimaryKey, SecondaryKey> Key;
256			typedef TwoKeyAVLTreeNodeStrategy<PrimaryKey, SecondaryKey, Value,
257				PrimaryKeyCompare, SecondaryKeyCompare, GetPrimaryKey,
258				GetSecondaryKey> NodeStrategy;
259			typedef AVLTreeMap<Key, Value, NodeStrategy> TreeMap;
260
261			typedef typename TreeMap::Iterator	TreeMapIterator;
262			typedef typename NodeStrategy::Node Node;
263
264			class Iterator;
265
266public:
267								TwoKeyAVLTree();
268								TwoKeyAVLTree(
269									const PrimaryKeyCompare& primaryCompare,
270									const GetPrimaryKey& getPrimary,
271									const SecondaryKeyCompare& secondaryCompare,
272									const GetSecondaryKey& getSecondary);
273								~TwoKeyAVLTree();
274
275	inline	int					CountItems() const	{ return fTreeMap.Count(); }
276
277			Node*				Previous(Node* node) const;
278			Node*				Next(Node* node) const;
279
280			Value*				FindFirst(const PrimaryKey& key,
281									Iterator* iterator = NULL);
282			Value*				FindFirstClosest(const PrimaryKey& key,
283									bool less, Iterator* iterator = NULL);
284			Value*				FindLast(const PrimaryKey& key,
285									Iterator* iterator = NULL);
286	inline	Value*				Find(const PrimaryKey& primaryKey,
287									const SecondaryKey& secondaryKey,
288									Iterator* iterator = NULL);
289
290	inline	void				GetIterator(Iterator* iterator);
291	inline	void				GetIterator(Node* node, Iterator* iterator);
292
293	inline	status_t			Insert(const Value& value,
294									Iterator* iterator);
295	inline	status_t			Insert(const Value& value,
296									Node** _node = NULL);
297	inline	status_t			Remove(const PrimaryKey& primaryKey,
298									const SecondaryKey& secondaryKey);
299	inline	status_t			Remove(Node* node);
300
301private:
302			Node*				_FindFirst(const PrimaryKey& key,
303									Node** _parent) const;
304
305private:
306			TreeMap				fTreeMap;
307			PrimaryKeyCompare	fPrimaryKeyCompare;
308			GetPrimaryKey		fGetPrimaryKey;
309};
310
311
312// #pragma mark - Iterator
313
314
315TWO_KEY_AVL_TREE_TEMPLATE_LIST
316class TWO_KEY_AVL_TREE_CLASS_NAME::Iterator {
317public:
318	typedef typename TWO_KEY_AVL_TREE_CLASS_NAME::TreeMapIterator
319		TreeMapIterator;
320
321	inline Iterator()
322		:
323		fIterator()
324	{
325	}
326
327	inline ~Iterator()
328	{
329	}
330
331	inline Value* Current()
332	{
333		return fIterator.CurrentValuePointer();
334	}
335
336	inline Node* CurrentNode()
337	{
338		return fIterator.CurrentNode();
339	}
340
341	inline Value* Next()
342	{
343		fIterator.Next();
344		return Current();
345	}
346
347	inline Value* Previous()
348	{
349		fIterator.Previous();
350		return Current();
351	}
352
353	inline void Remove()
354	{
355		fIterator.Remove();
356	}
357
358private:
359	friend class TWO_KEY_AVL_TREE_CLASS_NAME;
360
361	Iterator(const Iterator& other)
362		:
363		fIterator(other.fIterator)
364	{
365	}
366
367	Iterator& operator=(const Iterator& other)
368	{
369		fIterator = other.fIterator;
370	}
371
372	Iterator(const TreeMapIterator& iterator)
373		:
374		fIterator()
375	{
376	}
377
378	inline void _SetTo(const TreeMapIterator& iterator)
379	{
380		fIterator = iterator;
381	}
382
383private:
384	TWO_KEY_AVL_TREE_CLASS_NAME::TreeMapIterator fIterator;
385};
386
387
388TWO_KEY_AVL_TREE_TEMPLATE_LIST
389TWO_KEY_AVL_TREE_CLASS_NAME::TwoKeyAVLTree()
390	:
391	fTreeMap(NodeStrategy(PrimaryKeyCompare(), SecondaryKeyCompare(),
392		GetPrimaryKey(), GetSecondaryKey())),
393	fPrimaryKeyCompare(PrimaryKeyCompare()),
394	fGetPrimaryKey(GetPrimaryKey())
395{
396}
397
398
399TWO_KEY_AVL_TREE_TEMPLATE_LIST
400TWO_KEY_AVL_TREE_CLASS_NAME::TwoKeyAVLTree(
401	const PrimaryKeyCompare& primaryCompare, const GetPrimaryKey& getPrimary,
402	const SecondaryKeyCompare& secondaryCompare,
403	const GetSecondaryKey& getSecondary)
404	:
405	fTreeMap(NodeStrategy(primaryCompare, secondaryCompare, getPrimary,
406		getSecondary)),
407	fPrimaryKeyCompare(primaryCompare),
408	fGetPrimaryKey(getPrimary)
409{
410}
411
412
413TWO_KEY_AVL_TREE_TEMPLATE_LIST
414TWO_KEY_AVL_TREE_CLASS_NAME::~TwoKeyAVLTree()
415{
416}
417
418
419TWO_KEY_AVL_TREE_TEMPLATE_LIST
420typename TWO_KEY_AVL_TREE_CLASS_NAME::Node*
421TWO_KEY_AVL_TREE_CLASS_NAME::Previous(Node* node) const
422{
423	return fTreeMap.Previous(node);
424}
425
426
427TWO_KEY_AVL_TREE_TEMPLATE_LIST
428typename TWO_KEY_AVL_TREE_CLASS_NAME::Node*
429TWO_KEY_AVL_TREE_CLASS_NAME::Next(Node* node) const
430{
431	return fTreeMap.Next(node);
432}
433
434
435TWO_KEY_AVL_TREE_TEMPLATE_LIST
436Value*
437TWO_KEY_AVL_TREE_CLASS_NAME::FindFirst(const PrimaryKey& key,
438	Iterator* iterator)
439{
440	const NodeStrategy& strategy = fTreeMap.GetNodeStrategy();
441
442	Node* node = _FindFirst(key, NULL);
443	if (node == NULL)
444		return NULL;
445
446	if (iterator != NULL)
447		iterator->_SetTo(fTreeMap.GetIterator(node));
448
449	return &strategy.GetValue(node);
450}
451
452
453TWO_KEY_AVL_TREE_TEMPLATE_LIST
454Value*
455TWO_KEY_AVL_TREE_CLASS_NAME::FindFirstClosest(const PrimaryKey& key, bool less,
456	Iterator* iterator)
457{
458	const NodeStrategy& strategy = fTreeMap.GetNodeStrategy();
459
460	Node* parent;
461	Node* node = _FindFirst(key, &parent);
462	if (node == NULL) {
463		// not found -- try to get the closest node
464		if (parent == NULL)
465			return NULL;
466
467		node = parent;
468		int expectedCmp = less ? 1 : -1;
469		int cmp = fPrimaryKeyCompare(key,
470			fGetPrimaryKey(strategy.GetValue(strategy.GetNode(node))));
471
472		if (cmp != expectedCmp) {
473			// The node's value is less although we were asked for a greater
474			// value, or the other way around. We need to iterate to the next
475			// node in the respective direction. If there is no node, we fail.
476			node = less ? Previous(node) : Next(node);
477			if (node == NULL)
478				return NULL;
479		}
480	}
481
482	if (iterator != NULL)
483		iterator->_SetTo(fTreeMap.GetIterator(node));
484
485	return &strategy.GetValue(node);
486}
487
488
489TWO_KEY_AVL_TREE_TEMPLATE_LIST
490Value*
491TWO_KEY_AVL_TREE_CLASS_NAME::FindLast(const PrimaryKey& key,
492	Iterator* iterator)
493{
494	const NodeStrategy& strategy = fTreeMap.GetNodeStrategy();
495	Node* node = fTreeMap.RootNode();
496
497	while (node) {
498		int cmp = fPrimaryKeyCompare(key, fGetPrimaryKey(
499			strategy.GetValue(node)));
500		if (cmp == 0) {
501			// found a matching node, now get the right-most node with that key
502			while (node->right && fPrimaryKeyCompare(key,
503				   	fGetPrimaryKey(strategy.GetValue(
504						strategy.GetNode(node->right)))) == 0) {
505				node = strategy.GetNode(node->right);
506			}
507			if (iterator)
508				iterator->_SetTo(fTreeMap.GetIterator(node));
509			return &strategy.GetValue(node);
510		}
511
512		if (cmp < 0)
513			node = strategy.GetNode(node->left);
514		else
515			node = strategy.GetNode(node->right);
516	}
517	return NULL;
518}
519
520
521TWO_KEY_AVL_TREE_TEMPLATE_LIST
522Value*
523TWO_KEY_AVL_TREE_CLASS_NAME::Find(const PrimaryKey& primaryKey,
524	const SecondaryKey& secondaryKey, Iterator* iterator)
525{
526	TreeMapIterator it = fTreeMap.Find(Key(primaryKey, secondaryKey));
527	if (iterator)
528		iterator->_SetTo(it);
529	return it.CurrentValuePointer();
530}
531
532
533TWO_KEY_AVL_TREE_TEMPLATE_LIST
534void
535TWO_KEY_AVL_TREE_CLASS_NAME::GetIterator(Iterator* iterator)
536{
537	TreeMapIterator it = fTreeMap.GetIterator();
538	it.Next();
539		// Our iterator needs to point to the first entry already.
540	iterator->_SetTo(it);
541}
542
543
544TWO_KEY_AVL_TREE_TEMPLATE_LIST
545void
546TWO_KEY_AVL_TREE_CLASS_NAME::GetIterator(Node* node, Iterator* iterator)
547{
548	iterator->_SetTo(fTreeMap.GetIterator(node));
549}
550
551
552TWO_KEY_AVL_TREE_TEMPLATE_LIST
553status_t
554TWO_KEY_AVL_TREE_CLASS_NAME::Insert(const Value& value, Iterator* iterator)
555{
556	NodeStrategy& strategy
557		= const_cast<NodeStrategy&>(fTreeMap.GetNodeStrategy());
558
559	TreeMapIterator it;
560	status_t status = fTreeMap.Insert(strategy.GetValueKey(value), value, &it);
561	if (status != B_OK || !iterator)
562		return status;
563
564	iterator->_SetTo(it);
565	return B_OK;
566}
567
568
569TWO_KEY_AVL_TREE_TEMPLATE_LIST
570status_t
571TWO_KEY_AVL_TREE_CLASS_NAME::Insert(const Value& value, Node** _node)
572{
573	NodeStrategy& strategy
574		= const_cast<NodeStrategy&>(fTreeMap.GetNodeStrategy());
575
576	return fTreeMap.Insert(strategy.GetValueKey(value), value, _node);
577}
578
579
580TWO_KEY_AVL_TREE_TEMPLATE_LIST
581status_t
582TWO_KEY_AVL_TREE_CLASS_NAME::Remove(const PrimaryKey& primaryKey,
583	const SecondaryKey& secondaryKey)
584{
585	return fTreeMap.Remove(Key(primaryKey, secondaryKey));
586}
587
588
589TWO_KEY_AVL_TREE_TEMPLATE_LIST
590status_t
591TWO_KEY_AVL_TREE_CLASS_NAME::Remove(Node* node)
592{
593	return fTreeMap.Remove(node);
594}
595
596
597TWO_KEY_AVL_TREE_TEMPLATE_LIST
598typename TWO_KEY_AVL_TREE_CLASS_NAME::Node*
599TWO_KEY_AVL_TREE_CLASS_NAME::_FindFirst(const PrimaryKey& key,
600	Node** _parent) const
601{
602	const NodeStrategy& strategy = fTreeMap.GetNodeStrategy();
603	Node* node = fTreeMap.RootNode();
604	Node* parent = NULL;
605
606	while (node) {
607		int cmp = fPrimaryKeyCompare(key, fGetPrimaryKey(
608			strategy.GetValue(node)));
609		if (cmp == 0) {
610			// found a matching node, now get the left-most node with that key
611			while (node->left && fPrimaryKeyCompare(key,
612				   	fGetPrimaryKey(strategy.GetValue(
613						strategy.GetNode(node->left)))) == 0) {
614				node = strategy.GetNode(node->left);
615			}
616
617			return node;
618		}
619
620		parent = node;
621
622		if (cmp < 0)
623			node = strategy.GetNode(node->left);
624		else
625			node = strategy.GetNode(node->right);
626	}
627
628	if (_parent != NULL)
629		*_parent = parent;
630
631	return NULL;
632}
633
634
635#endif	// TWO_KEY_AVL_TREE_H
636