Optimize allocations and memory copying in subspace-archiving with special wrapper and a bit of unsafe

Author: nazar-pc

crates/subspace-archiving/src/archiver.rs

@@ -12,17 +12,19 @@
 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 // See the License for the specific language governing permissions and
 // limitations under the License.
+
+mod record_shards;
+
 extern crate alloc;
 
+use crate::archiver::record_shards::RecordShards;
 use crate::merkle_tree::{MerkleTree, Witness};
-use crate::utils;
-use crate::utils::{Gf16Element, GF_16_ELEMENT_BYTES};
+use crate::utils::GF_16_ELEMENT_BYTES;
 use alloc::borrow::Cow;
 use alloc::collections::VecDeque;
 use alloc::vec;
 use alloc::vec::Vec;
 use core::cmp::Ordering;
-use core::iter;
 use parity_scale_codec::{Compact, CompactLen, Decode, Encode};
 use reed_solomon_erasure::galois_16::ReedSolomon;
 use subspace_core_primitives::objects::{
@@ -135,6 +137,9 @@ pub enum ArchiverInstantiationError {
         error("Segment size is not bigger than record size")
     )]
     SegmentSizeTooSmall,
+    /// Segment size must be multiple of two
+    #[cfg_attr(feature = "thiserror", error("Segment size must be multiple of two"))]
+    SegmentSizeNotMultipleOfTwo,
     /// Segment size is not a multiple of record size
     #[cfg_attr(
         feature = "thiserror",
@@ -220,6 +225,9 @@ impl Archiver {
         if segment_size <= record_size {
             return Err(ArchiverInstantiationError::SegmentSizeTooSmall);
         }
+        if segment_size % GF_16_ELEMENT_BYTES != 0 {
+            return Err(ArchiverInstantiationError::SegmentSizeNotMultipleOfTwo);
+        }
         if segment_size % record_size != 0 {
             return Err(ArchiverInstantiationError::SegmentSizesNotMultipleOfRecordSize);
         }
@@ -611,64 +619,54 @@ impl Archiver {
             }
             corrected_object_mapping
         };
