Prepare subspace-farmer-components to retrieve pieces in batches

crates/subspace-archiving/src/piece_reconstructor.rs

@@ -212,6 +212,7 @@ impl PiecesReconstructor {
             return Err(ReconstructorError::IncorrectPiecePosition);
         }
 
+        // TODO: Early exit if position already exists and doesn't need reconstruction
         let (reconstructed_records, polynomial) = self.reconstruct_shards(segment_pieces)?;
 
         let mut piece = Piece::from(&reconstructed_records[piece_position]);

crates/subspace-farmer-components/src/lib.rs

@@ -6,6 +6,7 @@
     array_chunks,
     const_option,
     const_trait_impl,
+    exact_size_is_empty,
     int_roundings,
     iter_collect_into,
     never_type,
@@ -25,6 +26,8 @@ mod segment_reconstruction;
 
 use crate::file_ext::FileExt;
 use async_trait::async_trait;
+use futures::stream::FuturesUnordered;
+use futures::Stream;
 use parity_scale_codec::{Decode, Encode};
 use serde::{Deserialize, Serialize};
 use static_assertions::const_assert;
@@ -40,6 +43,28 @@ use subspace_core_primitives::segments::{ArchivedHistorySegment, HistorySize};
 pub trait PieceGetter {
     /// Get piece by index
     async fn get_piece(&self, piece_index: PieceIndex) -> anyhow::Result<Option<Piece>>;
+
+    /// Get pieces with provided indices
+    async fn get_pieces<'a, PieceIndices>(
+        &'a self,
+        piece_indices: PieceIndices,
+    ) -> anyhow::Result<
+        Box<dyn Stream<Item = (PieceIndex, anyhow::Result<Option<Piece>>)> + Send + Unpin + 'a>,
+    >
+    where
+        PieceIndices: IntoIterator<Item = PieceIndex, IntoIter: Send> + Send + 'a,
+    {
+        // TODO: Remove default impl here
+        Ok(Box::new(
+            piece_indices
+                .into_iter()
+                .map(|piece_index| async move {
+                    let result = self.get_piece(piece_index).await;
+                    (piece_index, result)
+                })
+                .collect::<FuturesUnordered<_>>(),
+        ) as Box<_>)
+    }
 }
 
 #[async_trait]
@@ -50,6 +75,18 @@ where
     async fn get_piece(&self, piece_index: PieceIndex) -> anyhow::Result<Option<Piece>> {
         self.as_ref().get_piece(piece_index).await
     }
+
+    async fn get_pieces<'a, PieceIndices>(
+        &'a self,
+        piece_indices: PieceIndices,
+    ) -> anyhow::Result<
+        Box<dyn Stream<Item = (PieceIndex, anyhow::Result<Option<Piece>>)> + Send + Unpin + 'a>,
+    >
+    where
+        PieceIndices: IntoIterator<Item = PieceIndex, IntoIter: Send> + Send + 'a,
+    {
+        self.as_ref().get_pieces(piece_indices).await
+    }
 }
 
 #[async_trait]
@@ -59,6 +96,26 @@ impl PieceGetter for ArchivedHistorySegment {
 
         Ok(self.pieces().nth(position))
     }
+
+    async fn get_pieces<'a, PieceIndices>(
+        &'a self,
+        piece_indices: PieceIndices,
+    ) -> anyhow::Result<
+        Box<dyn Stream<Item = (PieceIndex, anyhow::Result<Option<Piece>>)> + Send + Unpin + 'a>,
+    >
+    where
+        PieceIndices: IntoIterator<Item = PieceIndex, IntoIter: Send> + Send + 'a,
+    {
+        Ok(Box::new(
+            piece_indices
+                .into_iter()
+                .map(|piece_index| async move {
+                    let result = self.get_piece(piece_index).await;
+                    (piece_index, result)
+                })
+                .collect::<FuturesUnordered<_>>(),
+        ) as Box<_>)
+    }
 }
 
 /// Enum to encapsulate the selection between [`ReadAtSync`] and [`ReadAtAsync]` variants