feat: use iterator in select_block_headers (#667)

crates/store/src/db/mod.rs

@@ -250,11 +250,14 @@ impl Db {
 
     /// Loads multiple block headers from the DB.
     #[instrument(target = COMPONENT, skip_all, ret(level = "debug"), err)]
-    pub async fn select_block_headers(&self, blocks: Vec<BlockNumber>) -> Result<Vec<BlockHeader>> {
+    pub async fn select_block_headers(
+        &self,
+        blocks: impl Iterator<Item = BlockNumber> + Send + 'static,
+    ) -> Result<Vec<BlockHeader>> {
         self.pool
             .get()
             .await?
-            .interact(move |conn| sql::select_block_headers(conn, &blocks))
+            .interact(move |conn| sql::select_block_headers(conn, blocks))
             .await
             .map_err(|err| {
                 DatabaseError::InteractError(format!(

crates/store/src/db/sql/mod.rs

@@ -1059,11 +1059,11 @@ pub fn select_block_header_by_block_num(
 /// A vector of [`BlockHeader`] or an error.
 pub fn select_block_headers(
     conn: &mut Connection,
-    blocks: &[BlockNumber],
+    blocks: impl Iterator<Item = BlockNumber> + Send,
 ) -> Result<Vec<BlockHeader>> {
-    let mut headers = Vec::with_capacity(blocks.len());
+    let blocks: Vec<Value> = blocks.map(|b| b.as_u32().into()).collect();
 
-    let blocks: Vec<Value> = blocks.iter().copied().map(|b| b.as_u32().into()).collect();
+    let mut headers = Vec::with_capacity(blocks.len());
     let mut stmt = conn
         .prepare_cached("SELECT block_header FROM block_headers WHERE block_num IN rarray(?1);")?;
     let mut rows = stmt.query(params![Rc::new(blocks)])?;

crates/store/src/state.rs

@@ -450,9 +450,7 @@ impl State {
         let blocks = note_proofs
             .values()
             .map(|proof| proof.location().block_num())
-            .collect::<BTreeSet<_>>()
-            .into_iter()
-            .collect::<Vec<_>>();
+            .collect::<BTreeSet<_>>();
 
         // Grab the block merkle paths from the inner state.
         //
@@ -479,7 +477,8 @@ impl State {
             (chain_length.into(), paths)
         };
 
-        let headers = self.db.select_block_headers(blocks).await?;
+        let headers = self.db.select_block_headers(blocks.into_iter()).await?;
+
         let headers = headers
             .into_iter()
             .map(|header| (header.block_num(), header))
@@ -596,15 +595,11 @@ impl State {
             (latest_block_num, partial_mmr)
         };
 
-        // TODO: Unnecessary conversion. We should change the select_block_headers function to take
-        // an impl Iterator instead to avoid this allocation.
-        let mut blocks: Vec<_> = blocks.into_iter().collect();
         // Fetch the reference block of the batch as part of this query, so we can avoid looking it
         // up in a separate DB access.
-        blocks.push(batch_reference_block);
         let mut headers = self
             .db
-            .select_block_headers(blocks)
+            .select_block_headers(blocks.into_iter().chain(std::iter::once(batch_reference_block)))
             .await
             .map_err(GetBatchInputsError::SelectBlockHeaderError)?;