1/*
2 * Copyright 2003-2009, Ingo Weinhold <ingo_weinhold@gmx.de>.
3 * Distributed under the terms of the MIT License.
4 */
5#ifndef _KERNEL_UTIL_AVL_TREE_MAP_H
6#define _KERNEL_UTIL_AVL_TREE_MAP_H
7
8
9#include <util/MallocFreeAllocator.h>
10#include <util/AVLTreeBase.h>
11
12
13// strategies
14namespace AVLTreeMapStrategy {
15	// key orders
16	template<typename Value> class Ascending;
17	template<typename Value> class Descending;
18
19	//! Automatic node strategy (works like STL containers do)
20	template <typename Key, typename Value,
21			  typename KeyOrder = Ascending<Key>,
22			  template <typename> class Allocator = MallocFreeAllocator>
23	class Auto;
24
25/*!	NodeStrategy must implement this interface:
26	typename Node;
27	inline Node* Allocate(const Key& key, const Value& value)
28	inline void Free(Node* node)
29	inline const Key GetKey(const Node* node) const
30	inline Value& GetValue(Node* node) const
31	inline AVLTreeNode* GetAVLTreeNode(Node* node) const
32	inline Node* GetNode(AVLTreeNode* node) const
33	inline int CompareKeyNode(const Key& a, const Node* b)
34	inline int CompareNodes(const Node* a, const Node* b)
35*/
36}
37
38
39// for convenience
40#define _AVL_TREE_MAP_TEMPLATE_LIST template<typename Key, typename Value, \
41	typename NodeStrategy>
42#define _AVL_TREE_MAP_CLASS_NAME AVLTreeMap<Key, Value, NodeStrategy>
43
44
45// AVLTreeMap
46template<typename Key, typename Value,
47		 typename NodeStrategy = AVLTreeMapStrategy::Auto<Key, Value> >
48class AVLTreeMap : protected AVLTreeCompare {
49private:
50	typedef typename NodeStrategy::Node	Node;
51	typedef _AVL_TREE_MAP_CLASS_NAME	Class;
52
53public:
54	class Iterator;
55	class ConstIterator;
56
57public:
58								AVLTreeMap(const NodeStrategy& strategy
59									= NodeStrategy());
60	virtual						~AVLTreeMap();
61
62	inline	int					Count() const	{ return fTree.Count(); }
63	inline	bool				IsEmpty() const	{ return fTree.IsEmpty(); }
64	inline	void				MakeEmpty();
65
66			Node*				RootNode() const;
67
68			Node*				Previous(Node* node) const;
69			Node*				Next(Node* node) const;
70
71	inline	Iterator			GetIterator();
72	inline	ConstIterator		GetIterator() const;
73
74	inline	Iterator			GetIterator(Node* node);
75	inline	ConstIterator		GetIterator(Node* node) const;
76
77			Iterator			Find(const Key& key);
78			Iterator			FindClose(const Key& key, bool less);
79
80			status_t			Insert(const Key& key, const Value& value,
81									Iterator* iterator);
82			status_t			Insert(const Key& key, const Value& value,
83									Node** _node = NULL);
84			status_t			Remove(const Key& key);
85			status_t			Remove(Node* node);
86
87			const NodeStrategy&	GetNodeStrategy() const	{ return fStrategy; }
88
89protected:
90	// AVLTreeCompare
91	virtual	int					CompareKeyNode(const void* key,
92									const AVLTreeNode* node);
93	virtual	int					CompareNodes(const AVLTreeNode* node1,
94									const AVLTreeNode* node2);
95
96			void				_FreeTree(AVLTreeNode* node);
97
98	// strategy shortcuts
99	inline	Node*				_Allocate(const Key& key, const Value& value);
100	inline	void				_Free(Node* node);
101	inline	Key					_GetKey(Node* node) const;
102	inline	Value&				_GetValue(Node* node) const;
103	inline	AVLTreeNode*		_GetAVLTreeNode(const Node* node) const;
104	inline	Node*				_GetNode(const AVLTreeNode* node) const;
105	inline	int					_CompareKeyNode(const Key& a, const Node* b);
106	inline	int					_CompareNodes(const Node* a, const Node* b);
107
108protected:
109			friend class Iterator;
110			friend class ConstIterator;
111
112			AVLTreeBase			fTree;
113			NodeStrategy		fStrategy;
114
115public:
116	// Iterator
117	// (need to implement it here, otherwise gcc 2.95.3 chokes)
118	class Iterator : public ConstIterator {
119	public:
120		inline Iterator()
121			: ConstIterator()
122		{
123		}
124
125		inline Iterator(const Iterator& other)
126			: ConstIterator(other)
127		{
128		}
129
130		inline void Remove()
131		{
132			if (AVLTreeNode* node = ConstIterator::fTreeIterator.Remove()) {
133				_AVL_TREE_MAP_CLASS_NAME* parent
134					= const_cast<_AVL_TREE_MAP_CLASS_NAME*>(
135						ConstIterator::fParent);
136				parent->_Free(parent->_GetNode(node));
137			}
138		}
139
140	private:
141		inline Iterator(_AVL_TREE_MAP_CLASS_NAME* parent,
142			const AVLTreeIterator& treeIterator)
143			: ConstIterator(parent, treeIterator)
144		{
145		}
146
147		friend class _AVL_TREE_MAP_CLASS_NAME;
148	};
149};
150
151
152// ConstIterator
153_AVL_TREE_MAP_TEMPLATE_LIST
154class _AVL_TREE_MAP_CLASS_NAME::ConstIterator {
155public:
156	inline ConstIterator()
157		: fParent(NULL),
158		  fTreeIterator()
159	{
160	}
161
162	inline ConstIterator(const ConstIterator& other)
163		: fParent(other.fParent),
164		  fTreeIterator(other.fTreeIterator)
165	{
166	}
167
168	inline bool HasCurrent() const
169	{
170		return fTreeIterator.Current();
171	}
172
173	inline Key CurrentKey()
174	{
175		if (AVLTreeNode* node = fTreeIterator.Current())
176			return fParent->_GetKey(fParent->_GetNode(node));
177		return Key();
178	}
179
180	inline Value Current()
181	{
182		if (AVLTreeNode* node = fTreeIterator.Current())
183			return fParent->_GetValue(fParent->_GetNode(node));
184		return Value();
185	}
186
187	inline Value* CurrentValuePointer()
188	{
189		if (AVLTreeNode* node = fTreeIterator.Current())
190			return &fParent->_GetValue(fParent->_GetNode(node));
191		return NULL;
192	}
193
194	inline Node* CurrentNode()
195	{
196		if (AVLTreeNode* node = fTreeIterator.Current())
197			return fParent->_GetNode(node);
198		return NULL;
199	}
200
201	inline bool HasNext() const
202	{
203		return fTreeIterator.HasNext();
204	}
205
206	inline Value Next()
207	{
208		if (AVLTreeNode* node = fTreeIterator.Next())
209			return fParent->_GetValue(fParent->_GetNode(node));
210		return Value();
211	}
212
213	inline Value* NextValuePointer()
214	{
215		if (AVLTreeNode* node = fTreeIterator.Next())
216			return &fParent->_GetValue(fParent->_GetNode(node));
217		return NULL;
218	}
219
220	inline Value Previous()
221	{
222		if (AVLTreeNode* node = fTreeIterator.Previous())
223			return fParent->_GetValue(fParent->_GetNode(node));
224		return Value();
225	}
226
227	inline ConstIterator& operator=(const ConstIterator& other)
228	{
229		fParent = other.fParent;
230		fTreeIterator = other.fTreeIterator;
231		return *this;
232	}
233
234protected:
235	inline ConstIterator(const _AVL_TREE_MAP_CLASS_NAME* parent,
236		const AVLTreeIterator& treeIterator)
237	{
238		fParent = parent;
239		fTreeIterator = treeIterator;
240	}
241
242	friend class _AVL_TREE_MAP_CLASS_NAME;
243
244	const _AVL_TREE_MAP_CLASS_NAME*	fParent;
245	AVLTreeIterator					fTreeIterator;
246};
247
248
249// constructor
250_AVL_TREE_MAP_TEMPLATE_LIST
251_AVL_TREE_MAP_CLASS_NAME::AVLTreeMap(const NodeStrategy& strategy)
252	: fTree(this),
253	  fStrategy(strategy)
254{
255}
256
257
258// destructor
259_AVL_TREE_MAP_TEMPLATE_LIST
260_AVL_TREE_MAP_CLASS_NAME::~AVLTreeMap()
261{
262}
263
264
265// MakeEmpty
266_AVL_TREE_MAP_TEMPLATE_LIST
267inline void
268_AVL_TREE_MAP_CLASS_NAME::MakeEmpty()
269{
270	AVLTreeNode* root = fTree.Root();
271	_FreeTree(root);
272	fTree.MakeEmpty();
273}
274
275
276// RootNode
277_AVL_TREE_MAP_TEMPLATE_LIST
278inline typename _AVL_TREE_MAP_CLASS_NAME::Node*
279_AVL_TREE_MAP_CLASS_NAME::RootNode() const
280{
281	if (AVLTreeNode* root = fTree.Root())
282		return _GetNode(root);
283	return NULL;
284}
285
286
287// Previous
288_AVL_TREE_MAP_TEMPLATE_LIST
289inline typename _AVL_TREE_MAP_CLASS_NAME::Node*
290_AVL_TREE_MAP_CLASS_NAME::Previous(Node* node) const
291{
292	if (node == NULL)
293		return NULL;
294
295	AVLTreeNode* treeNode = fTree.Previous(_GetAVLTreeNode(node));
296	return treeNode != NULL ? _GetNode(treeNode) : NULL;
297}
298
299
300// Next
301_AVL_TREE_MAP_TEMPLATE_LIST
302inline typename _AVL_TREE_MAP_CLASS_NAME::Node*
303_AVL_TREE_MAP_CLASS_NAME::Next(Node* node) const
304{
305	if (node == NULL)
306		return NULL;
307
308	AVLTreeNode* treeNode = fTree.Next(_GetAVLTreeNode(node));
309	return treeNode != NULL ? _GetNode(treeNode) : NULL;
310}
311
312
313// GetIterator
314_AVL_TREE_MAP_TEMPLATE_LIST
315inline typename _AVL_TREE_MAP_CLASS_NAME::Iterator
316_AVL_TREE_MAP_CLASS_NAME::GetIterator()
317{
318	return Iterator(this, fTree.GetIterator());
319}
320
321
322// GetIterator
323_AVL_TREE_MAP_TEMPLATE_LIST
324inline typename _AVL_TREE_MAP_CLASS_NAME::ConstIterator
325_AVL_TREE_MAP_CLASS_NAME::GetIterator() const
326{
327	return ConstIterator(this, fTree.GetIterator());
328}
329
330
331// GetIterator
332_AVL_TREE_MAP_TEMPLATE_LIST
333inline typename _AVL_TREE_MAP_CLASS_NAME::Iterator
334_AVL_TREE_MAP_CLASS_NAME::GetIterator(Node* node)
335{
336	return Iterator(this, fTree.GetIterator(_GetAVLTreeNode(node)));
337}
338
339
340// GetIterator
341_AVL_TREE_MAP_TEMPLATE_LIST
342inline typename _AVL_TREE_MAP_CLASS_NAME::ConstIterator
343_AVL_TREE_MAP_CLASS_NAME::GetIterator(Node* node) const
344{
345	return ConstIterator(this, fTree.GetIterator(_GetAVLTreeNode(node)));
346}
347
348
349// Find
350_AVL_TREE_MAP_TEMPLATE_LIST
351typename _AVL_TREE_MAP_CLASS_NAME::Iterator
352_AVL_TREE_MAP_CLASS_NAME::Find(const Key& key)
353{
354	if (AVLTreeNode* node = fTree.Find(&key))
355		return Iterator(this, fTree.GetIterator(node));
356	return Iterator();
357}
358
359
360// FindClose
361_AVL_TREE_MAP_TEMPLATE_LIST
362typename _AVL_TREE_MAP_CLASS_NAME::Iterator
363_AVL_TREE_MAP_CLASS_NAME::FindClose(const Key& key, bool less)
364{
365	if (AVLTreeNode* node = fTree.FindClosest(&key, less))
366		return Iterator(this, fTree.GetIterator(node));
367	return Iterator();
368}
369
370
371// Insert
372_AVL_TREE_MAP_TEMPLATE_LIST
373status_t
374_AVL_TREE_MAP_CLASS_NAME::Insert(const Key& key, const Value& value,
375	Iterator* iterator)
376{
377	// allocate a node
378	Node* userNode = _Allocate(key, value);
379	if (!userNode)
380		return B_NO_MEMORY;
381
382	// insert node
383	AVLTreeNode* node = _GetAVLTreeNode(userNode);
384	status_t error = fTree.Insert(node);
385	if (error != B_OK) {
386		_Free(userNode);
387		return error;
388	}
389
390	if (iterator)
391		*iterator = Iterator(this, fTree.GetIterator(node));
392
393	return B_OK;
394}
395
396
397// Insert
398_AVL_TREE_MAP_TEMPLATE_LIST
399status_t
400_AVL_TREE_MAP_CLASS_NAME::Insert(const Key& key, const Value& value,
401	Node** _node)
402{
403	// allocate a node
404	Node* userNode = _Allocate(key, value);
405	if (!userNode)
406		return B_NO_MEMORY;
407
408	// insert node
409	AVLTreeNode* node = _GetAVLTreeNode(userNode);
410	status_t error = fTree.Insert(node);
411	if (error != B_OK) {
412		_Free(userNode);
413		return error;
414	}
415
416	if (_node != NULL)
417		*_node = userNode;
418
419	return B_OK;
420}
421
422
423// Remove
424_AVL_TREE_MAP_TEMPLATE_LIST
425status_t
426_AVL_TREE_MAP_CLASS_NAME::Remove(const Key& key)
427{
428	AVLTreeNode* node = fTree.Remove(&key);
429	if (!node)
430		return B_ENTRY_NOT_FOUND;
431
432	_Free(_GetNode(node));
433	return B_OK;
434}
435
436
437// Remove
438_AVL_TREE_MAP_TEMPLATE_LIST
439status_t
440_AVL_TREE_MAP_CLASS_NAME::Remove(Node* node)
441{
442	if (!fTree.Remove(node))
443		return B_ENTRY_NOT_FOUND;
444
445	_Free(node);
446	return B_OK;
447}
448
449
450// CompareKeyNode
451_AVL_TREE_MAP_TEMPLATE_LIST
452int
453_AVL_TREE_MAP_CLASS_NAME::CompareKeyNode(const void* key,
454	const AVLTreeNode* node)
455{
456	return _CompareKeyNode(*(const Key*)key, _GetNode(node));
457}
458
459
460// CompareNodes
461_AVL_TREE_MAP_TEMPLATE_LIST
462int
463_AVL_TREE_MAP_CLASS_NAME::CompareNodes(const AVLTreeNode* node1,
464	const AVLTreeNode* node2)
465{
466	return _CompareNodes(_GetNode(node1), _GetNode(node2));
467}
468
469
470// _Allocate
471_AVL_TREE_MAP_TEMPLATE_LIST
472inline typename _AVL_TREE_MAP_CLASS_NAME::Node*
473_AVL_TREE_MAP_CLASS_NAME::_Allocate(const Key& key, const Value& value)
474{
475	return fStrategy.Allocate(key, value);
476}
477
478
479// _Free
480_AVL_TREE_MAP_TEMPLATE_LIST
481inline void
482_AVL_TREE_MAP_CLASS_NAME::_Free(Node* node)
483{
484	fStrategy.Free(node);
485}
486
487
488// _GetKey
489_AVL_TREE_MAP_TEMPLATE_LIST
490inline Key
491_AVL_TREE_MAP_CLASS_NAME::_GetKey(Node* node) const
492{
493	return fStrategy.GetKey(node);
494}
495
496
497// _GetValue
498_AVL_TREE_MAP_TEMPLATE_LIST
499inline Value&
500_AVL_TREE_MAP_CLASS_NAME::_GetValue(Node* node) const
501{
502	return fStrategy.GetValue(node);
503}
504
505
506// _GetAVLTreeNode
507_AVL_TREE_MAP_TEMPLATE_LIST
508inline AVLTreeNode*
509_AVL_TREE_MAP_CLASS_NAME::_GetAVLTreeNode(const Node* node) const
510{
511	return fStrategy.GetAVLTreeNode(const_cast<Node*>(node));
512}
513
514
515// _GetNode
516_AVL_TREE_MAP_TEMPLATE_LIST
517inline typename _AVL_TREE_MAP_CLASS_NAME::Node*
518_AVL_TREE_MAP_CLASS_NAME::_GetNode(const AVLTreeNode* node) const
519{
520	return fStrategy.GetNode(const_cast<AVLTreeNode*>(node));
521}
522
523
524// _CompareKeyNode
525_AVL_TREE_MAP_TEMPLATE_LIST
526inline int
527_AVL_TREE_MAP_CLASS_NAME::_CompareKeyNode(const Key& a, const Node* b)
528{
529	return fStrategy.CompareKeyNode(a, b);
530}
531
532
533// _CompareNodes
534_AVL_TREE_MAP_TEMPLATE_LIST
535inline int
536_AVL_TREE_MAP_CLASS_NAME::_CompareNodes(const Node* a, const Node* b)
537{
538	return fStrategy.CompareNodes(a, b);
539}
540
541
542// _FreeTree
543_AVL_TREE_MAP_TEMPLATE_LIST
544void
545_AVL_TREE_MAP_CLASS_NAME::_FreeTree(AVLTreeNode* node)
546{
547	if (node) {
548		_FreeTree(node->left);
549		_FreeTree(node->right);
550		_Free(_GetNode(node));
551	}
552}
553
554
555// #pragma mark - strategy parameters
556
557// Ascending
558namespace AVLTreeMapStrategy {
559template<typename Value>
560class Ascending {
561public:
562	inline int operator()(const Value &a, const Value &b) const
563	{
564		if (a < b)
565			return -1;
566		else if (a > b)
567			return 1;
568		return 0;
569	}
570};
571}
572
573
574// Descending
575namespace AVLTreeMapStrategy {
576template<typename Value>
577class Descending {
578public:
579	inline int operator()(const Value &a, const Value &b) const
580	{
581		if (a < b)
582			return -1;
583		else if (a > b)
584			return 1;
585		return 0;
586	}
587};
588}
589
590
591// #pragma mark - strategies
592
593
594// Auto
595namespace AVLTreeMapStrategy {
596template <typename Key, typename Value, typename KeyOrder,
597		  template <typename> class Allocator>
598class Auto {
599public:
600	struct Node : AVLTreeNode {
601		Node(const Key &key, const Value &value)
602			: AVLTreeNode(),
603			  key(key),
604			  value(value)
605		{
606		}
607
608		Key		key;
609		Value	value;
610	};
611
612	inline Node* Allocate(const Key& key, const Value& value)
613	{
614		Node* result = fAllocator.Allocate();
615		if (result)
616			fAllocator.Construct(result, key, value);
617		return result;
618	}
619
620	inline void Free(Node* node)
621	{
622		fAllocator.Destruct(node);
623		fAllocator.Deallocate(node);
624	}
625
626	inline const Key& GetKey(const Node* node) const
627	{
628		return node->key;
629	}
630
631	inline Value& GetValue(Node* node) const
632	{
633		return node->value;
634	}
635
636	inline AVLTreeNode* GetAVLTreeNode(Node* node) const
637	{
638		return node;
639	}
640
641	inline Node* GetNode(AVLTreeNode* node) const
642	{
643		return static_cast<Node*>(node);
644	}
645
646	inline int CompareKeyNode(const Key& a, const Node* b) const
647	{
648		return fCompare(a, GetKey(b));
649	}
650
651	inline int CompareNodes(const Node* a, const Node* b) const
652	{
653		return fCompare(GetKey(a), GetKey(b));
654	}
655
656private:
657	KeyOrder		fCompare;
658	Allocator<Node>	fAllocator;
659};
660}
661
662#endif	// _KERNEL_UTIL_AVL_TREE_MAP_H
663