auto merge of #15968 : nham/rust/keys_values_other_maps, r=alexcrichton

Adds methods for obtaining iterators over the keys or values of a SmallIntMap/TreeMap/TrieMap.

Closes #14376

3 files changed, 131 insertions, 0 deletions

diff --git a/src/libcollections/smallintmap.rs b/src/libcollections/smallintmap.rs
index 994a6d6c5f3..e1980e1549a 100644
--- a/src/libcollections/smallintmap.rs
+++ b/src/libcollections/smallintmap.rs
@@ -17,6 +17,7 @@ use core::prelude::*;
 
 use core::default::Default;
 use core::fmt;
+use core::iter;
 use core::iter::{Enumerate, FilterMap};
 use core::mem::replace;
 
@@ -194,6 +195,18 @@ impl<V> SmallIntMap<V> {
         self.find(key).expect("key not present")
     }
 
+    /// An iterator visiting all keys in ascending order by the keys.
+    /// Iterator element type is `uint`.
+    pub fn keys<'r>(&'r self) -> Keys<'r, V> {
+        self.iter().map(|(k, _v)| k)
+    }
+
+    /// An iterator visiting all values in ascending order by the keys.
+    /// Iterator element type is `&'r V`.
+    pub fn values<'r>(&'r self) -> Values<'r, V> {
+        self.iter().map(|(_k, v)| v)
+    }
+
     /// An iterator visiting all key-value pairs in ascending order by the keys.
     /// Iterator element type is `(uint, &'r V)`.
     ///
@@ -422,6 +435,14 @@ pub struct MutEntries<'a, T> {
 iterator!(impl MutEntries -> (uint, &'a mut T), get_mut_ref)
 double_ended_iterator!(impl MutEntries -> (uint, &'a mut T), get_mut_ref)
 
+/// Forward iterator over the keys of a map
+pub type Keys<'a, T> =
+    iter::Map<'static, (uint, &'a T), uint, Entries<'a, T>>;
+
+/// Forward iterator over the values of a map
+pub type Values<'a, T> =
+    iter::Map<'static, (uint, &'a T), &'a T, Entries<'a, T>>;
+
 #[cfg(test)]
 mod test_map {
     use std::prelude::*;
@@ -518,6 +539,32 @@ mod test_map {
     }
 
     #[test]
+    fn test_keys() {
+        let mut map = SmallIntMap::new();
+        map.insert(1, 'a');
+        map.insert(2, 'b');
+        map.insert(3, 'c');
+        let keys = map.keys().collect::<Vec<uint>>();
+        assert_eq!(keys.len(), 3);
+        assert!(keys.contains(&1));
+        assert!(keys.contains(&2));
+        assert!(keys.contains(&3));
+    }
+
+    #[test]
+    fn test_values() {
+        let mut map = SmallIntMap::new();
+        map.insert(1, 'a');
+        map.insert(2, 'b');
+        map.insert(3, 'c');
+        let values = map.values().map(|&v| v).collect::<Vec<char>>();
+        assert_eq!(values.len(), 3);
+        assert!(values.contains(&'a'));
+        assert!(values.contains(&'b'));
+        assert!(values.contains(&'c'));
+    }
+
+    #[test]
     fn test_iterator() {
         let mut m = SmallIntMap::new();
 
diff --git a/src/libcollections/treemap.rs b/src/libcollections/treemap.rs
index 5658d07a1d1..9bc0a1abbc5 100644
--- a/src/libcollections/treemap.rs
+++ b/src/libcollections/treemap.rs
@@ -141,6 +141,16 @@ impl<K: Ord, V> TreeMap<K, V> {
     /// Create an empty TreeMap
     pub fn new() -> TreeMap<K, V> { TreeMap{root: None, length: 0} }
 
+    /// Get a lazy iterator over the keys in the map.
+    pub fn keys<'a>(&'a self) -> Keys<'a, K, V> {
+        self.iter().map(|(k, _v)| k)
+    }
+
+    /// Get a lazy iterator over the values in the map.
+    pub fn values<'a>(&'a self) -> Values<'a, K, V> {
+        self.iter().map(|(_k, v)| v)
+    }
+
     /// Get a lazy iterator over the key-value pairs in the map.
     /// Requires that it be frozen (immutable).
     pub fn iter<'a>(&'a self) -> Entries<'a, K, V> {
@@ -381,6 +391,15 @@ pub struct RevMutEntries<'a, K, V> {
 }
 
 
