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}
273
274
275// RootNode
276_AVL_TREE_MAP_TEMPLATE_LIST
277inline typename _AVL_TREE_MAP_CLASS_NAME::Node*
278_AVL_TREE_MAP_CLASS_NAME::RootNode() const
279{
280	if (AVLTreeNode* root = fTree.Root())
281		return _GetNode(root);
282	return NULL;
283}
284
285
286// Previous
287_AVL_TREE_MAP_TEMPLATE_LIST
288inline typename _AVL_TREE_MAP_CLASS_NAME::Node*
289_AVL_TREE_MAP_CLASS_NAME::Previous(Node* node) const
290{
291	if (node == NULL)
292		return NULL;
293
294	AVLTreeNode* treeNode = fTree.Previous(_GetAVLTreeNode(node));
295	return treeNode != NULL ? _GetNode(treeNode) : NULL;
296}
297
298
299// Next
300_AVL_TREE_MAP_TEMPLATE_LIST
301inline typename _AVL_TREE_MAP_CLASS_NAME::Node*
302_AVL_TREE_MAP_CLASS_NAME::Next(Node* node) const
303{
304	if (node == NULL)
305		return NULL;
306
307	AVLTreeNode* treeNode = fTree.Next(_GetAVLTreeNode(node));
308	return treeNode != NULL ? _GetNode(treeNode) : NULL;
309}
310
311
312// GetIterator
313_AVL_TREE_MAP_TEMPLATE_LIST
314inline typename _AVL_TREE_MAP_CLASS_NAME::Iterator
315_AVL_TREE_MAP_CLASS_NAME::GetIterator()
316{
317	return Iterator(this, fTree.GetIterator());
318}
319
320
321// GetIterator
322_AVL_TREE_MAP_TEMPLATE_LIST
323inline typename _AVL_TREE_MAP_CLASS_NAME::ConstIterator
324_AVL_TREE_MAP_CLASS_NAME::GetIterator() const
325{
326	return ConstIterator(this, fTree.GetIterator());
327}
328
329
330// GetIterator
331_AVL_TREE_MAP_TEMPLATE_LIST
332inline typename _AVL_TREE_MAP_CLASS_NAME::Iterator
333_AVL_TREE_MAP_CLASS_NAME::GetIterator(Node* node)
334{
335	return Iterator(this, fTree.GetIterator(_GetAVLTreeNode(node)));
336}
337
338
339// GetIterator
340_AVL_TREE_MAP_TEMPLATE_LIST
341inline typename _AVL_TREE_MAP_CLASS_NAME::ConstIterator
342_AVL_TREE_MAP_CLASS_NAME::GetIterator(Node* node) const
343{
344	return ConstIterator(this, fTree.GetIterator(_GetAVLTreeNode(node)));
345}
346
347
348// Find
349_AVL_TREE_MAP_TEMPLATE_LIST
350typename _AVL_TREE_MAP_CLASS_NAME::Iterator
351_AVL_TREE_MAP_CLASS_NAME::Find(const Key& key)
352{
353	if (AVLTreeNode* node = fTree.Find(&key))
354		return Iterator(this, fTree.GetIterator(node));
355	return Iterator();
356}
357
358
359// FindClose
360_AVL_TREE_MAP_TEMPLATE_LIST
361typename _AVL_TREE_MAP_CLASS_NAME::Iterator
362_AVL_TREE_MAP_CLASS_NAME::FindClose(const Key& key, bool less)
363{
364	if (AVLTreeNode* node = fTree.FindClosest(&key, less))
365		return Iterator(this, fTree.GetIterator(node));
366	return Iterator();
367}
368
369
370// Insert
371_AVL_TREE_MAP_TEMPLATE_LIST
372status_t
373_AVL_TREE_MAP_CLASS_NAME::Insert(const Key& key, const Value& value,
374	Iterator* iterator)
375{
376	// allocate a node
377	Node* userNode = _Allocate(key, value);
378	if (!userNode)
379		return B_NO_MEMORY;
380
381	// insert node
382	AVLTreeNode* node = _GetAVLTreeNode(userNode);
383	status_t error = fTree.Insert(node);
384	if (error != B_OK) {
385		_Free(userNode);
386		return error;
387	}
388
389	if (iterator)
390		*iterator = Iterator(this, fTree.GetIterator(node));
391
392	return B_OK;
393}
394
395
396// Insert
397_AVL_TREE_MAP_TEMPLATE_LIST
398status_t
399_AVL_TREE_MAP_CLASS_NAME::Insert(const Key& key, const Value& value,
400	Node** _node)
401{
402	// allocate a node
403	Node* userNode = _Allocate(key, value);
404	if (!userNode)
405		return B_NO_MEMORY;
406
407	// insert node
408	AVLTreeNode* node = _GetAVLTreeNode(userNode);
409	status_t error = fTree.Insert(node);
410	if (error != B_OK) {
411		_Free(userNode);
412		return error;
413	}
414
415	if (_node != NULL)
416		*_node = userNode;
417
418	return B_OK;
419}
420
421
422// Remove
423_AVL_TREE_MAP_TEMPLATE_LIST
424status_t
425_AVL_TREE_MAP_CLASS_NAME::Remove(const Key& key)
426{
427	AVLTreeNode* node = fTree.Remove(&key);
428	if (!node)
429		return B_ENTRY_NOT_FOUND;
430
431	_Free(_GetNode(node));
432	return B_OK;
433}
434
435
436// Remove
437_AVL_TREE_MAP_TEMPLATE_LIST
438status_t
439_AVL_TREE_MAP_CLASS_NAME::Remove(Node* node)
440{
441	if (!fTree.Remove(node))
442		return B_ENTRY_NOT_FOUND;
443
444	_Free(node);
445	return B_OK;
446}
447
448
449// CompareKeyNode
450_AVL_TREE_MAP_TEMPLATE_LIST
451int
452_AVL_TREE_MAP_CLASS_NAME::CompareKeyNode(const void* key,
453	const AVLTreeNode* node)
454{
455	return _CompareKeyNode(*(const Key*)key, _GetNode(node));
456}
457
458
459// CompareNodes
460_AVL_TREE_MAP_TEMPLATE_LIST
461int
462_AVL_TREE_MAP_CLASS_NAME::CompareNodes(const AVLTreeNode* node1,
463	const AVLTreeNode* node2)
464{
465	return _CompareNodes(_GetNode(node1), _GetNode(node2));
466}
467
468
469// _Allocate
470_AVL_TREE_MAP_TEMPLATE_LIST
471inline typename _AVL_TREE_MAP_CLASS_NAME::Node*
472_AVL_TREE_MAP_CLASS_NAME::_Allocate(const Key& key, const Value& value)
473{
474	return fStrategy.Allocate(key, value);
475}
476
477
478// _Free
479_AVL_TREE_MAP_TEMPLATE_LIST
480inline void
481_AVL_TREE_MAP_CLASS_NAME::_Free(Node* node)
482{
483	fStrategy.Free(node);
484}
485
486
487// _GetKey
488_AVL_TREE_MAP_TEMPLATE_LIST
489inline Key
490_AVL_TREE_MAP_CLASS_NAME::_GetKey(Node* node) const
491{
492	return fStrategy.GetKey(node);
493}
494
495
496// _GetValue
497_AVL_TREE_MAP_TEMPLATE_LIST
498inline Value&
499_AVL_TREE_MAP_CLASS_NAME::_GetValue(Node* node) const
500{
501	return fStrategy.GetValue(node);
502}
503
504
505// _GetAVLTreeNode
506_AVL_TREE_MAP_TEMPLATE_LIST
507inline AVLTreeNode*
508_AVL_TREE_MAP_CLASS_NAME::_GetAVLTreeNode(const Node* node) const
509{
510	return fStrategy.GetAVLTreeNode(const_cast<Node*>(node));
511}
512
513
514// _GetNode
515_AVL_TREE_MAP_TEMPLATE_LIST
516inline typename _AVL_TREE_MAP_CLASS_NAME::Node*
517_AVL_TREE_MAP_CLASS_NAME::_GetNode(const AVLTreeNode* node) const
518{
519	return fStrategy.GetNode(const_cast<AVLTreeNode*>(node));
520}
521
522
523// _CompareKeyNode
524_AVL_TREE_MAP_TEMPLATE_LIST
525inline int
526_AVL_TREE_MAP_CLASS_NAME::_CompareKeyNode(const Key& a, const Node* b)
527{
528	return fStrategy.CompareKeyNode(a, b);
529}
530
531
532// _CompareNodes
533_AVL_TREE_MAP_TEMPLATE_LIST
534inline int
535_AVL_TREE_MAP_CLASS_NAME::_CompareNodes(const Node* a, const Node* b)
536{
537	return fStrategy.CompareNodes(a, b);
538}
539
540
541// _FreeTree
542_AVL_TREE_MAP_TEMPLATE_LIST
543void
544_AVL_TREE_MAP_CLASS_NAME::_FreeTree(AVLTreeNode* node)
545{
546	if (node) {
547		_FreeTree(node->left);
548		_FreeTree(node->right);
549		_Free(_GetNode(node));
550	}
551}
552
553
554// #pragma mark - strategy parameters
555
556// Ascending
557namespace AVLTreeMapStrategy {
558template<typename Value>
559class Ascending {
560public:
561	inline int operator()(const Value &a, const Value &b) const
562	{
563		if (a < b)
564			return -1;
565		else if (a > b)
566			return 1;
567		return 0;
568	}
569};
570}
571
572
573// Descending
574namespace AVLTreeMapStrategy {
575template<typename Value>
576class Descending {
577public:
578	inline int operator()(const Value &a, const Value &b) const
579	{
580		if (a < b)
581			return -1;
582		else if (a > b)
583			return 1;
584		return 0;
585	}
586};
587}
588
589
590// #pragma mark - strategies
591
592
593// Auto
594namespace AVLTreeMapStrategy {
595template <typename Key, typename Value, typename KeyOrder,
596		  template <typename> class Allocator>
597class Auto {
598public:
599	struct Node : AVLTreeNode {
600		Node(const Key &key, const Value &value)
601			: AVLTreeNode(),
602			  key(key),
603			  value(value)
604		{
605		}
606
607		Key		key;
608		Value	value;
609	};
610
611	inline Node* Allocate(const Key& key, const Value& value)
612	{
613		Node* result = fAllocator.Allocate();
614		if (result)
615			fAllocator.Construct(result, key, value);
616		return result;
617	}
618
619	inline void Free(Node* node)
620	{
621		fAllocator.Destruct(node);
622		fAllocator.Deallocate(node);
623	}
624
625	inline const Key& GetKey(const Node* node) const
626	{
627		return node->key;
628	}
629
630	inline Value& GetValue(Node* node) const
631	{
632		return node->value;
633	}
634
635	inline AVLTreeNode* GetAVLTreeNode(Node* node) const
636	{
637		return node;
638	}
639
640	inline Node* GetNode(AVLTreeNode* node) const
641	{
642		return static_cast<Node*>(node);
643	}
644
645	inline int CompareKeyNode(const Key& a, const Node* b) const
646	{
647		return fCompare(a, GetKey(b));
648	}
649
650	inline int CompareNodes(const Node* a, const Node* b) const
651	{
652		return fCompare(GetKey(a), GetKey(b));
653	}
654
655private:
656	KeyOrder		fCompare;
657	Allocator<Node>	fAllocator;
658};
659}
660
661#endif	// _KERNEL_UTIL_AVL_TREE_MAP_H
662