diff --git a/src/libstd/trie.rs b/src/libstd/trie.rs index 679d36b87f1..93a11bdea98 100644 --- a/src/libstd/trie.rs +++ b/src/libstd/trie.rs @@ -14,6 +14,7 @@ use prelude::*; use iterator::{IteratorUtil, FromIterator, Extendable}; use uint; use util::{swap, replace}; +use vec; // FIXME: #5244: need to manually update the TrieNode constructor static SHIFT: uint = 4; @@ -146,6 +147,15 @@ impl TrieMap { pub fn each_value_reverse(&self, f: &fn(&T) -> bool) -> bool { self.each_reverse(|_, v| f(v)) } + + /// Get an iterator over the key-value pairs in the map + pub fn iter<'a>(&'a self) -> TrieMapIterator<'a, T> { + TrieMapIterator { + stack: ~[self.root.children.iter()], + remaining_min: self.length, + remaining_max: self.length + } + } } impl> FromIterator<(uint, T), Iter> for TrieMap { @@ -217,6 +227,12 @@ impl TrieSet { pub fn each_reverse(&self, f: &fn(&uint) -> bool) -> bool { self.map.each_key_reverse(f) } + + /// Get an iterator over the values in the set + #[inline] + pub fn iter<'a>(&'a self) -> TrieSetIterator<'a> { + TrieSetIterator{iter: self.map.iter()} + } } impl> FromIterator for TrieSet { @@ -366,6 +382,61 @@ fn remove(count: &mut uint, child: &mut Child, key: uint, return ret; } +/// Forward iterator over a map +pub struct TrieMapIterator<'self, T> { + priv stack: ~[vec::VecIterator<'self, Child>], + priv remaining_min: uint, + priv remaining_max: uint +} + +impl<'self, T> Iterator<(uint, &'self T)> for TrieMapIterator<'self, T> { + fn next(&mut self) -> Option<(uint, &'self T)> { + while !self.stack.is_empty() { + match self.stack[self.stack.len() - 1].next() { + None => { + self.stack.pop(); + } + Some(ref child) => { + match **child { + Internal(ref node) => { + self.stack.push(node.children.iter()); + } + External(key, ref value) => { + self.remaining_max -= 1; + if self.remaining_min > 0 { + self.remaining_min -= 1; + } + return Some((key, value)); + } + Nothing => {} + } + } + } + } + return None; + } + + #[inline] + fn size_hint(&self) -> (uint, Option) { + (self.remaining_min, Some(self.remaining_max)) + } +} + +/// Forward iterator over a set +pub struct TrieSetIterator<'self> { + priv iter: TrieMapIterator<'self, ()> +} + +impl<'self> Iterator for TrieSetIterator<'self> { + fn next(&mut self) -> Option { + do self.iter.next().map |&(key, _)| { key } + } + + fn size_hint(&self) -> (uint, Option) { + self.iter.size_hint() + } +} + #[cfg(test)] pub fn check_integrity(trie: &TrieNode) { assert!(trie.count != 0); @@ -553,6 +624,31 @@ mod test_map { assert_eq!(map.find(&k), Some(&v)); } } + + #[test] + fn test_iteration() { + let empty_map : TrieMap = TrieMap::new(); + assert_eq!(empty_map.iter().next(), None); + + let first = uint::max_value - 10000; + let last = uint::max_value; + + let mut map = TrieMap::new(); + do uint::range_rev(last, first) |x| { + printfln!("%x", x); + map.insert(x, x / 2); + true + }; + + let mut i = 0; + for (k, &v) in map.iter() { + printfln!("%x", k); + assert_eq!(k, first + i); + assert_eq!(v, k / 2); + i += 1; + } + assert_eq!(i, last - first); + } } #[cfg(test)]