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