Refactor FlatPieces API to make it nicer to use and helpful in more cases

Author: nazar-pc

crates/sc-consensus-subspace/src/tests.rs

@@ -67,7 +67,7 @@ use std::sync::Arc;
 use std::{cell::RefCell, task::Poll, time::Duration};
 use subspace_archiving::archiver::Archiver;
 use subspace_core_primitives::objects::BlockObjectMapping;
-use subspace_core_primitives::{FlatPieces, LocalChallenge, Piece, Signature, Tag, PIECE_SIZE};
+use subspace_core_primitives::{FlatPieces, LocalChallenge, Piece, Signature, Tag};
 use subspace_solving::{SubspaceCodec, SOLUTION_SIGNING_CONTEXT};
 use substrate_test_runtime::{Block as TestBlock, Hash};
 
@@ -538,7 +538,7 @@ fn run_one_test(mutator: impl Fn(&mut TestHeader, Stage) + Send + Sync + 'static
                 .await
                 .unwrap()
                 .iter()
-                .flat_map(|flat_pieces| flat_pieces.chunks_exact(PIECE_SIZE))
+                .flat_map(|flat_pieces| flat_pieces.as_pieces())
                 .enumerate()
                 .choose(&mut rand::thread_rng())
                 .map(|(piece_index, piece)| (piece_index as u64, Piece::try_from(piece).unwrap()))

crates/subspace-archiving/tests/integration/archiver.rs

@@ -130,7 +130,7 @@ fn archiver() {
             .chain(iter::repeat(block_1.as_ref()).zip(block_1_object_mapping.objects.iter()));
         let piece_objects = first_archived_segment
             .pieces
-            .chunks_exact(PIECE_SIZE)
+            .as_pieces()
             .zip(&first_archived_segment.object_mapping)
             .flat_map(|(piece, object_mapping)| {
                 iter::repeat(piece.as_ref()).zip(&object_mapping.objects)
@@ -140,11 +140,7 @@ fn archiver() {
     }
 
     // Check that all pieces are valid
-    for (position, piece) in first_archived_segment
-        .pieces
-        .chunks_exact(PIECE_SIZE)
-        .enumerate()
-    {
+    for (position, piece) in first_archived_segment.pieces.as_pieces().enumerate() {
         assert!(archiver::is_piece_valid(
             piece,
             first_archived_segment.root_block.records_root(),
@@ -193,7 +189,7 @@ fn archiver() {
             iter::repeat(block_1.as_ref()).zip(block_1_object_mapping.objects.iter().skip(2));
         let piece_objects = archived_segments[0]
             .pieces
-            .chunks_exact(PIECE_SIZE)
+            .as_pieces()
             .zip(&archived_segments[0].object_mapping)
             .flat_map(|(piece, object_mapping)| {
                 iter::repeat(piece.as_ref()).zip(&object_mapping.objects)
@@ -231,7 +227,7 @@ fn archiver() {
             previous_root_block_hash
         );
 
-        for (position, piece) in archived_segment.pieces.chunks_exact(PIECE_SIZE).enumerate() {
+        for (position, piece) in archived_segment.pieces.as_pieces().enumerate() {
             assert!(archiver::is_piece_valid(
                 piece,
                 archived_segment.root_block.records_root(),
@@ -274,7 +270,7 @@ fn archiver() {
         assert_eq!(last_archived_block.number, 3);
         assert_eq!(last_archived_block.partial_archived(), None);
 
-        for (position, piece) in archived_segment.pieces.chunks_exact(PIECE_SIZE).enumerate() {
+        for (position, piece) in archived_segment.pieces.as_pieces().enumerate() {
             assert!(archiver::is_piece_valid(
                 piece,
                 archived_segment.root_block.records_root(),

crates/subspace-archiving/tests/integration/reconstructor.rs

@@ -16,7 +16,7 @@ const SEGMENT_SIZE: usize = RECORD_SIZE * MERKLE_NUM_LEAVES / 2;
 
 fn flat_pieces_to_regular(pieces: &FlatPieces) -> Vec<Piece> {
     pieces
-        .chunks_exact(PIECE_SIZE)
+        .as_pieces()
         .map(|piece| piece.try_into().unwrap())
         .collect()
 }

crates/subspace-core-primitives/src/lib.rs

@@ -256,13 +256,11 @@ impl AsMut<[u8]> for Piece {
 pub struct FlatPieces(Vec<u8>);
 
 impl FlatPieces {
-    /// Create new instance, returns error if `pieces` is not multiple of full pieces.
-    pub fn new(pieces: Vec<u8>) -> Result<Self, Vec<u8>> {
-        if pieces.len() % PIECE_SIZE != 0 {
-            return Err(pieces);
-        }
-
-        Ok(Self(pieces))
+    /// Allocate `FlatPieces` that will hold `piece_count` pieces filled with zeroes.
+    pub fn new(piece_count: usize) -> Self {
+        let mut pieces = Vec::with_capacity(piece_count * PIECE_SIZE);
+        pieces.resize(pieces.capacity(), 0);
+        Self(pieces)
     }
 
     /// Number of pieces contained.
@@ -274,13 +272,27 @@ impl FlatPieces {
     pub fn into_inner(self) -> Vec<u8> {
         self.0
     }
+
+    /// Iterator over individual pieces as byte slices.
+    pub fn as_pieces(&self) -> impl ExactSizeIterator<Item = &[u8]> {
+        self.0.chunks_exact(PIECE_SIZE)
+    }
+
+    /// Iterator over individual pieces as byte slices.
+    pub fn as_pieces_mut(&mut self) -> impl ExactSizeIterator<Item = &mut [u8]> {
+        self.0.chunks_exact_mut(PIECE_SIZE)
+    }
 }
 
 impl TryFrom<Vec<u8>> for FlatPieces {
     type Error = Vec<u8>;
 
     fn try_from(value: Vec<u8>) -> Result<Self, Self::Error> {
-        Self::new(value)
+        if value.len() % PIECE_SIZE != 0 {
+            return Err(value);
+        }
+
+        Ok(Self(value))
     }
 }
 

crates/subspace-solving/tests/integration/codec.rs

@@ -1,5 +1,5 @@
 use std::iter;
-use subspace_core_primitives::PIECE_SIZE;
+use subspace_core_primitives::{FlatPieces, PIECE_SIZE};
 use subspace_solving::SubspaceCodec;
 
 #[test]
@@ -25,10 +25,10 @@ fn batch() {
     // Use 2.5 batches worth of pieces
     let piece_count = subspace_codec.batch_size() * 2 + subspace_codec.batch_size() / 2;
 
-    let mut pieces: Vec<u8> = vec![0u8; PIECE_SIZE * piece_count];
-    pieces.chunks_exact_mut(PIECE_SIZE).for_each(|piece| {
+    let mut pieces = FlatPieces::new(piece_count);
+    for piece in pieces.as_pieces_mut() {
         piece.copy_from_slice(&rand::random::<[u8; PIECE_SIZE]>());
-    });
+    }
     let original_pieces = pieces.clone();
     let piece_indexes: Vec<u64> = iter::repeat_with(|| rand::random())
         .take(piece_count)
@@ -39,8 +39,8 @@ fn batch() {
         .unwrap();
 
     for ((original_piece, piece), piece_index) in original_pieces
-        .chunks_exact(PIECE_SIZE)
-        .zip(pieces.chunks_exact_mut(PIECE_SIZE))
+        .as_pieces()
+        .zip(pieces.as_pieces_mut())
         .zip(piece_indexes)
     {
         assert_ne!(original_piece, piece);