From 71ea05c1768a195f4f10498731e3dd506be5917a Mon Sep 17 00:00:00 2001
From: Tao Zhu <tao@solana.com>
Date: Fri, 18 Mar 2022 09:44:42 -0500
Subject: [PATCH] replace nested for_each with flat_map

---
 core/src/unprocessed_packet_batches.rs | 34 +++++++++++++-------------
 1 file changed, 17 insertions(+), 17 deletions(-)

diff --git a/core/src/unprocessed_packet_batches.rs b/core/src/unprocessed_packet_batches.rs
index c891b1185f..87af6f0445 100644
--- a/core/src/unprocessed_packet_batches.rs
+++ b/core/src/unprocessed_packet_batches.rs
@@ -107,32 +107,32 @@ impl UnprocessedPacketBatches {
             .sum()
     }
 
-    /// Iterates packets in buffered packet_batches, returns all unprocessed packet's stake,
-    /// and its locator
+    /// Iterates the inner `Vec<DeserializedPacketBatch>`.
+    /// Returns the flattened result of mapping each
+    /// `DeserializedPacketBatch` to a list the batch's inner
+    /// packets' sender's stake and their `PacketLocator`'s within the
+    /// `Vec<DeserializedPacketBatch>`.
     #[allow(dead_code)]
     fn get_stakes_and_locators(&self) -> (Vec<u64>, Vec<PacketLocator>) {
-        let num_unprocessed_packets = self.get_unprocessed_packets_count();
-        let mut stakes = Vec::<u64>::with_capacity(num_unprocessed_packets);
-        let mut locators = Vec::<PacketLocator>::with_capacity(num_unprocessed_packets);
-
         self.iter()
             .enumerate()
-            .for_each(|(batch_index, deserialized_packet_batch)| {
+            .flat_map(|(batch_index, deserialized_packet_batch)| {
                 let packet_batch = &deserialized_packet_batch.packet_batch;
                 deserialized_packet_batch
                     .unprocessed_packets
                     .keys()
-                    .for_each(|packet_index| {
+                    .map(move |packet_index| {
                         let p = &packet_batch.packets[*packet_index];
-                        stakes.push(p.meta.sender_stake);
-                        locators.push(PacketLocator {
-                            batch_index,
-                            packet_index: *packet_index,
-                        });
-                    });
-            });
-
-        (stakes, locators)
+                        (
+                            p.meta.sender_stake,
+                            PacketLocator {
+                                batch_index,
+                                packet_index: *packet_index,
+                            },
+                        )
+                    })
+            })
+            .unzip()
     }
 }