xref: /haiku/headers/private/kernel/util/AVLTreeMap.h (revision 9760dcae2038d47442f4658c2575844c6cf92c40)
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
14 namespace 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 Key GetKey(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
46 template<typename Key, typename Value,
47 		 typename NodeStrategy = AVLTreeMapStrategy::Auto<Key, Value> >
48 class AVLTreeMap : protected AVLTreeCompare {
49 private:
50 	typedef typename NodeStrategy::Node	Node;
51 	typedef _AVL_TREE_MAP_CLASS_NAME	Class;
52 
53 public:
54 	class Iterator;
55 	class ConstIterator;
56 
57 public:
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 	inline	Iterator			GetIterator();
69 	inline	ConstIterator		GetIterator() const;
70 
71 	inline	Iterator			GetIterator(Node* node);
72 	inline	ConstIterator		GetIterator(Node* node) const;
73 
74 			Iterator			Find(const Key& key);
75 			Iterator			FindClose(const Key& key, bool less);
76 
77 			status_t			Insert(const Key& key, const Value& value,
78 									Iterator* iterator);
79 			status_t			Remove(const Key& key);
80 
81 			const NodeStrategy&	GetNodeStrategy() const	{ return fStrategy; }
82 
83 protected:
84 	// AVLTreeCompare
85 	virtual	int					CompareKeyNode(const void* key,
86 									const AVLTreeNode* node);
87 	virtual	int					CompareNodes(const AVLTreeNode* node1,
88 									const AVLTreeNode* node2);
89 
90 			void				_FreeTree(AVLTreeNode* node);
91 
92 	// strategy shortcuts
93 	inline	Node*				_Allocate(const Key& key, const Value& value);
94 	inline	void				_Free(Node* node);
95 	inline	const Key&			_GetKey(Node* node) const;
96 	inline	Value&				_GetValue(Node* node) const;
97 	inline	AVLTreeNode*		_GetAVLTreeNode(const Node* node) const;
98 	inline	Node*				_GetNode(const AVLTreeNode* node) const;
99 	inline	int					_CompareKeyNode(const Key& a, const Node* b);
100 	inline	int					_CompareNodes(const Node* a, const Node* b);
101 
102 protected:
103 			friend class Iterator;
104 			friend class ConstIterator;
105 
106 			AVLTreeBase			fTree;
107 			NodeStrategy		fStrategy;
108 
109 public:
110 	// Iterator
111 	// (need to implement it here, otherwise gcc 2.95.3 chokes)
112 	class Iterator : public ConstIterator {
113 	public:
114 		inline Iterator()
115 			: ConstIterator()
116 		{
117 		}
118 
119 		inline Iterator(const Iterator& other)
120 			: ConstIterator(other)
121 		{
122 		}
123 
124 		inline void Remove()
125 		{
126 			if (AVLTreeNode* node = ConstIterator::fTreeIterator.Remove()) {
127 				_AVL_TREE_MAP_CLASS_NAME* parent
128 					= const_cast<_AVL_TREE_MAP_CLASS_NAME*>(
129 						ConstIterator::fParent);
130 				parent->_Free(parent->_GetNode(node));
131 			}
132 		}
133 
134 	private:
135 		inline Iterator(_AVL_TREE_MAP_CLASS_NAME* parent,
136 			const AVLTreeIterator& treeIterator)
137 			: ConstIterator(parent, treeIterator)
138 		{
139 		}
140 
141 		friend class _AVL_TREE_MAP_CLASS_NAME;
142 	};
143 };
144 
145 
146 // ConstIterator
147 _AVL_TREE_MAP_TEMPLATE_LIST
148 class _AVL_TREE_MAP_CLASS_NAME::ConstIterator {
149 public:
150 	inline ConstIterator()
151 		: fParent(NULL),
152 		  fTreeIterator()
153 	{
154 	}
155 
156 	inline ConstIterator(const ConstIterator& other)
157 		: fParent(other.fParent),
158 		  fTreeIterator(other.fTreeIterator)
159 	{
160 	}
161 
162 	inline bool HasCurrent() const
163 	{
164 		return fTreeIterator.Current();
165 	}
166 
167 	inline Key CurrentKey()
168 	{
169 		if (AVLTreeNode* node = fTreeIterator.Current())
170 			return fParent->_GetKey(fParent->_GetNode(node));
171 		return Key();
172 	}
173 
174 	inline Value Current()
175 	{
176 		if (AVLTreeNode* node = fTreeIterator.Current())
177 			return fParent->_GetValue(fParent->_GetNode(node));
178 		return Value();
179 	}
180 
181 	inline Value* CurrentValuePointer()
182 	{
183 		if (AVLTreeNode* node = fTreeIterator.Current())
184 			return &fParent->_GetValue(fParent->_GetNode(node));
185 		return NULL;
186 	}
187 
188 	inline bool HasNext() const
189 	{
190 		return fTreeIterator.HasNext();
191 	}
192 
193 	inline Value Next()
194 	{
195 		if (AVLTreeNode* node = fTreeIterator.Next())
196 			return fParent->_GetValue(fParent->_GetNode(node));
197 		return Value();
198 	}
199 
200 	inline Value* NextValuePointer()
201 	{
202 		if (AVLTreeNode* node = fTreeIterator.Next())
203 			return &fParent->_GetValue(fParent->_GetNode(node));
204 		return NULL;
205 	}
206 
207 	inline Value Previous()
208 	{
209 		if (AVLTreeNode* node = fTreeIterator.Previous())
210 			return fParent->_GetValue(fParent->_GetNode(node));
211 		return Value();
212 	}
213 
214 	inline ConstIterator& operator=(const ConstIterator& other)
215 	{
216 		fParent = other.fParent;
217 		fTreeIterator = other.fTreeIterator;
218 		return *this;
219 	}
220 
221 protected:
222 	inline ConstIterator(const _AVL_TREE_MAP_CLASS_NAME* parent,
223 		const AVLTreeIterator& treeIterator)
224 	{
225 		fParent = parent;
226 		fTreeIterator = treeIterator;
227 	}
228 
229 	friend class _AVL_TREE_MAP_CLASS_NAME;
230 
231 	const _AVL_TREE_MAP_CLASS_NAME*	fParent;
232 	AVLTreeIterator					fTreeIterator;
233 };
234 
235 
236 // constructor
237 _AVL_TREE_MAP_TEMPLATE_LIST
238 _AVL_TREE_MAP_CLASS_NAME::AVLTreeMap(const NodeStrategy& strategy)
239 	: fTree(this),
240 	  fStrategy(strategy)
241 {
242 }
243 
244 
245 // destructor
246 _AVL_TREE_MAP_TEMPLATE_LIST
247 _AVL_TREE_MAP_CLASS_NAME::~AVLTreeMap()
248 {
249 }
250 
251 
252 // MakeEmpty
253 _AVL_TREE_MAP_TEMPLATE_LIST
254 inline void
255 _AVL_TREE_MAP_CLASS_NAME::MakeEmpty()
256 {
257 	AVLTreeNode* root = fTree.Root();
258 	_FreeTree(root);
259 }
260 
261 
262 // RootNode
263 _AVL_TREE_MAP_TEMPLATE_LIST
264 inline typename _AVL_TREE_MAP_CLASS_NAME::Node*
265 _AVL_TREE_MAP_CLASS_NAME::RootNode() const
266 {
267 	if (AVLTreeNode* root = fTree.Root())
268 		return _GetNode(root);
269 	return NULL;
270 }
271 
272 
273 // GetIterator
274 _AVL_TREE_MAP_TEMPLATE_LIST
275 inline typename _AVL_TREE_MAP_CLASS_NAME::Iterator
276 _AVL_TREE_MAP_CLASS_NAME::GetIterator()
277 {
278 	return Iterator(this, fTree.GetIterator());
279 }
280 
281 
282 // GetIterator
283 _AVL_TREE_MAP_TEMPLATE_LIST
284 inline typename _AVL_TREE_MAP_CLASS_NAME::ConstIterator
285 _AVL_TREE_MAP_CLASS_NAME::GetIterator() const
286 {
287 	return ConstIterator(this, fTree.GetIterator());
288 }
289 
290 
291 // GetIterator
292 _AVL_TREE_MAP_TEMPLATE_LIST
293 inline typename _AVL_TREE_MAP_CLASS_NAME::Iterator
294 _AVL_TREE_MAP_CLASS_NAME::GetIterator(Node* node)
295 {
296 	return Iterator(this, fTree.GetIterator(_GetAVLTreeNode(node)));
297 }
298 
299 
300 // GetIterator
301 _AVL_TREE_MAP_TEMPLATE_LIST
302 inline typename _AVL_TREE_MAP_CLASS_NAME::ConstIterator
303 _AVL_TREE_MAP_CLASS_NAME::GetIterator(Node* node) const
304 {
305 	return ConstIterator(this, fTree.GetIterator(_GetAVLTreeNode(node)));
306 }
307 
308 
309 // Find
310 _AVL_TREE_MAP_TEMPLATE_LIST
311 typename _AVL_TREE_MAP_CLASS_NAME::Iterator
312 _AVL_TREE_MAP_CLASS_NAME::Find(const Key& key)
313 {
314 	if (AVLTreeNode* node = fTree.Find(&key))
315 		return Iterator(this, fTree.GetIterator(node));
316 	return Iterator();
317 }
318 
319 
320 // FindClose
321 _AVL_TREE_MAP_TEMPLATE_LIST
322 typename _AVL_TREE_MAP_CLASS_NAME::Iterator
323 _AVL_TREE_MAP_CLASS_NAME::FindClose(const Key& key, bool less)
324 {
325 	if (AVLTreeNode* node = fTree.FindClosest(&key, less))
326 		return Iterator(this, fTree.GetIterator(node));
327 	return Iterator();
328 }
329 
330 
331 // Insert
332 _AVL_TREE_MAP_TEMPLATE_LIST
333 status_t
334 _AVL_TREE_MAP_CLASS_NAME::Insert(const Key& key, const Value& value,
335 	Iterator* iterator)
336 {
337 	// allocate a node
338 	Node* userNode = _Allocate(key, value);
339 	if (!userNode)
340 		return B_NO_MEMORY;
341 
342 	// insert node
343 	AVLTreeNode* node = _GetAVLTreeNode(userNode);
344 	status_t error = fTree.Insert(node);
345 	if (error != B_OK) {
346 		_Free(userNode);
347 		return error;
348 	}
349 
350 	if (iterator)
351 		*iterator = Iterator(this, fTree.GetIterator(node));
352 
353 	return B_OK;
354 }
355 
356 
357 // Remove
358 _AVL_TREE_MAP_TEMPLATE_LIST
359 status_t
360 _AVL_TREE_MAP_CLASS_NAME::Remove(const Key& key)
361 {
362 	AVLTreeNode* node = fTree.Remove(&key);
363 	if (!node)
364 		return B_ENTRY_NOT_FOUND;
365 
366 	_Free(_GetNode(node));
367 	return B_OK;
368 }
369 
370 
371 // CompareKeyNode
372 _AVL_TREE_MAP_TEMPLATE_LIST
373 int
374 _AVL_TREE_MAP_CLASS_NAME::CompareKeyNode(const void* key,
375 	const AVLTreeNode* node)
376 {
377 	return _CompareKeyNode(*(const Key*)key, _GetNode(node));
378 }
379 
380 
381 // CompareNodes
382 _AVL_TREE_MAP_TEMPLATE_LIST
383 int
384 _AVL_TREE_MAP_CLASS_NAME::CompareNodes(const AVLTreeNode* node1,
385 	const AVLTreeNode* node2)
386 {
387 	return _CompareNodes(_GetNode(node1), _GetNode(node2));
388 }
389 
390 
391 // _Allocate
392 _AVL_TREE_MAP_TEMPLATE_LIST
393 inline typename _AVL_TREE_MAP_CLASS_NAME::Node*
394 _AVL_TREE_MAP_CLASS_NAME::_Allocate(const Key& key, const Value& value)
395 {
396 	return fStrategy.Allocate(key, value);
397 }
398 
399 
400 // _Free
401 _AVL_TREE_MAP_TEMPLATE_LIST
402 inline void
403 _AVL_TREE_MAP_CLASS_NAME::_Free(Node* node)
404 {
405 	fStrategy.Free(node);
406 }
407 
408 
409 // _GetKey
410 _AVL_TREE_MAP_TEMPLATE_LIST
411 inline const Key&
412 _AVL_TREE_MAP_CLASS_NAME::_GetKey(Node* node) const
413 {
414 	return fStrategy.GetKey(node);
415 }
416 
417 
418 // _GetValue
419 _AVL_TREE_MAP_TEMPLATE_LIST
420 inline Value&
421 _AVL_TREE_MAP_CLASS_NAME::_GetValue(Node* node) const
422 {
423 	return fStrategy.GetValue(node);
424 }
425 
426 
427 // _GetAVLTreeNode
428 _AVL_TREE_MAP_TEMPLATE_LIST
429 inline AVLTreeNode*
430 _AVL_TREE_MAP_CLASS_NAME::_GetAVLTreeNode(const Node* node) const
431 {
432 	return fStrategy.GetAVLTreeNode(const_cast<Node*>(node));
433 }
434 
435 
436 // _GetNode
437 _AVL_TREE_MAP_TEMPLATE_LIST
438 inline typename _AVL_TREE_MAP_CLASS_NAME::Node*
439 _AVL_TREE_MAP_CLASS_NAME::_GetNode(const AVLTreeNode* node) const
440 {
441 	return fStrategy.GetNode(const_cast<AVLTreeNode*>(node));
442 }
443 
444 
445 // _CompareKeyNode
446 _AVL_TREE_MAP_TEMPLATE_LIST
447 inline int
448 _AVL_TREE_MAP_CLASS_NAME::_CompareKeyNode(const Key& a, const Node* b)
449 {
450 	return fStrategy.CompareKeyNode(a, b);
451 }
452 
453 
454 // _CompareNodes
455 _AVL_TREE_MAP_TEMPLATE_LIST
456 inline int
457 _AVL_TREE_MAP_CLASS_NAME::_CompareNodes(const Node* a, const Node* b)
458 {
459 	return fStrategy.CompareNodes(a, b);
460 }
461 
462 
463 // _FreeTree
464 _AVL_TREE_MAP_TEMPLATE_LIST
465 void
466 _AVL_TREE_MAP_CLASS_NAME::_FreeTree(AVLTreeNode* node)
467 {
468 	if (node) {
469 		_FreeTree(node->left);
470 		_FreeTree(node->right);
471 		_Free(_GetNode(node));
472 	}
473 }
474 
475 
476 // #pragma mark - strategy parameters
477 
478 // Ascending
479 namespace AVLTreeMapStrategy {
480 template<typename Value>
481 class Ascending {
482 public:
483 	inline int operator()(const Value &a, const Value &b) const
484 	{
485 		if (a < b)
486 			return -1;
487 		else if (a > b)
488 			return 1;
489 		return 0;
490 	}
491 };
492 }
493 
494 
495 // Descending
496 namespace AVLTreeMapStrategy {
497 template<typename Value>
498 class Descending {
499 public:
500 	inline int operator()(const Value &a, const Value &b) const
501 	{
502 		if (a < b)
503 			return -1;
504 		else if (a > b)
505 			return 1;
506 		return 0;
507 	}
508 };
509 }
510 
511 
512 // #pragma mark - strategies
513 
514 
515 // Auto
516 namespace AVLTreeMapStrategy {
517 template <typename Key, typename Value, typename KeyOrder,
518 		  template <typename> class Allocator>
519 class Auto {
520 public:
521 	struct Node : AVLTreeNode {
522 		Node(const Key &key, const Value &value)
523 			: AVLTreeNode(),
524 			  key(key),
525 			  value(value)
526 		{
527 		}
528 
529 		Key		key;
530 		Value	value;
531 	};
532 
533 	inline Node* Allocate(const Key& key, const Value& value)
534 	{
535 		Node* result = fAllocator.Allocate();
536 		if (result)
537 			fAllocator.Construct(result, key, value);
538 		return result;
539 	}
540 
541 	inline void Free(Node* node)
542 	{
543 		fAllocator.Destruct(node);
544 		fAllocator.Deallocate(node);
545 	}
546 
547 	inline Key& GetKey(Node* node) const
548 	{
549 		return node->key;
550 	}
551 
552 	inline Value& GetValue(Node* node) const
553 	{
554 		return node->value;
555 	}
556 
557 	inline AVLTreeNode* GetAVLTreeNode(Node* node) const
558 	{
559 		return node;
560 	}
561 
562 	inline Node* GetNode(AVLTreeNode* node) const
563 	{
564 		return static_cast<Node*>(node);
565 	}
566 
567 	inline int CompareKeyNode(const Key& a, const Node* b) const
568 	{
569 		return fCompare(a, _GetKey(b));
570 	}
571 
572 	inline int CompareNodes(const Node* a, const Node* b) const
573 	{
574 		return fCompare(_GetKey(a), _GetKey(b));
575 	}
576 
577 private:
578 	KeyOrder		fCompare;
579 	Allocator<Node>	fAllocator;
580 };
581 }
582 
583 #endif	// _KERNEL_UTIL_AVL_TREE_MAP_H
584