-        let segment = {
-            let mut segment = segment.encode();
-            // Add a few bytes of padding if needed to get constant size (caused by compact length
-            // encoding of scale codec)
-            assert!(self.segment_size >= segment.len());
-            segment.resize(self.segment_size, 0u8);
-            segment
+        let mut record_shards = {
+            // Allocate record shards that can hold both data and parity record shards (hence `*2`)
+            let mut record_shards =
+                RecordShards::new(self.segment_size / self.record_size * 2, self.record_size);
+            segment.encode_to(&mut record_shards);
+            record_shards
         };
 
-        let data_shards: Vec<Vec<Gf16Element>> = segment
-            .chunks_exact(self.record_size)
-            .map(utils::slice_to_arrays)
-            .collect();
-
         drop(segment);
 
-        let mut parity_shards: Vec<Vec<Gf16Element>> =
-            iter::repeat(vec![Gf16Element::default(); self.record_size / 2])
-                .take(data_shards.len())
-                .collect();
+        let mut record_shards_slices = record_shards.as_mut_slices();
 
         // Apply erasure coding to to create parity shards/records
         self.reed_solomon
-            .encode_sep(&data_shards, &mut parity_shards)
-            .expect("Encoding is running with fixed parameters and should never fail");
+            .encode(&mut record_shards_slices)
+            .expect("Encoding is running with fixed parameters and should never fail; qed");
 
-        let mut pieces = vec![0u8; (data_shards.len() + parity_shards.len()) * PIECE_SIZE];
-        // Combine data and parity records back into flat vector of pieces (witness part of piece is
-        // still zeroes after this and is filled later)
-        pieces
-            .chunks_exact_mut(PIECE_SIZE)
-            .zip(data_shards.into_iter().chain(parity_shards))
-            .flat_map(|(piece, shard)| {
-                piece
-                    .chunks_exact_mut(GF_16_ELEMENT_BYTES)
-                    .zip(shard.into_iter())
-            })
-            .for_each(|(piece_chunk, shard_chunk)| {
-                piece_chunk.copy_from_slice(&shard_chunk);
-            });
+        let mut pieces = FlatPieces::new(record_shards_slices.len());
+        drop(record_shards_slices);
 
         // Build a Merkle tree over all records
         let merkle_tree = MerkleTree::from_data(
-            pieces
-                .chunks_exact(PIECE_SIZE)
-                .map(|piece| &piece[..self.record_size]),
+            record_shards
+                .as_bytes()
+                .as_ref()
+                .chunks_exact(self.record_size),
         );
 
-        // Fill witnesses (Merkle proofs) in pieces created earlier
+        // Combine data and parity records back into flat vector of pieces along with corresponding
+        // witnesses (Merkle proofs) created above.
         pieces
-            .chunks_exact_mut(PIECE_SIZE)
+            .as_pieces_mut()
             .enumerate()
-            .for_each(|(position, piece)| {
-                let witness = merkle_tree
-                    .get_witness(position)
-                    .expect("We use the same indexes as during Merkle tree creation; qed");
-
-                piece[self.record_size..].copy_from_slice(&witness);
+            .zip(
+                record_shards
+                    .as_bytes()
+                    .as_ref()
+                    .chunks_exact(self.record_size),
+            )
+            .for_each(|((position, piece), shard_chunk)| {
+                let (record_part, witness_part) = piece.split_at_mut(self.record_size);
+
+                record_part.copy_from_slice(shard_chunk);
+                witness_part.copy_from_slice(
+                    &merkle_tree
+                        .get_witness(position)
+                        .expect("We use the same indexes as during Merkle tree creation; qed"),
+                );
             });
 
         // Now produce root block
@@ -689,9 +687,7 @@ impl Archiver {
 
         ArchivedSegment {
             root_block,
-            pieces: pieces
-                .try_into()
-                .expect("Pieces length is correct as created above; qed"),
+            pieces,
             object_mapping,
         }
     }

crates/subspace-archiving/src/archiver/record_shards.rs

@@ -0,0 +1,89 @@
+use crate::utils::{Gf16Element, GF_16_ELEMENT_BYTES};
+use std::io;
+use std::mem::ManuallyDrop;
+
+/// Wrapper data structure for record shards that correspond to the same recorded history segment
+/// for more convenient management.
+///
+/// Allows to accessing underlying data both as list of shards for erasure coding and regular slice
+/// of bytes for other purposes, also implements [`std::io::Write`] so that it can be used with
+/// `parity-scale-codec` to write encoded data right into [`RecordShards`].
+pub(super) struct RecordShards {
+    shards: Vec<Gf16Element>,
+    cursor: usize,
+    /// Shard size in bytes
+    shard_size: usize,
+}
+
+impl io::Write for RecordShards {
+    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
+        // SAFETY:
+        // Lifetime is the same as source, de-allocation is prevented with `ManuallyDrop`, no
+        // re-allocation will happen so interpreting one vector as different vector is fine.
+        unsafe {
+            let mut target = ManuallyDrop::new(Vec::from_raw_parts(
+                self.shards.as_mut_ptr() as *mut u8,
+                self.shards.len() * GF_16_ELEMENT_BYTES,
+                self.shards.len() * GF_16_ELEMENT_BYTES,
+            ));
+
+            let written_bytes = target[self.cursor..].as_mut().write(buf)?;
+
+            self.cursor += written_bytes;
+
+            Ok(written_bytes)
+        }
+    }
+
+    fn flush(&mut self) -> io::Result<()> {
+        Ok(())
+    }
+}
+
+impl RecordShards {
+    /// Create new `Shards` struct for specified number of shards and shard size in bytes.
+    ///
+    /// Panics if shard size is not multiple of 2.
+    pub(super) fn new(number_of_shards: usize, shard_size: usize) -> Self {
+        assert_eq!(shard_size % GF_16_ELEMENT_BYTES, 0);
+
+        let mut shards = Vec::with_capacity(number_of_shards * shard_size / GF_16_ELEMENT_BYTES);
+        shards.resize(shards.capacity(), Gf16Element::default());
+
+        Self {
+            shards,
+            cursor: 0,
+            shard_size,
+        }
+    }
+
+    /// Access internal record shards as contiguous memory slice.
+    pub(super) fn as_bytes(&mut self) -> impl AsRef<[u8]> + '_ {
+        // SAFETY:
+        // Returned lifetime is the same as source, de-allocation is prevented with `ManuallyDrop`
+        // and specific type is erased from API, so nothing except reading bytes is possible.
+        unsafe {
+            /// Private wrapper just so that `AsRef<[u8]>` can be implemented on it.
+            struct AsRefRecordShards(ManuallyDrop<Vec<u8>>);
+
+            impl AsRef<[u8]> for AsRefRecordShards {
+                fn as_ref(&self) -> &[u8] {
+                    &self.0
+                }
+            }
+
+            AsRefRecordShards(ManuallyDrop::new(Vec::from_raw_parts(
+                self.shards.as_mut_ptr() as *mut u8,
+                self.shards.len() * GF_16_ELEMENT_BYTES,
+                self.shards.len() * GF_16_ELEMENT_BYTES,
+            )))
+        }
+    }
+
+    /// Access internal record shards as vector of mutable shards.
+    pub(super) fn as_mut_slices(&mut self) -> Vec<&mut [Gf16Element]> {
+        self.shards
+            .chunks_exact_mut(self.shard_size / GF_16_ELEMENT_BYTES)
+            .collect()
+    }
+}