+/// TreeMap keys iterator
+pub type Keys<'a, K, V> =
+    iter::Map<'static, (&'a K, &'a V), &'a K, Entries<'a, K, V>>;
+
+/// TreeMap values iterator
+pub type Values<'a, K, V> =
+    iter::Map<'static, (&'a K, &'a V), &'a V, Entries<'a, K, V>>;
+
+
 // FIXME #5846 we want to be able to choose between &x and &mut x
 // (with many different `x`) below, so we need to optionally pass mut
 // as a tt, but the only thing we can do with a `tt` is pass them to
@@ -1471,6 +1490,28 @@ mod test_treemap {
     }
 
     #[test]
+    fn test_keys() {
+        let vec = vec![(1i, 'a'), (2i, 'b'), (3i, 'c')];
+        let map = vec.move_iter().collect::<TreeMap<int, char>>();
+        let keys = map.keys().map(|&k| k).collect::<Vec<int>>();
+        assert_eq!(keys.len(), 3);
+        assert!(keys.contains(&1));
+        assert!(keys.contains(&2));
+        assert!(keys.contains(&3));
+    }
+
+    #[test]
+    fn test_values() {
+        let vec = vec![(1i, 'a'), (2i, 'b'), (3i, 'c')];
+        let map = vec.move_iter().collect::<TreeMap<int, char>>();
+        let values = map.values().map(|&v| v).collect::<Vec<char>>();
+        assert_eq!(values.len(), 3);
+        assert!(values.contains(&'a'));
+        assert!(values.contains(&'b'));
+        assert!(values.contains(&'c'));
+    }
+
+    #[test]
     fn test_eq() {
         let mut a = TreeMap::new();
         let mut b = TreeMap::new();
diff --git a/src/libcollections/trie.rs b/src/libcollections/trie.rs
index 5c17dd91225..14ab122f0a4 100644
--- a/src/libcollections/trie.rs
+++ b/src/libcollections/trie.rs
@@ -18,6 +18,7 @@ use core::default::Default;
 use core::mem::zeroed;
 use core::mem;
 use core::uint;
+use core::iter;
 use std::hash::{Writer, Hash};
 
 use {Collection, Mutable, Map, MutableMap, Set, MutableSet};
@@ -196,6 +197,18 @@ impl<T> TrieMap<T> {
         self.root.each_reverse(f)
     }
 
+    /// Get an iterator visiting all keys in ascending order by the keys.
+    /// Iterator element type is `uint`.
+    pub fn keys<'r>(&'r self) -> Keys<'r, T> {
+        self.iter().map(|(k, _v)| k)
+    }
+
+    /// Get an iterator visiting all values in ascending order by the keys.
+    /// Iterator element type is `&'r T`.
+    pub fn values<'r>(&'r self) -> Values<'r, T> {
+        self.iter().map(|(_k, v)| v)
+    }
+
     /// Get an iterator over the key-value pairs in the map, ordered by keys.
     ///
     /// # Example
@@ -783,6 +796,14 @@ pub struct MutEntries<'a, T> {
     remaining_max: uint
 }
 
+/// Forward iterator over the keys of a map
+pub type Keys<'a, T> =
+    iter::Map<'static, (uint, &'a T), uint, Entries<'a, T>>;
+
+/// Forward iterator over the values of a map
+pub type Values<'a, T> =
+    iter::Map<'static, (uint, &'a T), &'a T, Entries<'a, T>>;
+
 // FIXME #5846: see `addr!` above.
 macro_rules! item { ($i:item) => {$i}}
 
@@ -1071,6 +1092,28 @@ mod test_map {
     }
 
     #[test]
+    fn test_keys() {
+        let vec = vec![(1, 'a'), (2, 'b'), (3, 'c')];
+        let map = vec.move_iter().collect::<TrieMap<char>>();
+        let keys = map.keys().collect::<Vec<uint>>();
+        assert_eq!(keys.len(), 3);
+        assert!(keys.contains(&1));
+        assert!(keys.contains(&2));
+        assert!(keys.contains(&3));
+    }
+
+    #[test]
+    fn test_values() {
+        let vec = vec![(1, 'a'), (2, 'b'), (3, 'c')];
+        let map = vec.move_iter().collect::<TrieMap<char>>();
+        let values = map.values().map(|&v| v).collect::<Vec<char>>();
+        assert_eq!(values.len(), 3);
+        assert!(values.contains(&'a'));
+        assert!(values.contains(&'b'));
+        assert!(values.contains(&'c'));
+    }
+
+    #[test]
     fn test_iteration() {
         let empty_map : TrieMap<uint> = TrieMap::new();
         assert_eq!(empty_map.iter().next(), None);