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		// There is no way to overload operator delete with extra parameters.
211		node->~Node();
212		object_cache_free(fObjectCache, node, 0);
213	}
214
215	// internal use (not part of the strategy)
216	inline Key GetValueKey(const Value& value) const
217	{
218		return Key(fGetPrimaryKey(value), fGetSecondaryKey(value));
219	}
220
221	inline Key GetKey(Node* node) const
222	{
223		return GetValueKey(node->value);
224	}
225
226	inline Value& GetValue(Node* node) const
227	{
228		return node->value;
229	}
230
231	inline AVLTreeNode* GetAVLTreeNode(Node* node) const
232	{
233		return node;
234	}
235
236	inline Node* GetNode(AVLTreeNode* node) const
237	{
238		return static_cast<Node*>(node);
239	}
240
241	inline int CompareKeyNode(const Key& a, const Node* b) const
242	{
243		return _CompareKeys(a, GetKey(const_cast<Node*>(b)));
244	}
245
246	inline int CompareNodes(const Node* a, const Node* b) const
247	{
248		return _CompareKeys(GetKey(const_cast<Node*>(a)),
249			GetKey(const_cast<Node*>(b)));
250	}
251
252private:
253	inline int _CompareKeys(const Key& a, const Key& b) const
254	{
255		int result = fPrimaryKeyCompare(a.primary, b.primary);
256		if (result == 0 && a.use_secondary && b.use_secondary)
257			result = fSecondaryKeyCompare(a.secondary, b.secondary);
258		return result;
259	}
260
261	PrimaryKeyCompare	fPrimaryKeyCompare;
262	SecondaryKeyCompare	fSecondaryKeyCompare;
263	GetPrimaryKey		fGetPrimaryKey;
264	GetSecondaryKey		fGetSecondaryKey;
265
266	object_cache*		fObjectCache;
267	int32*				fObjectCacheRefs;
268};
269
270
271// for convenience
272#define TWO_KEY_AVL_TREE_TEMPLATE_LIST template<typename Value, \
273	typename PrimaryKey, typename PrimaryKeyCompare, \
274	typename GetPrimaryKey, typename SecondaryKey, \
275	typename SecondaryKeyCompare, typename GetSecondaryKey>
276#define TWO_KEY_AVL_TREE_CLASS_NAME TwoKeyAVLTree<Value, PrimaryKey, \
277	PrimaryKeyCompare, GetPrimaryKey, SecondaryKey, \
278	SecondaryKeyCompare, GetSecondaryKey>
279
280
281// #pragma mark - TwoKeyAVLTree
282
283
284template<typename Value, typename PrimaryKey,
285	typename PrimaryKeyCompare, typename GetPrimaryKey,
286	typename SecondaryKey = Value,
287	typename SecondaryKeyCompare = TwoKeyAVLTreeStandardCompare<SecondaryKey>,
288	typename GetSecondaryKey
289		= TwoKeyAVLTreeStandardGetKey<Value, SecondaryKey> >
290class TwoKeyAVLTree {
291public:
292			typedef TwoKeyAVLTreeKey<PrimaryKey, SecondaryKey> Key;
293			typedef TwoKeyAVLTreeNodeStrategy<PrimaryKey, SecondaryKey, Value,
294				PrimaryKeyCompare, SecondaryKeyCompare, GetPrimaryKey,
295				GetSecondaryKey> NodeStrategy;
296			typedef AVLTreeMap<Key, Value, NodeStrategy> TreeMap;
297
298			typedef typename TreeMap::Iterator	TreeMapIterator;
299			typedef typename NodeStrategy::Node Node;
300
301			class Iterator;
302
303public:
304								TwoKeyAVLTree();
305								TwoKeyAVLTree(
306									const PrimaryKeyCompare& primaryCompare,
307									const GetPrimaryKey& getPrimary,
308									const SecondaryKeyCompare& secondaryCompare,
309									const GetSecondaryKey& getSecondary);
310								~TwoKeyAVLTree();
311
312	inline	int					CountItems() const	{ return fTreeMap.Count(); }
313
314			Node*				Previous(Node* node) const;
315			Node*				Next(Node* node) const;
316
317			Value*				FindFirst(const PrimaryKey& key,
318									Iterator* iterator = NULL);
319			Value*				FindFirstClosest(const PrimaryKey& key,
320									bool less, Iterator* iterator = NULL);
321			Value*				FindLast(const PrimaryKey& key,
322									Iterator* iterator = NULL);
323	inline	Value*				Find(const PrimaryKey& primaryKey,
324									const SecondaryKey& secondaryKey,
325									Iterator* iterator = NULL);
326
327	inline	void				GetIterator(Iterator* iterator);
328	inline	void				GetIterator(Node* node, Iterator* iterator);
329
330	inline	status_t			Insert(const Value& value,
331									Iterator* iterator);
332	inline	status_t			Insert(const Value& value,
333									Node** _node = NULL);
334	inline	status_t			Remove(const PrimaryKey& primaryKey,
335									const SecondaryKey& secondaryKey);
336	inline	status_t			Remove(Node* node);
337
338private:
339			Node*				_FindFirst(const PrimaryKey& key,
340									Node** _parent) const;
341
342private:
343			TreeMap				fTreeMap;
344			PrimaryKeyCompare	fPrimaryKeyCompare;
345			GetPrimaryKey		fGetPrimaryKey;
346};
347
348
349// #pragma mark - Iterator
350
351
352TWO_KEY_AVL_TREE_TEMPLATE_LIST
353class TWO_KEY_AVL_TREE_CLASS_NAME::Iterator {
354public:
355	typedef typename TWO_KEY_AVL_TREE_CLASS_NAME::TreeMapIterator
356		TreeMapIterator;
357
358	inline Iterator()
359		:
360		fIterator()
361	{
362	}
363
364	inline ~Iterator()
365	{
366	}
367
368	inline Value* Current()
369	{
370		return fIterator.CurrentValuePointer();
371	}
372
373	inline Node* CurrentNode()
374	{
375		return fIterator.CurrentNode();
376	}
377
378	inline Value* Next()
379	{
380		fIterator.Next();
381		return Current();
382	}
383
384	inline Value* Previous()
385	{
386		fIterator.Previous();
387		return Current();
388	}
389
390	inline void Remove()
391	{
392		fIterator.Remove();
393	}
394
395private:
396	friend class TWO_KEY_AVL_TREE_CLASS_NAME;
397
398	Iterator(const Iterator& other)
399		:
400		fIterator(other.fIterator)
401	{
402	}
403
404	Iterator& operator=(const Iterator& other)
405	{
406		fIterator = other.fIterator;
407	}
408
409	Iterator(const TreeMapIterator& iterator)
410		:
411		fIterator()
412	{
413	}
414
415	inline void _SetTo(const TreeMapIterator& iterator)
416	{
417		fIterator = iterator;
418	}
419
420private:
421	TWO_KEY_AVL_TREE_CLASS_NAME::TreeMapIterator fIterator;
422};
423
424
425TWO_KEY_AVL_TREE_TEMPLATE_LIST
426TWO_KEY_AVL_TREE_CLASS_NAME::TwoKeyAVLTree()
427	:
428	fTreeMap(NodeStrategy(PrimaryKeyCompare(), SecondaryKeyCompare(),
429		GetPrimaryKey(), GetSecondaryKey())),
430	fPrimaryKeyCompare(PrimaryKeyCompare()),
431	fGetPrimaryKey(GetPrimaryKey())
432{
433}
434
435
436TWO_KEY_AVL_TREE_TEMPLATE_LIST
437TWO_KEY_AVL_TREE_CLASS_NAME::TwoKeyAVLTree(
438	const PrimaryKeyCompare& primaryCompare, const GetPrimaryKey& getPrimary,
439	const SecondaryKeyCompare& secondaryCompare,
440	const GetSecondaryKey& getSecondary)
441	:
442	fTreeMap(NodeStrategy(primaryCompare, secondaryCompare, getPrimary,
443		getSecondary)),
444	fPrimaryKeyCompare(primaryCompare),
445	fGetPrimaryKey(getPrimary)
446{
447}
448
449
450TWO_KEY_AVL_TREE_TEMPLATE_LIST
451TWO_KEY_AVL_TREE_CLASS_NAME::~TwoKeyAVLTree()
452{
453}
454
455
456TWO_KEY_AVL_TREE_TEMPLATE_LIST
457typename TWO_KEY_AVL_TREE_CLASS_NAME::Node*
458TWO_KEY_AVL_TREE_CLASS_NAME::Previous(Node* node) const
459{
460	return fTreeMap.Previous(node);
461}
462
463
464TWO_KEY_AVL_TREE_TEMPLATE_LIST
465typename TWO_KEY_AVL_TREE_CLASS_NAME::Node*
466TWO_KEY_AVL_TREE_CLASS_NAME::Next(Node* node) const
467{
468	return fTreeMap.Next(node);
469}
470
471
472TWO_KEY_AVL_TREE_TEMPLATE_LIST
473Value*
474TWO_KEY_AVL_TREE_CLASS_NAME::FindFirst(const PrimaryKey& key,
475	Iterator* iterator)
476{
477	const NodeStrategy& strategy = fTreeMap.GetNodeStrategy();
478
479	Node* node = _FindFirst(key, NULL);
480	if (node == NULL)
481		return NULL;
482
483	if (iterator != NULL)
484		iterator->_SetTo(fTreeMap.GetIterator(node));
485
486	return &strategy.GetValue(node);
487}
488
489
490TWO_KEY_AVL_TREE_TEMPLATE_LIST
491Value*
492TWO_KEY_AVL_TREE_CLASS_NAME::FindFirstClosest(const PrimaryKey& key, bool less,
493	Iterator* iterator)
494{
495	const NodeStrategy& strategy = fTreeMap.GetNodeStrategy();
496
497	Node* parent = NULL;
498	Node* node = _FindFirst(key, &parent);
499	if (node == NULL) {
500		// not found -- try to get the closest node
501		if (parent == NULL)
502			return NULL;
503
504		node = parent;
505		int expectedCmp = less ? 1 : -1;
506		int cmp = fPrimaryKeyCompare(key,
507			fGetPrimaryKey(strategy.GetValue(strategy.GetNode(node))));
508
509		if (cmp != expectedCmp) {
510			// The node's value is less although we were asked for a greater
511			// value, or the other way around. We need to iterate to the next
512			// node in the respective direction. If there is no node, we fail.
513			node = less ? Previous(node) : Next(node);
514			if (node == NULL)
515				return NULL;
516		}
517	}
518
519	if (iterator != NULL)
520		iterator->_SetTo(fTreeMap.GetIterator(node));
521
522	return &strategy.GetValue(node);
523}
524
525
526TWO_KEY_AVL_TREE_TEMPLATE_LIST
527Value*
528TWO_KEY_AVL_TREE_CLASS_NAME::FindLast(const PrimaryKey& key,
529	Iterator* iterator)
530{
531	const NodeStrategy& strategy = fTreeMap.GetNodeStrategy();
532	Node* node = fTreeMap.RootNode();
533
534	while (node) {
535		int cmp = fPrimaryKeyCompare(key, fGetPrimaryKey(
536			strategy.GetValue(node)));
537		if (cmp == 0) {
538			// found a matching node, now get the right-most node with that key
539			while (node->right && fPrimaryKeyCompare(key,
540				   	fGetPrimaryKey(strategy.GetValue(
541						strategy.GetNode(node->right)))) == 0) {
542				node = strategy.GetNode(node->right);
543			}
544			if (iterator)
545				iterator->_SetTo(fTreeMap.GetIterator(node));
546			return &strategy.GetValue(node);
547		}
548
549		if (cmp < 0)
550			node = strategy.GetNode(node->left);
551		else
552			node = strategy.GetNode(node->right);
553	}
554	return NULL;
555}
556
557
558TWO_KEY_AVL_TREE_TEMPLATE_LIST
559Value*
560TWO_KEY_AVL_TREE_CLASS_NAME::Find(const PrimaryKey& primaryKey,
561	const SecondaryKey& secondaryKey, Iterator* iterator)
562{
563	TreeMapIterator it = fTreeMap.Find(Key(primaryKey, secondaryKey));
564	if (iterator)
565		iterator->_SetTo(it);
566	return it.CurrentValuePointer();
567}
568
569
570TWO_KEY_AVL_TREE_TEMPLATE_LIST
571void
572TWO_KEY_AVL_TREE_CLASS_NAME::GetIterator(Iterator* iterator)
573{
574	TreeMapIterator it = fTreeMap.GetIterator();
575	it.Next();
576		// Our iterator needs to point to the first entry already.
577	iterator->_SetTo(it);
578}
579
580
581TWO_KEY_AVL_TREE_TEMPLATE_LIST
582void
583TWO_KEY_AVL_TREE_CLASS_NAME::GetIterator(Node* node, Iterator* iterator)
584{
585	iterator->_SetTo(fTreeMap.GetIterator(node));
586}
587
588
589TWO_KEY_AVL_TREE_TEMPLATE_LIST
590status_t
591TWO_KEY_AVL_TREE_CLASS_NAME::Insert(const Value& value, Iterator* iterator)
592{
593	NodeStrategy& strategy
594		= const_cast<NodeStrategy&>(fTreeMap.GetNodeStrategy());
595
596	TreeMapIterator it;
597	status_t status = fTreeMap.Insert(strategy.GetValueKey(value), value, &it);
598	if (status != B_OK || !iterator)
599		return status;
600
601	iterator->_SetTo(it);
602	return B_OK;
603}
604
605
606TWO_KEY_AVL_TREE_TEMPLATE_LIST
607status_t
608TWO_KEY_AVL_TREE_CLASS_NAME::Insert(const Value& value, Node** _node)
609{
610	NodeStrategy& strategy
611		= const_cast<NodeStrategy&>(fTreeMap.GetNodeStrategy());
612
613	return fTreeMap.Insert(strategy.GetValueKey(value), value, _node);
614}
615
616
617TWO_KEY_AVL_TREE_TEMPLATE_LIST
618status_t
619TWO_KEY_AVL_TREE_CLASS_NAME::Remove(const PrimaryKey& primaryKey,
620	const SecondaryKey& secondaryKey)
621{
622	return fTreeMap.Remove(Key(primaryKey, secondaryKey));
623}
624
625
626TWO_KEY_AVL_TREE_TEMPLATE_LIST
627status_t
628TWO_KEY_AVL_TREE_CLASS_NAME::Remove(Node* node)
629{
630	return fTreeMap.Remove(node);
631}
632
633
634TWO_KEY_AVL_TREE_TEMPLATE_LIST
635typename TWO_KEY_AVL_TREE_CLASS_NAME::Node*
636TWO_KEY_AVL_TREE_CLASS_NAME::_FindFirst(const PrimaryKey& key,
637	Node** _parent) const
638{
639	const NodeStrategy& strategy = fTreeMap.GetNodeStrategy();
640	Node* node = fTreeMap.RootNode();
641	Node* parent = NULL;
642
643	while (node) {
644		int cmp = fPrimaryKeyCompare(key, fGetPrimaryKey(
645			strategy.GetValue(node)));
646		if (cmp == 0) {
647			// found a matching node, now get the left-most node with that key
648			while (node->left && fPrimaryKeyCompare(key,
649				   	fGetPrimaryKey(strategy.GetValue(
650						strategy.GetNode(node->left)))) == 0) {
651				node = strategy.GetNode(node->left);
652			}
653
654			return node;
655		}
656
657		parent = node;
658
659		if (cmp < 0)
660			node = strategy.GetNode(node->left);
661		else
662			node = strategy.GetNode(node->right);
663	}
664
665	if (_parent != NULL)
666		*_parent = parent;
667
668	return NULL;
669}
670
671
672#endif	// TWO_KEY_AVL_TREE_H
673