From f717fda9a33aa47ccc205145b1785ec402a0d19e Mon Sep 17 00:00:00 2001
From: Edgar Xi <edgarxi97@gmail.com>
Date: Mon, 21 Mar 2022 18:48:51 -0400
Subject: [PATCH] modify get_protobuf_or_bincode_cells to accept and return an
 iterator

---
 storage-bigtable/src/bigtable.rs | 13 ++++++-------
 storage-bigtable/src/lib.rs      | 27 ++++++++++++++-------------
 2 files changed, 20 insertions(+), 20 deletions(-)

diff --git a/storage-bigtable/src/bigtable.rs b/storage-bigtable/src/bigtable.rs
index 23c6719872..683aee3bca 100644
--- a/storage-bigtable/src/bigtable.rs
+++ b/storage-bigtable/src/bigtable.rs
@@ -680,17 +680,17 @@ impl<F: FnMut(Request<()>) -> InterceptedRequestResult> BigTable<F> {
         deserialize_protobuf_or_bincode_cell_data(&row_data, table, key)
     }
 
-    pub async fn get_protobuf_or_bincode_cells<B, P>(
+    pub async fn get_protobuf_or_bincode_cells<'a, B, P>(
         &mut self,
-        table: &str,
-        row_keys: &[RowKey],
-    ) -> Result<Vec<(RowKey, CellData<B, P>)>>
+        table: &'a str,
+        row_keys: impl Iterator<Item = String>,
+    ) -> Result<impl Iterator<Item = (RowKey, CellData<B, P>)> + 'a>
     where
         B: serde::de::DeserializeOwned,
         P: prost::Message + Default,
     {
         Ok(self
-            .get_multi_row_data(table, row_keys)
+            .get_multi_row_data(table, row_keys.collect::<Vec<RowKey>>().as_slice())
             .await?
             .into_iter()
             .map(|(key, row_data)| {
@@ -699,8 +699,7 @@ impl<F: FnMut(Request<()>) -> InterceptedRequestResult> BigTable<F> {
                     key,
                     deserialize_protobuf_or_bincode_cell_data(&row_data, table, key_str).unwrap(),
                 )
-            })
-            .collect())
+            }))
     }
 
     pub async fn put_bincode_cells<T>(
diff --git a/storage-bigtable/src/lib.rs b/storage-bigtable/src/lib.rs
index 123793a20b..4f61b9f478 100644
--- a/storage-bigtable/src/lib.rs
+++ b/storage-bigtable/src/lib.rs
@@ -466,21 +466,22 @@ impl LedgerStorage {
         );
         inc_new_counter_debug!("storage-bigtable-query", 1);
         let mut bigtable = self.connection.client();
-        let row_keys: Vec<RowKey> = slots.iter().copied().map(slot_to_blocks_key).collect();
+        let row_keys = slots.iter().copied().map(slot_to_blocks_key);
         let data: Vec<(Slot, ConfirmedBlock)> = bigtable
-            .get_protobuf_or_bincode_cells::<StoredConfirmedBlock, generated::ConfirmedBlock>(
-                "blocks",
-                row_keys.as_slice(),
-            )
+            .get_protobuf_or_bincode_cells("blocks", row_keys)
             .await?
-            .into_iter()
-            .filter_map(|(row_key, block_cell_data)| {
-                let block = match block_cell_data {
-                    bigtable::CellData::Bincode(block) => block.into(),
-                    bigtable::CellData::Protobuf(block) => block.try_into().ok()?,
-                };
-                Some((key_to_slot(&row_key).unwrap(), block))
-            })
+            .filter_map(
+                |(row_key, block_cell_data): (
+                    RowKey,
+                    bigtable::CellData<StoredConfirmedBlock, generated::ConfirmedBlock>,
+                )| {
+                    let block = match block_cell_data {
+                        bigtable::CellData::Bincode(block) => block.into(),
+                        bigtable::CellData::Protobuf(block) => block.try_into().ok()?,
+                    };
+                    Some((key_to_slot(&row_key).unwrap(), block))
+                },
+            )
             .collect();
         Ok(data)
     }