ab_farmer_components/
sector.rs

1//! Sector-related data structures
2//!
3//! Sectors and corresponding metadata created by functions in [`plotting`](crate::plotting) module
4//! have a specific structure, represented by data structured in this module.
5//!
6//! It is typically not needed to construct these data structures explicitly outside of this crate,
7//! instead they will be returned as a result of certain operations (like plotting).
8
9use ab_core_primitives::checksum::Blake3Checksummed;
10use ab_core_primitives::hashes::Blake3Hash;
11use ab_core_primitives::pieces::{PieceOffset, Record, RecordChunksRoot, RecordProof, RecordRoot};
12use ab_core_primitives::sectors::{SBucket, SectorIndex};
13use ab_core_primitives::segments::{HistorySize, SegmentIndex};
14use ab_io_type::trivial_type::TrivialType;
15use parity_scale_codec::{Decode, Encode};
16use rayon::prelude::*;
17use std::ops::{Deref, DerefMut};
18use thiserror::Error;
19use tracing::debug;
20
21/// Size of the part of the plot containing record chunks (s-buckets).
22///
23/// Total size of the plot can be computed with [`sector_size()`].
24#[inline]
25pub const fn sector_record_chunks_size(pieces_in_sector: u16) -> usize {
26    pieces_in_sector as usize * Record::SIZE
27}
28
29/// Size of the part of the plot containing record metadata.
30///
31/// Total size of the plot can be computed with [`sector_size()`].
32#[inline]
33pub const fn sector_record_metadata_size(pieces_in_sector: u16) -> usize {
34    pieces_in_sector as usize * RecordMetadata::encoded_size()
35}
36
37/// Exact sector plot size (sector contents map, record chunks, record metadata).
38///
39/// NOTE: Each sector also has corresponding fixed size metadata whose size can be obtained with
40/// [`SectorMetadataChecksummed::encoded_size()`], size of the record chunks (s-buckets) with
41/// [`sector_record_chunks_size()`] and size of record roots and proofs with
42/// [`sector_record_metadata_size()`]. This function just combines those three together for
43/// convenience.
44#[inline]
45pub const fn sector_size(pieces_in_sector: u16) -> usize {
46    sector_record_chunks_size(pieces_in_sector)
47        + sector_record_metadata_size(pieces_in_sector)
48        + SectorContentsMap::encoded_size(pieces_in_sector)
49        + Blake3Hash::SIZE
50}
51
52/// Metadata of the plotted sector
53#[derive(Debug, Encode, Decode, Clone)]
54pub struct SectorMetadata {
55    /// Sector index
56    pub sector_index: SectorIndex,
57    /// Number of pieces stored in this sector
58    pub pieces_in_sector: u16,
59    /// S-bucket sizes in a sector
60    pub s_bucket_sizes: Box<[u16; Record::NUM_S_BUCKETS]>,
61    /// Size of the blockchain history at time of sector creation
62    pub history_size: HistorySize,
63}
64
65impl SectorMetadata {
66    /// Returns offsets of each s-bucket relatively to the beginning of the sector (in chunks)
67    pub fn s_bucket_offsets(&self) -> Box<[u32; Record::NUM_S_BUCKETS]> {
68        let s_bucket_offsets = self
69            .s_bucket_sizes
70            .iter()
71            .map({
72                let mut base_offset = 0;
73
74                move |s_bucket_size| {
75                    let offset = base_offset;
76                    base_offset += u32::from(*s_bucket_size);
77                    offset
78                }
79            })
80            .collect::<Box<_>>();
81
82        assert_eq!(s_bucket_offsets.len(), Record::NUM_S_BUCKETS);
83        // SAFETY: Number of elements checked above
84        unsafe {
85            Box::from_raw(Box::into_raw(s_bucket_offsets).cast::<[u32; Record::NUM_S_BUCKETS]>())
86        }
87    }
88}
89
90/// Same as [`SectorMetadata`], but with checksums verified during SCALE encoding/decoding
91#[derive(Debug, Clone, Encode, Decode)]
92pub struct SectorMetadataChecksummed(Blake3Checksummed<SectorMetadata>);
93
94impl From<SectorMetadata> for SectorMetadataChecksummed {
95    #[inline]
96    fn from(value: SectorMetadata) -> Self {
97        Self(Blake3Checksummed(value))
98    }
99}
100
101impl Deref for SectorMetadataChecksummed {
102    type Target = SectorMetadata;
103
104    #[inline]
105    fn deref(&self) -> &Self::Target {
106        &self.0.0
107    }
108}
109
110impl DerefMut for SectorMetadataChecksummed {
111    #[inline]
112    fn deref_mut(&mut self) -> &mut Self::Target {
113        &mut self.0.0
114    }
115}
116
117impl SectorMetadataChecksummed {
118    /// Size of encoded checksummed sector metadata.
119    ///
120    /// For sector plot size use [`sector_size()`].
121    #[inline]
122    pub fn encoded_size() -> usize {
123        let default = SectorMetadataChecksummed::from(SectorMetadata {
124            sector_index: SectorIndex::ZERO,
125            pieces_in_sector: 0,
126            // TODO: Should have been just `::new()`, but https://github.com/rust-lang/rust/issues/53827
127            // SAFETY: Data structure filled with zeroes is a valid invariant
128            s_bucket_sizes: unsafe { Box::new_zeroed().assume_init() },
129            history_size: HistorySize::from(SegmentIndex::ZERO),
130        });
131
132        default.encoded_size()
133    }
134}
135
136/// Root and proof corresponding to the same record
137#[derive(Debug, Default, Clone, Encode, Decode)]
138pub(crate) struct RecordMetadata {
139    /// Record root
140    pub(crate) root: RecordRoot,
141    /// Parity chunks root
142    pub(crate) parity_chunks_root: RecordChunksRoot,
143    /// Record proof
144    pub(crate) proof: RecordProof,
145    /// Checksum (hash) of the whole piece
146    pub(crate) piece_checksum: Blake3Hash,
147}
148
149impl RecordMetadata {
150    pub(crate) const fn encoded_size() -> usize {
151        RecordProof::SIZE + RecordRoot::SIZE + RecordChunksRoot::SIZE + Blake3Hash::SIZE
152    }
153}
154
155/// Raw sector before it is transformed and written to plot, used during plotting
156#[derive(Debug, Clone)]
157pub(crate) struct RawSector {
158    /// List of records, likely downloaded from the network
159    pub(crate) records: Vec<Record>,
160    /// Metadata (root and proof) corresponding to the same record
161    pub(crate) metadata: Vec<RecordMetadata>,
162}
163
164impl RawSector {
165    /// Create new raw sector with internal vectors being allocated and filled with default values
166    pub(crate) fn new(pieces_in_sector: u16) -> Self {
167        Self {
168            records: Record::new_zero_vec(usize::from(pieces_in_sector)),
169            metadata: vec![RecordMetadata::default(); usize::from(pieces_in_sector)],
170        }
171    }
172}
173
174/// S-buckets at which proofs were found.
175///
176/// S-buckets are grouped by 8, within each `u8` bits right to left (LSB) indicate the presence
177/// of a proof for corresponding s-bucket, so that the whole array of bytes can be thought as a
178/// large set of bits.
179///
180/// There will be at most [`Record::NUM_CHUNKS`] proofs produced/bits set to `1`.
181pub type FoundProofs = [u8; Record::NUM_S_BUCKETS / u8::BITS as usize];
182
183/// Error happening when trying to create [`SectorContentsMap`] from bytes
184#[derive(Debug, Error, Copy, Clone, Eq, PartialEq)]
185pub enum SectorContentsMapFromBytesError {
186    /// Invalid bytes length
187    #[error("Invalid bytes length, expected {expected}, actual {actual}")]
188    InvalidBytesLength {
189        /// Expected length
190        expected: usize,
191        /// Actual length
192        actual: usize,
193    },
194    /// Checksum mismatch
195    #[error("Checksum mismatch")]
196    ChecksumMismatch,
197}
198
199/// Error happening when trying to encode [`SectorContentsMap`] into bytes
200#[derive(Debug, Error, Copy, Clone, Eq, PartialEq)]
201pub enum SectorContentsMapEncodeIntoError {
202    /// Invalid bytes length
203    #[error("Invalid bytes length, expected {expected}, actual {actual}")]
204    InvalidBytesLength {
205        /// Expected length
206        expected: usize,
207        /// Actual length
208        actual: usize,
209    },
210}
211
212/// Error happening when trying to create [`SectorContentsMap`] from bytes
213#[derive(Debug, Error, Copy, Clone, Eq, PartialEq)]
214pub enum SectorContentsMapIterationError {
215    /// S-bucket provided is out of range
216    #[error("S-bucket provided {provided} is out of range, max {max}")]
217    SBucketOutOfRange {
218        /// Provided s-bucket
219        provided: usize,
220        /// Max s-bucket
221        max: usize,
222    },
223}
224
225/// Map of sector contents.
226///
227/// Abstraction on top of bitfields that allow making sense of sector contents that contain
228/// encoded (meaning erasure coded and encoded with existing PoSpace proof) chunks used at the same
229/// time both in records (before writing to plot) and s-buckets (written into the plot) format
230#[derive(Debug, Clone, Eq, PartialEq)]
231pub struct SectorContentsMap {
232    /// Bitfields for each record, each bit is `true` if record chunk at corresponding position was
233    /// used
234    record_chunks_used: Vec<FoundProofs>,
235}
236
237impl SectorContentsMap {
238    /// Create new sector contents map initialized with zeroes to store data for `pieces_in_sector`
239    /// records
240    pub fn new(pieces_in_sector: u16) -> Self {
241        Self {
242            record_chunks_used: vec![[0; _]; usize::from(pieces_in_sector)],
243        }
244    }
245
246    /// Reconstruct sector contents map from bytes.
247    ///
248    /// Returns error if length of the vector doesn't match [`Self::encoded_size()`] for
249    /// `pieces_in_sector`.
250    pub fn from_bytes(
251        bytes: &[u8],
252        pieces_in_sector: u16,
253    ) -> Result<Self, SectorContentsMapFromBytesError> {
254        if bytes.len() != Self::encoded_size(pieces_in_sector) {
255            return Err(SectorContentsMapFromBytesError::InvalidBytesLength {
256                expected: Self::encoded_size(pieces_in_sector),
257                actual: bytes.len(),
258            });
259        }
260
261        let (single_records_bit_arrays, expected_checksum) =
262            bytes.split_at(bytes.len() - Blake3Hash::SIZE);
263        // SAFETY: All bit patterns are valid
264        let expected_checksum = unsafe {
265            Blake3Hash::from_bytes(expected_checksum).expect("No alignment requirements; qed")
266        };
267        // Verify checksum
268        let actual_checksum = Blake3Hash::from(blake3::hash(single_records_bit_arrays));
269        if &actual_checksum != expected_checksum {
270            debug!(
271                %actual_checksum,
272                %expected_checksum,
273                "Hash doesn't match, corrupted bytes"
274            );
275
276            return Err(SectorContentsMapFromBytesError::ChecksumMismatch);
277        }
278
279        let mut record_chunks_used = vec![[0; _]; pieces_in_sector.into()];
280
281        record_chunks_used
282            .as_flattened_mut()
283            .copy_from_slice(single_records_bit_arrays);
284
285        Ok(Self { record_chunks_used })
286    }
287
288    /// Size of sector contents map when encoded and stored in the plot for specified number of
289    /// pieces in sector
290    pub const fn encoded_size(pieces_in_sector: u16) -> usize {
291        size_of::<FoundProofs>() * pieces_in_sector as usize + Blake3Hash::SIZE
292    }
293
294    /// Encode internal contents into `output`
295    pub fn encode_into(&self, output: &mut [u8]) -> Result<(), SectorContentsMapEncodeIntoError> {
296        if output.len() != Self::encoded_size(self.record_chunks_used.len() as u16) {
297            return Err(SectorContentsMapEncodeIntoError::InvalidBytesLength {
298                expected: Self::encoded_size(self.record_chunks_used.len() as u16),
299                actual: output.len(),
300            });
301        }
302
303        let slice = self.record_chunks_used.as_flattened();
304        // Write data and checksum
305        output[..slice.len()].copy_from_slice(slice);
306        output[slice.len()..].copy_from_slice(blake3::hash(slice).as_bytes());
307
308        Ok(())
309    }
310
311    /// Iterate over individual record chunks (s-buckets) that were used
312    pub fn iter_record_chunks_used(&self) -> &[FoundProofs] {
313        &self.record_chunks_used
314    }
315
316    /// Iterate mutably over individual record chunks (s-buckets) that were used
317    pub fn iter_record_chunks_used_mut(&mut self) -> &mut [FoundProofs] {
318        &mut self.record_chunks_used
319    }
320
321    /// Returns sizes of each s-bucket
322    pub fn s_bucket_sizes(&self) -> Box<[u16; Record::NUM_S_BUCKETS]> {
323        // Rayon doesn't support iteration over custom types yet
324        let s_bucket_sizes = (u16::from(SBucket::ZERO)..=u16::from(SBucket::MAX))
325            .into_par_iter()
326            .map(SBucket::from)
327            .map(|s_bucket| {
328                self.iter_s_bucket_piece_offsets(s_bucket)
329                    .expect("S-bucket guaranteed to be in range; qed")
330                    .count() as u16
331            })
332            .collect::<Box<_>>();
333
334        assert_eq!(s_bucket_sizes.len(), Record::NUM_S_BUCKETS);
335
336        // SAFETY: Number of elements checked above
337        unsafe {
338            Box::from_raw(Box::into_raw(s_bucket_sizes).cast::<[u16; Record::NUM_S_BUCKETS]>())
339        }
340    }
341
342    /// Creates an iterator of `(s_bucket, chunk_location)`, where `s_bucket` is the position of the
343    /// chunk in the erasure coded record and `chunk_location` is the offset of the chunk in the
344    /// plot (across all s-buckets).
345    pub fn iter_record_chunk_to_plot(
346        &self,
347        piece_offset: PieceOffset,
348    ) -> impl Iterator<Item = (SBucket, usize)> + '_ {
349        // Iterate over all s-buckets
350        (SBucket::ZERO..=SBucket::MAX)
351            // In each s-bucket map all records used
352            .flat_map(|s_bucket| {
353                self.iter_s_bucket_piece_offsets(s_bucket)
354                    .expect("S-bucket guaranteed to be in range; qed")
355                    .map(move |current_piece_offset| (s_bucket, current_piece_offset))
356            })
357            // We've got contents of all s-buckets in a flat iterator, enumerating them so it is
358            // possible to find in the plot later if desired
359            .enumerate()
360            // Everything about the piece offset we care about
361            .filter_map(move |(chunk_location, (s_bucket, current_piece_offset))| {
362                // In case record for `piece_offset` is found, return necessary information
363                (current_piece_offset == piece_offset).then_some((s_bucket, chunk_location))
364            })
365            // Tiny optimization in case we have found chunks for all records already
366            .take(Record::NUM_CHUNKS)
367    }
368
369    /// Creates an iterator of `Option<chunk_offset>`, where each entry corresponds
370    /// s-bucket/position of the chunk in the erasure coded record, `chunk_offset` is the offset of
371    /// the chunk in the corresponding s-bucket.
372    ///
373    /// Similar to `Self::iter_record_chunk_to_plot()`, but runs in parallel, returns entries for
374    /// all s-buckets and offsets are within corresponding s-buckets rather than the whole plot.
375    pub fn par_iter_record_chunk_to_plot(
376        &self,
377        piece_offset: PieceOffset,
378    ) -> impl IndexedParallelIterator<Item = Option<usize>> + '_ {
379        let piece_offset = usize::from(piece_offset);
380        (u16::from(SBucket::ZERO)..=u16::from(SBucket::MAX))
381            .into_par_iter()
382            .map(SBucket::from)
383            // In each s-bucket map all records used
384            .map(move |s_bucket| {
385                let byte_offset = usize::from(s_bucket) / u8::BITS as usize;
386                let bit_mask = 1 << (usize::from(s_bucket) % u8::BITS as usize);
387
388                if self.record_chunks_used[piece_offset][byte_offset] & bit_mask == 0 {
389                    return None;
390                }
391
392                // How many other record chunks we have in s-bucket before piece offset we care
393                // about
394                let chunk_offset = self
395                    .record_chunks_used
396                    .iter()
397                    .take(piece_offset)
398                    .filter(move |record_chunks_used| {
399                        record_chunks_used[byte_offset] & bit_mask != 0
400                    })
401                    .count();
402
403                Some(chunk_offset)
404            })
405    }
406
407    /// Creates an iterator of piece offsets to which corresponding chunks belong.
408    ///
409    /// Returns error if `s_bucket` is outside of [`Record::NUM_S_BUCKETS`] range.
410    pub fn iter_s_bucket_piece_offsets(
411        &self,
412        s_bucket: SBucket,
413    ) -> Result<impl Iterator<Item = PieceOffset> + '_, SectorContentsMapIterationError> {
414        let s_bucket = usize::from(s_bucket);
415
416        if s_bucket >= Record::NUM_S_BUCKETS {
417            return Err(SectorContentsMapIterationError::SBucketOutOfRange {
418                provided: s_bucket,
419                max: Record::NUM_S_BUCKETS,
420            });
421        }
422
423        Ok((PieceOffset::ZERO..)
424            .zip(&self.record_chunks_used)
425            .filter_map(move |(piece_offset, record_chunks_used)| {
426                let byte_offset = s_bucket / u8::BITS as usize;
427                let bit_mask = 1 << (s_bucket % u8::BITS as usize);
428
429                (record_chunks_used[byte_offset] & bit_mask != 0).then_some(piece_offset)
430            }))
431    }
432
433    /// Iterate over chunks of s-bucket indicating if record chunk is used at corresponding
434    /// position.
435    ///
436    /// ## Panics
437    /// Panics if `s_bucket` is outside of [`Record::NUM_S_BUCKETS`] range.
438    pub fn iter_s_bucket_used_record_chunks_used(
439        &self,
440        s_bucket: SBucket,
441    ) -> Result<impl Iterator<Item = bool> + '_, SectorContentsMapIterationError> {
442        let s_bucket = usize::from(s_bucket);
443
444        if s_bucket >= Record::NUM_S_BUCKETS {
445            return Err(SectorContentsMapIterationError::SBucketOutOfRange {
446                provided: s_bucket,
447                max: Record::NUM_S_BUCKETS,
448            });
449        }
450
451        Ok(self
452            .record_chunks_used
453            .iter()
454            .map(move |record_chunks_used| {
455                let byte_offset = s_bucket / u8::BITS as usize;
456                let bit_mask = 1 << (s_bucket % u8::BITS as usize);
457
458                record_chunks_used[byte_offset] & bit_mask != 0
459            }))
460    }
461}