From 0017a43364904ee09392fea4dfa6fee627d72c4a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tomasz=20Drwi=C4=99ga?= <tomusdrw@users.noreply.github.com>
Date: Thu, 20 Oct 2016 14:49:44 +0200
Subject: [PATCH] Making Trie.iter non-recursive (#2733)

---
 util/benches/trie.rs    | 28 +++++++++++++-
 util/src/trie/triedb.rs | 86 ++++++++++++++++++++++-------------------
 2 files changed, 73 insertions(+), 41 deletions(-)

diff --git a/util/benches/trie.rs b/util/benches/trie.rs
index 1c57ab969..925139032 100644
--- a/util/benches/trie.rs
+++ b/util/benches/trie.rs
@@ -21,7 +21,7 @@ extern crate ethcore_util;
 #[macro_use]
 extern crate log;
 
-use test::Bencher;
+use test::{Bencher, black_box};
 use ethcore_util::hash::*;
 use ethcore_util::bytes::*;
 use ethcore_util::trie::*;
@@ -78,6 +78,32 @@ fn trie_insertions_32_mir_1k(b: &mut Bencher) {
 	});
 //	println!("hash_count: {}", hash_count);
 }
+#[bench]
+fn trie_iter(b: &mut Bencher) {
+	let st = StandardMap {
+		alphabet: Alphabet::All,
+		min_key: 32,
+		journal_key: 0,
+		value_mode: ValueMode::Mirror,
+		count: 1000,
+	};
+	let d = st.make();
+	let mut memdb = MemoryDB::new();
+	let mut root = H256::new();
+	{
+		let mut t = TrieDBMut::new(&mut memdb, &mut root);
+		for i in d.iter() {
+			t.insert(&i.0, &i.1).unwrap();
+		}
+	}
+
+	b.iter(&mut ||{
+		let t = TrieDB::new(&memdb, &root).unwrap();
+		for n in t.iter().unwrap() {
+			black_box(n).unwrap();
+		}
+	});
+}
 
 #[bench]
 fn triehash_insertions_32_mir_1k(b: &mut Bencher) {
diff --git a/util/src/trie/triedb.rs b/util/src/trie/triedb.rs
index f5de26f8e..5522af448 100644
--- a/util/src/trie/triedb.rs
+++ b/util/src/trie/triedb.rs
@@ -304,14 +304,6 @@ impl<'a> TrieDBIterator<'a> {
 		Ok(())
 	}
 
-	/// Descend into a payload and get the next item.
-	fn descend_next(&mut self, d: &'a [u8]) -> Option<TrieItem<'a>> {
-		match self.descend(d) {
-			Ok(()) => self.next(),
-			Err(e) => Some(Err(e)),
-		}
-	}
-
 	/// The present key.
 	fn key(&self) -> Bytes {
 		// collapse the key_nibbles down to bytes.
@@ -323,38 +315,52 @@ impl<'a> Iterator for TrieDBIterator<'a> {
 	type Item = TrieItem<'a>;
 
 	fn next(&mut self) -> Option<Self::Item> {
-		let b = match self.trail.last_mut() {
-			Some(mut b) => { b.increment(); b.clone() },
-			None => return None,
-		};
-		match (b.status, b.node) {
-			(Status::Exiting, n) => {
-				match n {
-					Node::Leaf(n, _) | Node::Extension(n, _) => {
-						let l = self.key_nibbles.len();
-						self.key_nibbles.truncate(l - n.len());
-					},
-					Node::Branch(_, _) => { self.key_nibbles.pop(); },
-					_ => {}
-				}
-				self.trail.pop();
-				self.next()
-			},
-			(Status::At, Node::Leaf(_, v)) | (Status::At, Node::Branch(_, Some(v))) => Some(Ok((self.key(), v))),
-			(Status::At, Node::Extension(_, d)) => self.descend_next(d),
-			(Status::At, Node::Branch(_, _)) => self.next(),
-			(Status::AtChild(i), Node::Branch(children, _)) if children[i].len() > 0 => {
-				match i {
-					0 => self.key_nibbles.push(0),
-					i => *self.key_nibbles.last_mut().unwrap() = i as u8,
-				}
-				self.descend_next(children[i])
-			},
-			(Status::AtChild(i), Node::Branch(_, _)) => {
-				if i == 0 { self.key_nibbles.push(0); }
-				self.next()
-			},
-			_ => panic!() // Should never see Entering or AtChild without a Branch here.
+		loop {
+			let b = match self.trail.last_mut() {
+				Some(mut b) => { b.increment(); b.clone() },
+				None => return None,
+			};
+			match (b.status, b.node) {
+				(Status::Exiting, n) => {
+					match n {
+						Node::Leaf(n, _) | Node::Extension(n, _) => {
+							let l = self.key_nibbles.len();
+							self.key_nibbles.truncate(l - n.len());
+						},
+						Node::Branch(_, _) => { self.key_nibbles.pop(); },
+						_ => {}
+					}
+					self.trail.pop();
+					// continue
+				},
+				(Status::At, Node::Leaf(_, v)) | (Status::At, Node::Branch(_, Some(v))) => {
+					return Some(Ok((self.key(), v)));
+				},
+				(Status::At, Node::Extension(_, d)) => {
+					if let Err(e) = self.descend(d) {
+						return Some(Err(e));
+					}
+					// continue
+				},
+				(Status::At, Node::Branch(_, _)) => {},
+				(Status::AtChild(i), Node::Branch(children, _)) if children[i].len() > 0 => {
+					match i {
+						0 => self.key_nibbles.push(0),
+						i => *self.key_nibbles.last_mut().unwrap() = i as u8,
+					}
+					if let Err(e) = self.descend(children[i]) {
+						return Some(Err(e));
+					}
+					// continue
+				},
+				(Status::AtChild(i), Node::Branch(_, _)) => {
+					if i == 0 {
+						self.key_nibbles.push(0);
+					}
+					// continue
+				},
+				_ => panic!() // Should never see Entering or AtChild without a Branch here.
+			}
 		}
 	}
 }