Implemented FromIterator for TrieMap and TrieSet

2 files changed, 91 insertions, 14 deletions

diff --git a/src/libstd/hashmap.rs b/src/libstd/hashmap.rs
index d62a445882e..c219886cd03 100644
--- a/src/libstd/hashmap.rs
+++ b/src/libstd/hashmap.rs
@@ -735,6 +735,20 @@ impl<T:Hash + Eq> HashSet<T> {
     }
 }
 
+impl<K: Eq + Hash, T: Iterator<K>> FromIterator<K, T> for HashSet<K> {
+    pub fn from_iterator(iter: &mut T) -> HashSet<K> {
+        let (lower, _) = iter.size_hint();
+        let mut set = HashSet::with_capacity(lower);
+
+        for iter.advance |k| {
+            set.insert(k);
+        }
+
+        set
+    }
+}
+  
+
 #[cfg(test)]
 mod test_map {
     use container::{Container, Map, Set};
@@ -1139,4 +1153,15 @@ mod test_set {
         }
         assert_eq!(i, expected.len());
     }
+
+    #[test]
+    fn test_from_iter() {
+        let xs = ~[1, 2, 3, 4, 5, 6, 7, 8, 9];
+        
+        let set: HashSet<int> = xs.iter().transform(|&x| x).collect();
+
+        for xs.iter().advance |x: &int| {
+            assert!(set.contains(x));
+        }
+    }
 }
diff --git a/src/libstd/trie.rs b/src/libstd/trie.rs
index 8ce02d59ab1..fd2cfa08b30 100644
--- a/src/libstd/trie.rs
+++ b/src/libstd/trie.rs
@@ -11,7 +11,7 @@
 //! An ordered map and set for integer keys implemented as a radix trie
 
 use prelude::*;
-use iterator::IteratorUtil;
+use iterator::{IteratorUtil, FromIterator};
 use uint;
 use util::{swap, replace};
 
@@ -171,6 +171,18 @@ impl<T> TrieMap<T> {
     }
 }
 
+impl<T, Iter: Iterator<(uint, T)>> FromIterator<(uint, T), Iter> for TrieMap<T> {
+    pub fn from_iterator(iter: &mut Iter) -> TrieMap<T> {
+        let mut map = TrieMap::new();
+
+        for iter.advance |(k, v)| {
+            map.insert(k, v);
+        }
+
+        map
+    }
+}
+
 #[allow(missing_doc)]
 pub struct TrieSet {
     priv map: TrieMap<()>
@@ -230,6 +242,18 @@ impl TrieSet {
     }
 }
 
+impl<Iter: Iterator<uint>> FromIterator<uint, Iter> for TrieSet {
+    pub fn from_iterator(iter: &mut Iter) -> TrieSet {
+        let mut set = TrieSet::new();
+
+        for iter.advance |elem| {
+            set.insert(elem);
+        }
+
+        set
+    }
+}
+
 struct TrieNode<T> {
     count: uint,
     children: [Child<T>, ..SIZE]
@@ -382,7 +406,7 @@ pub fn check_integrity<T>(trie: &TrieNode<T>) {
 }
 
 #[cfg(test)]
-mod tests {
+mod test_map {
     use super::*;
     use core::option::{Some, None};
     use uint;
@@ -511,6 +535,39 @@ mod tests {
     }
 
     #[test]
+    fn test_swap() {
+        let mut m = TrieMap::new();
+        assert_eq!(m.swap(1, 2), None);
+        assert_eq!(m.swap(1, 3), Some(2));
+        assert_eq!(m.swap(1, 4), Some(3));
+    }
+
+    #[test]
+    fn test_pop() {
+        let mut m = TrieMap::new();
+        m.insert(1, 2);
+        assert_eq!(m.pop(&1), Some(2));
+        assert_eq!(m.pop(&1), None);
+    }
+
+    #[test]
+    fn test_from_iter() {
+        let xs = ~[(1u, 1i), (2, 2), (3, 3), (4, 4), (5, 5), (6, 6)];
+        
+        let map: TrieMap<int> = xs.iter().transform(|&x| x).collect();
+
+        for xs.iter().advance |&(k, v)| {
+            assert_eq!(map.find(&k), Some(&v));
+        }
+    }
+}
+
+#[cfg(test)]
+mod test_set {
+    use super::*;
+    use uint;
+
+    #[test]
     fn test_sane_chunk() {
         let x = 1;
         let y = 1 << (uint::bits - 1);
@@ -533,18 +590,13 @@ mod tests {
     }
 
     #[test]
-    fn test_swap() {
-        let mut m = TrieMap::new();
-        assert_eq!(m.swap(1, 2), None);
-        assert_eq!(m.swap(1, 3), Some(2));
-        assert_eq!(m.swap(1, 4), Some(3));
-    }
+    fn test_from_iter() {
+        let xs = ~[9u, 8, 7, 6, 5, 4, 3, 2, 1];
+        
+        let set: TrieSet = xs.iter().transform(|&x| x).collect();
 
-    #[test]
-    fn test_pop() {
-        let mut m = TrieMap::new();
-        m.insert(1, 2);
-        assert_eq!(m.pop(&1), Some(2));
-        assert_eq!(m.pop(&1), None);
+        for xs.iter().advance |x| {
+            assert!(set.contains(x));
+        }
     }
 }