From f8e0f1bb3f1fa0838a6278ecb410f4209ff5c42d Mon Sep 17 00:00:00 2001
From: Horst Schirmeier <horst.schirmeier@tu-dortmund.de>
Date: Wed, 14 Jan 2015 23:39:33 +0100
Subject: [PATCH] util: add SumTree::iterator

Change-Id: I8304b64634fa3ab92a126fe5d942674b26334b3d
---
 src/core/util/SumTree.hpp            | 59 ++++++++++++++++++++++++++++
 src/core/util/testing/SumTreeTest.cc |  8 +++-
 2 files changed, 66 insertions(+), 1 deletion(-)

diff --git a/src/core/util/SumTree.hpp b/src/core/util/SumTree.hpp
index 9d50269b..d5df54da 100644
--- a/src/core/util/SumTree.hpp
+++ b/src/core/util/SumTree.hpp
@@ -4,6 +4,7 @@
 #include <assert.h>
 #include <stdint.h>
 #include <vector>
+#include <stack>
 
 // The SumTree implements an efficient tree data structure for
 // "roulette-wheel" sampling, or "sampling with fault expansion", i.e.,
@@ -44,6 +45,64 @@ class SumTree {
 		std::vector<T> elements;
 	};
 
+public:
+	//! Iterator
+	class TreeIterator : public std::iterator<std::input_iterator_tag, T> {
+		//! Buckets and corresponding element indexes down the tree
+		std::stack<std::pair<Bucket *, unsigned> > hierarchy;
+	public:
+		TreeIterator() {}
+//MyIterator(int* x) :p(x) {}
+		TreeIterator(const TreeIterator& i) : hierarchy(i.hierarchy) { }
+		TreeIterator(const SumTree<T, BUCKETSIZE>& tree)
+		{
+			// go down until we see leaves
+			hierarchy.push(std::pair<Bucket *, unsigned>(tree.m_root, 0));
+			while (!hierarchy.top().first->elements.size() && hierarchy.top().first->children.size() > 0) {
+				hierarchy.push(std::pair<Bucket *, unsigned>(hierarchy.top().first->children[hierarchy.top().second], 0));
+			}
+		}
+		TreeIterator& operator++()
+		{
+			// advance index in the current level
+			hierarchy.top().second++;
+			if (hierarchy.top().second < hierarchy.top().first->elements.size()) {
+				return *this;
+			}
+			// current level is exhausted, go back up to a not yet finished level
+			do {
+				hierarchy.pop();
+			} while (!hierarchy.empty()
+				&& ++hierarchy.top().second >= hierarchy.top().first->children.size());
+			// at the end?
+			if (hierarchy.empty()) {
+				return *this;
+			}
+			// go down until we see leaves again
+			do {
+				hierarchy.push(std::pair<Bucket *, unsigned>(hierarchy.top().first->children[hierarchy.top().second], 0));
+			} while (!hierarchy.top().first->elements.size() && hierarchy.top().first->children.size() > 0);
+			return *this;
+		}
+		TreeIterator operator++(int) { TreeIterator tmp(*this); operator++(); return tmp; }
+	    bool operator==(const TreeIterator& rhs) { return hierarchy == rhs.hierarchy; }
+		bool operator!=(const TreeIterator& rhs) { return hierarchy != rhs.hierarchy; }
+		T& operator*() { return hierarchy.top().first->elements[hierarchy.top().second]; }
+		T *operator->() { return &(operator*()); }
+	};
+	typedef TreeIterator iterator;
+
+	iterator begin()
+	{
+		return iterator(*this);
+	}
+
+	iterator end()
+	{
+		return iterator();
+	}
+
+private:
 	//! Root node
 	Bucket *m_root;
 	//! Tree depth: nodes at level m_depth are leaf nodes, others are inner nodes
diff --git a/src/core/util/testing/SumTreeTest.cc b/src/core/util/testing/SumTreeTest.cc
index 9824a034..ad9d717b 100644
--- a/src/core/util/testing/SumTreeTest.cc
+++ b/src/core/util/testing/SumTreeTest.cc
@@ -17,13 +17,19 @@ struct Pilot {
 
 int main()
 {
-	fail::SumTree<Pilot, 2> tree;
+	typedef fail::SumTree<Pilot, 2> sumtree_type;
+	sumtree_type tree;
 	for (int i = 0; i <= 20; ++i) {
 		Pilot p;
 		p.duration = i;
 		tree.add(p);
 	}
 
+	LOG << "tree contents:" << endl;
+	for (sumtree_type::iterator it = tree.begin(); it != tree.end(); ++it) {
+		LOG << it->size() << endl;
+	}
+
 	while (tree.get_size() > 0) {
 		uint64_t pos = tree.get_size() / 2;
 		LOG << "MAIN tree.get_size() = " << tree.get_size()