From d445953603e66eb8992a39b4e10fcafec8501f24 Mon Sep 17 00:00:00 2001
From: Alyssa Haroldsen <kupiakos@google.com>
Date: Mon, 5 Feb 2024 18:28:08 -0800
Subject: [PATCH] Add map key and value iterators

PiperOrigin-RevId: 604495084
---
 rust/map.rs                            | 36 +++++++++++++++++++++++++-
 rust/test/shared/accessors_map_test.rs | 32 ++++++++++++++++++++---
 2 files changed, 63 insertions(+), 5 deletions(-)

diff --git a/rust/map.rs b/rust/map.rs
index 5ca5af404f..76814f462a 100644
--- a/rust/map.rs
+++ b/rust/map.rs
@@ -215,10 +215,27 @@ where
         self.len() == 0
     }
 
-    /// An alias for `<Self as IntoIterator>::into_iterator`.
+    /// Returns an iterator visiting all key-value pairs in arbitrary order.
+    ///
+    /// The iterator element type is `(View<K>, View<V>)`.
+    /// This is an alias for `<Self as IntoIterator>::into_iter`.
     pub fn iter(self) -> MapIter<'msg, K, V> {
         self.into_iter()
     }
+
+    /// Returns an iterator visiting all keys in arbitrary order.
+    ///
+    /// The iterator element type is `View<K>`.
+    pub fn keys(self) -> impl Iterator<Item = View<'msg, K>> + 'msg {
+        self.into_iter().map(|(k, _)| k)
+    }
+
+    /// Returns an iterator visiting all values in arbitrary order.
+    ///
+    /// The iterator element type is `View<V>`.
+    pub fn values(self) -> impl Iterator<Item = View<'msg, V>> + 'msg {
+        self.into_iter().map(|(_, v)| v)
+    }
 }
 
 #[doc(hidden)]
@@ -286,9 +303,26 @@ where
         todo!("implement b/28530933");
     }
 
+    /// Returns an iterator visiting all key-value pairs in arbitrary order.
+    ///
+    /// The iterator element type is `(View<K>, View<V>)`.
     pub fn iter(&self) -> MapIter<'_, K, V> {
         self.into_iter()
     }
+
+    /// Returns an iterator visiting all keys in arbitrary order.
+    ///
+    /// The iterator element type is `View<K>`.
+    pub fn keys(&self) -> impl Iterator<Item = View<'_, K>> + '_ {
+        self.as_view().keys()
+    }
+
+    /// Returns an iterator visiting all values in arbitrary order.
+    ///
+    /// The iterator element type is `View<V>`.
+    pub fn values(&self) -> impl Iterator<Item = View<'_, V>> + '_ {
+        self.as_view().values()
+    }
 }
 
 /// An iterator visiting all key-value pairs in arbitrary order.
diff --git a/rust/test/shared/accessors_map_test.rs b/rust/test/shared/accessors_map_test.rs
index fd7284aac2..159d1e381f 100644
--- a/rust/test/shared/accessors_map_test.rs
+++ b/rust/test/shared/accessors_map_test.rs
@@ -25,6 +25,14 @@ macro_rules! generate_map_primitives_tests {
                     msg.[< map_ $k_field _ $v_field >]().iter().collect::<Vec<_>>(),
                     elements_are![]
                 );
+                assert_that!(
+                    msg.[< map_ $k_field _ $v_field >]().keys().collect::<Vec<_>>(),
+                    elements_are![]
+                );
+                assert_that!(
+                    msg.[< map_ $k_field _ $v_field >]().values().collect::<Vec<_>>(),
+                    elements_are![]
+                );
                 let k = <$k_type>::default();
                 let v = <$v_type>::default();
                 assert_that!(msg.[< map_ $k_field _ $v_field _mut>]().insert(k, v), eq(true));
@@ -34,18 +42,34 @@ macro_rules! generate_map_primitives_tests {
                     msg.[< map_ $k_field _ $v_field >]().iter().collect::<Vec<_>>(),
                     elements_are![eq((k, v))]
                 );
+                assert_that!(
+                    msg.[< map_ $k_field _ $v_field >]().keys().collect::<Vec<_>>(),
+                    elements_are![eq(k)]
+                );
+                assert_that!(
+                    msg.[< map_ $k_field _ $v_field >]().values().collect::<Vec<_>>(),
+                    elements_are![eq(v)]
+                );
 
-                let k: $k_type = $k_nonzero;
-                let v: $v_type = $v_nonzero;
-                assert_that!(msg.[< map_ $k_field _ $v_field _mut>]().insert(k, v), eq(true));
+                let k2: $k_type = $k_nonzero;
+                let v2: $v_type = $v_nonzero;
+                assert_that!(msg.[< map_ $k_field _ $v_field _mut>]().insert(k2, v2), eq(true));
                 assert_that!(msg.[< map_ $k_field _ $v_field >]().len(), eq(2));
                 assert_that!(
                     msg.[< map_ $k_field _ $v_field >]().iter().collect::<Vec<_>>(),
                     unordered_elements_are![
                         eq((k, v)),
-                        eq((<$k_type>::default(), <$v_type>::default())),
+                        eq((k2, v2)),
                     ]
                 );
+                assert_that!(
+                    msg.[< map_ $k_field _ $v_field >]().keys().collect::<Vec<_>>(),
+                    unordered_elements_are![eq(k), eq(k2)]
+                );
+                assert_that!(
+                    msg.[< map_ $k_field _ $v_field >]().values().collect::<Vec<_>>(),
+                    unordered_elements_are![eq(v), eq(v2)]
+                );
             }
         )* }
     };