From 9e3812c19fd656e46bbff5dea4d52eeb07f5d3da Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Tue, 10 Sep 2024 10:13:00 +1000
Subject: [PATCH] impl ObjectPieceGetter for NewArchivedSegment and
 Vec<(PieceIndex, Piece)>

---
 .../src/object_fetcher.rs                     | 35 +++++++++++++++++--
 1 file changed, 32 insertions(+), 3 deletions(-)

diff --git a/shared/subspace-object-fetching/src/object_fetcher.rs b/shared/subspace-object-fetching/src/object_fetcher.rs
index e409e72298..670b1cabb4 100644
--- a/shared/subspace-object-fetching/src/object_fetcher.rs
+++ b/shared/subspace-object-fetching/src/object_fetcher.rs
@@ -19,7 +19,7 @@ use async_trait::async_trait;
 use parity_scale_codec::{Compact, CompactLen, Decode, Encode};
 use std::fmt;
 use std::sync::Arc;
-use subspace_archiving::archiver::{Segment, SegmentItem};
+use subspace_archiving::archiver::{NewArchivedSegment, Segment, SegmentItem};
 use subspace_core_primitives::objects::GlobalObject;
 use subspace_core_primitives::{
     ArchivedHistorySegment, Piece, PieceIndex, RawRecord, RecordedHistorySegment, SegmentIndex,
@@ -55,7 +55,7 @@ where
     PG: ObjectPieceGetter + Sync,
 {
     async fn get_piece(&self, piece_index: PieceIndex) -> Option<Piece> {
-        for piece_getter in self.iter() {
+        for piece_getter in self {
             if let Some(piece) = piece_getter.get_piece(piece_index).await {
                 return Some(piece);
             }
@@ -80,12 +80,41 @@ impl<PV> ObjectPieceGetter for PieceProvider<PV>
 where
     PV: PieceValidator,
 {
+    // TODO: this doesn't seem to return any pieces (on a single node launched with --dev)
     async fn get_piece(&self, piece_index: PieceIndex) -> Option<Piece> {
         if let Some(piece) = self.get_piece_from_cache(piece_index).await {
             return Some(piece);
         }
 
-        self.get_piece_from_archival_storage(piece_index, MAX_RANDOM_WALK_ROUNDS).await
+        self.get_piece_from_archival_storage(piece_index, MAX_RANDOM_WALK_ROUNDS)
+            .await
+    }
+}
+
+// Convenience methods, mainly used in testing
+#[async_trait]
+impl ObjectPieceGetter for NewArchivedSegment {
+    async fn get_piece(&self, piece_index: PieceIndex) -> Option<Piece> {
+        if piece_index.segment_index() == self.segment_header.segment_index() {
+            return self
+                .pieces
+                .pieces()
+                .nth(piece_index.position() as usize)
+                .map(Into::into);
+        }
+
+        None
+    }
+}
+
+#[async_trait]
+impl ObjectPieceGetter for (PieceIndex, Piece) {
+    async fn get_piece(&self, piece_index: PieceIndex) -> Option<Piece> {
+        if self.0 == piece_index {
+            return Some(self.1.clone());
+        }
+
+        None
     }
 }