Merge pull request #2556 from subspace/improvements-around-unsafe

Improvements around unsafe

Author: nazar-pc

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

@@ -118,7 +118,10 @@ where
     S: ReadAtSync,
     A: ReadAtAsync,
 {
-    let mut record_chunks = vec![None; Record::NUM_S_BUCKETS];
+    // TODO: Should have been just `::new()`, but https://github.com/rust-lang/rust/issues/53827
+    // SAFETY: Data structure filled with zeroes is a valid invariant
+    let mut record_chunks =
+        unsafe { Box::<[Option<Scalar>; Record::NUM_S_BUCKETS]>::new_zeroed().assume_init() };
 
     let read_chunks_inputs = record_chunks
         .par_iter_mut()
@@ -244,13 +247,6 @@ where
         }
     }
 
-    let mut record_chunks = ManuallyDrop::new(record_chunks);
-
-    // SAFETY: Original memory is not dropped, layout is exactly what we need here
-    let record_chunks = unsafe {
-        Box::from_raw(record_chunks.as_mut_ptr() as *mut [Option<Scalar>; Record::NUM_S_BUCKETS])
-    };
-
     Ok(record_chunks)
 }
 
@@ -261,6 +257,8 @@ pub fn recover_extended_record_chunks(
     erasure_coding: &ErasureCoding,
 ) -> Result<Box<[Scalar; Record::NUM_S_BUCKETS]>, ReadingError> {
     // Restore source record scalars
+    // TODO: Recover into `Box<[Scalar; Record::NUM_S_BUCKETS]>` or else conversion into `Box` below
+    //  might leak memory
     let record_chunks = erasure_coding
         .recover(sector_record_chunks)
         .map_err(|error| ReadingError::FailedToErasureDecodeRecord {
@@ -276,12 +274,12 @@ pub fn recover_extended_record_chunks(
         });
     }
 
+    // Allocation in vector can be larger than contents, we need to make sure allocation is the same
+    // as the contents, this should also contain fast path if allocation matches contents
+    let record_chunks = record_chunks.into_iter().collect::<Box<_>>();
     let mut record_chunks = ManuallyDrop::new(record_chunks);
-
     // SAFETY: Original memory is not dropped, size of the data checked above
-    let record_chunks = unsafe {
-        Box::from_raw(record_chunks.as_mut_ptr() as *mut [Scalar; Record::NUM_S_BUCKETS])
-    };
+    let record_chunks = unsafe { Box::from_raw(record_chunks.as_mut_ptr() as *mut _) };
 
     Ok(record_chunks)
 }

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

@@ -1,6 +1,7 @@
 use bitvec::prelude::*;
 use parity_scale_codec::{Decode, Encode};
 use rayon::prelude::*;
+use std::mem::ManuallyDrop;
 use std::ops::{Deref, DerefMut};
 use std::{mem, slice};
 use subspace_core_primitives::checksum::Blake3Checksummed;
@@ -59,24 +60,24 @@ pub struct SectorMetadata {
 impl SectorMetadata {
     /// Returns offsets of each s-bucket relatively to the beginning of the sector (in chunks)
     pub fn s_bucket_offsets(&self) -> Box<[u32; Record::NUM_S_BUCKETS]> {
-        // TODO: Should have been just `::new()`, but https://github.com/rust-lang/rust/issues/53827
-        // SAFETY: Data structure filled with zeroes is a valid invariant
-        let mut s_bucket_offsets =
-            unsafe { Box::<[u32; Record::NUM_S_BUCKETS]>::new_zeroed().assume_init() };
-
-        self.s_bucket_sizes
+        let s_bucket_offsets = self
+            .s_bucket_sizes
             .iter()
-            .zip(s_bucket_offsets.iter_mut())
-            .for_each({
+            .map({
                 let mut base_offset = 0;
 
-                move |(s_bucket_size, s_bucket_offset)| {
-                    *s_bucket_offset = base_offset;
+                move |s_bucket_size| {
+                    let offset = base_offset;
                     base_offset += u32::from(*s_bucket_size);
+                    offset
                 }
-            });
+            })
+            .collect::<Box<_>>();
 
-        s_bucket_offsets
+        assert_eq!(s_bucket_offsets.len(), Record::NUM_S_BUCKETS);
+        let mut s_bucket_offsets = ManuallyDrop::new(s_bucket_offsets);
+        // SAFETY: Original memory is not dropped, number of elements checked above
+        unsafe { Box::from_raw(s_bucket_offsets.as_mut_ptr() as *mut [u32; Record::NUM_S_BUCKETS]) }
     }
 }
 
@@ -164,6 +165,7 @@ impl RawSector {
 
 // Bit array containing space for bits equal to the number of s-buckets in a record
 type SingleRecordBitArray = BitArray<[u8; Record::NUM_S_BUCKETS / u8::BITS as usize]>;
+
 const SINGLE_RECORD_BIT_ARRAY_SIZE: usize = mem::size_of::<SingleRecordBitArray>();
 
 // TODO: I really tried to avoid `count_ones()`, but wasn't able to with safe Rust due to lifetimes
@@ -400,23 +402,21 @@ impl SectorContentsMap {
 
     /// Returns sizes of each s-bucket
     pub fn s_bucket_sizes(&self) -> Box<[u16; Record::NUM_S_BUCKETS]> {
-        // TODO: Should have been just `::new()`, but https://github.com/rust-lang/rust/issues/53827
-        // SAFETY: Data structure filled with zeroes is a valid invariant
-        let mut s_bucket_sizes =
-            unsafe { Box::<[u16; Record::NUM_S_BUCKETS]>::new_zeroed().assume_init() };
         // Rayon doesn't support iteration over custom types yet
-        (u16::from(SBucket::ZERO)..=u16::from(SBucket::MAX))
+        let s_bucket_sizes = (u16::from(SBucket::ZERO)..=u16::from(SBucket::MAX))
             .into_par_iter()
             .map(SBucket::from)
-            .zip(s_bucket_sizes.par_iter_mut())
-            .for_each(|(s_bucket, s_bucket_size)| {
-                *s_bucket_size = self
-                    .iter_s_bucket_records(s_bucket)
+            .map(|s_bucket| {
+                self.iter_s_bucket_records(s_bucket)
                     .expect("S-bucket guaranteed to be in range; qed")
-                    .count() as u16;
-            });
+                    .count() as u16
+            })
+            .collect::<Box<_>>();
 
-        s_bucket_sizes
+        assert_eq!(s_bucket_sizes.len(), Record::NUM_S_BUCKETS);
+        let mut s_bucket_sizes = ManuallyDrop::new(s_bucket_sizes);
+        // SAFETY: Original memory is not dropped, number of elements checked above
+        unsafe { Box::from_raw(s_bucket_sizes.as_mut_ptr() as *mut [u16; Record::NUM_S_BUCKETS]) }
     }
 
     /// Creates an iterator of `(s_bucket, encoded_chunk_used, chunk_location)`, where `s_bucket` is