ab_archiving/

reconstructor.rs

use crate::archiver::{Segment, SegmentItem};
use ab_core_primitives::block::BlockNumber;
use ab_core_primitives::pieces::Piece;
use ab_core_primitives::segments::{
    ArchivedHistorySegment, LastArchivedBlock, LocalSegmentIndex, RecordedHistorySegment,
    SegmentHeader,
};
use ab_erasure_coding::{ErasureCoding, ErasureCodingError, ShardsBitmap};
use alloc::vec::Vec;
use core::{array, mem};
use parity_scale_codec::Decode;

/// Reconstructor-related instantiation error
#[derive(Debug, Clone, PartialEq, thiserror::Error)]
pub enum ReconstructorError {
    /// Error during data shards reconstruction
    #[error("Error during data shards reconstruction: {0}")]
    DataShardsReconstruction(#[from] ErasureCodingError),
    /// Not enough shards
    #[error("Not enough shards: {num_shards}")]
    NotEnoughShards { num_shards: usize },
    /// Segment size is not bigger than record size
    #[error("Error during segment decoding: {0}")]
    SegmentDecoding(parity_scale_codec::Error),
    /// Incorrect segment order, each next segment must directly follow the previous one
    #[error(
        "Incorrect segment order, expected index {expected_segment_index}, actual \
        {actual_segment_index}"
    )]
    IncorrectSegmentOrder {
        expected_segment_index: LocalSegmentIndex,
        actual_segment_index: LocalSegmentIndex,
    },
    /// Parent segment header is not the first item of the segment
    #[error("Parent segment header is not the first item of the segment")]
    UnexpectedParentSegmentHeader,
    /// Segment index overflow, no segment can follow the segment with the maximum index
    #[error("Segment index overflow, no segment can follow the segment with the maximum index")]
    SegmentIndexOverflow,
    /// Block number overflow
    #[error("Block number overflow")]
    BlockNumberOverflow,
}

/// Data structure that contains information reconstructed from given segment (potentially using
/// information from segments that were added previously)
#[derive(Debug, Default, Clone, Eq, PartialEq)]
pub struct ReconstructedContents {
    /// Segment header stored in a segment
    pub segment_header: Option<SegmentHeader>,
    /// Reconstructed encoded blocks with their block numbers
    pub blocks: Vec<(BlockNumber, Vec<u8>)>,
}

/// Reconstructor helps to retrieve blocks from archived pieces.
#[derive(Debug, Clone)]
pub struct Reconstructor {
    /// Erasure coding data structure
    erasure_coding: ErasureCoding,
    /// Index of the last segment added to reconstructor
    last_segment_index: Option<LocalSegmentIndex>,
    /// Partially reconstructed block waiting for more data, empty if there is none
    partial_block: Vec<u8>,
}

impl Reconstructor {
    /// Create a new instance
    pub fn new(erasure_coding: ErasureCoding) -> Self {
        Self {
            erasure_coding,
            last_segment_index: None,
            partial_block: Vec::new(),
        }
    }

    /// Given a set of pieces of a segment of the archived history (any half of all pieces are
    /// required to be present, the rest will be recovered automatically due to use of erasure
    /// coding if needed), reconstructs and returns the segment itself.
    ///
    /// Does not modify the internal state of the reconstructor.
    pub fn reconstruct_segment(
        &self,
        // TODO: Improve API to not use `Option` anymore
        segment_pieces: &[Option<Piece>],
    ) -> Result<Segment, ReconstructorError> {
        if segment_pieces.len() < ArchivedHistorySegment::NUM_PIECES {
            return Err(ReconstructorError::NotEnoughShards {
                num_shards: segment_pieces.len(),
            });
        }
        let mut segment_data = RecordedHistorySegment::new_boxed();

        if !segment_pieces
            .iter()
            .zip(segment_data.iter_mut())
            .all(|(maybe_piece, record)| {
                if let Some(piece) = maybe_piece {
                    // Fancy way to insert value to avoid going through stack (if naive
                    // dereferencing is used) and potentially causing stack overflow as the result
                    record.copy_from_slice(&*piece.record);
                    true
                } else {
                    false
                }
            })
        {
            // If not all data pieces are available, need to reconstruct data shards using erasure
            // coding
            let (source_segment_pieces, parity_segment_pieces) =
                segment_pieces.split_at(RecordedHistorySegment::NUM_RAW_RECORDS);

            let mut source_present = ShardsBitmap::none();
            for (index, (output_record, maybe_source_piece)) in segment_data
                .iter_mut()
                .zip(source_segment_pieces)
                .enumerate()
            {
                if let Some(input_piece) = maybe_source_piece {
                    // Fancy way to insert value to avoid going through stack (if naive
                    // dereferencing is used) and potentially causing stack overflow as the result
                    output_record.copy_from_slice(&*input_piece.record);
                    source_present.set(index);
                }
            }

            // Parity records are read-only here, so missing ones simply have no memory
            let parity_records = array::from_fn(|index| {
                parity_segment_pieces[index]
                    .as_ref()
                    .map(|input_piece| &input_piece.record)
            });

            #[cfg(not(feature = "parallel"))]
            let recover_source = ErasureCoding::recover_source_scattered;
            #[cfg(feature = "parallel")]
            let recover_source = ErasureCoding::par_recover_source_scattered;

            recover_source(
                &self.erasure_coding,
                {
                    let records: &mut [_; RecordedHistorySegment::NUM_RAW_RECORDS] =
                        segment_data.as_mut();
                    let mut records = records.iter_mut();

                    array::from_fn::<_, { RecordedHistorySegment::NUM_RAW_RECORDS }, _>(|_| {
                        records
                            .next()
                            .expect("Number of records matches the array size; qed")
                    })
                },
                &source_present,
                parity_records,
            )?;
        }

        let segment = Segment::decode(&mut AsRef::<[u8]>::as_ref(segment_data.as_ref()))
            .map_err(ReconstructorError::SegmentDecoding)?;

        Ok(segment)
    }

    /// Given a set of pieces of a segment of the archived history (any half of all pieces are
    /// required to be present, the rest will be recovered automatically due to use of erasure
    /// coding if needed), reconstructs and returns segment header and a list of encoded blocks with
    /// corresponding block numbers.
    ///
    /// The reconstructor is consumed and returned together with reconstructed contents if the
    /// segment was added successfully. It is possible to start with any segment, but when next
    /// segment is pushed, it needs to follow the previous one or else error will be returned.
    pub fn add_segment(
        mut self,
        segment_pieces: &[Option<Piece>],
    ) -> Result<(Self, ReconstructedContents), ReconstructorError> {
        let segment = self.reconstruct_segment(segment_pieces)?;
        let mut segment_items = segment.items.into_iter().peekable();

        let mut reconstructed_contents = ReconstructedContents::default();
        let mut next_block_number = BlockNumber::ZERO;

        // Only the first segment has no parent segment header, all other segments start with it
        let segment_index = if let Some(&SegmentItem::ParentSegmentHeader(parent_segment_header)) =
            segment_items.peek()
        {
            segment_items.next();

            let LastArchivedBlock {
                number,
                archived_progress,
            } = parent_segment_header.last_archived_block;

            reconstructed_contents
                .segment_header
                .replace(parent_segment_header);

            match archived_progress.partial() {
                None => {
                    // The last archived block ended exactly with the previous segment, so if its
                    // beginning was reconstructed from previous segments, it is complete now
                    if !self.partial_block.is_empty() {
                        reconstructed_contents
                            .blocks
                            .push((number.as_inner(), mem::take(&mut self.partial_block)));
                    }

                    next_block_number = number
                        .as_inner()
                        .checked_add(BlockNumber::ONE)
                        .ok_or(ReconstructorError::BlockNumberOverflow)?;
                }
                Some(_bytes) => {
                    next_block_number = number.as_inner();

                    if self.partial_block.is_empty() {
                        // Will not be able to recover full block, bump right away
                        next_block_number = next_block_number
                            .checked_add(BlockNumber::ONE)
                            .ok_or(ReconstructorError::BlockNumberOverflow)?;
                    }
                }
            }

            parent_segment_header
                .index
                .as_inner()
                .checked_add(LocalSegmentIndex::ONE)
                .ok_or(ReconstructorError::SegmentIndexOverflow)?
        } else {
            LocalSegmentIndex::ZERO
        };

        if let Some(last_segment_index) = self.last_segment_index {
            let expected_segment_index = last_segment_index
                .checked_add(LocalSegmentIndex::ONE)
                .ok_or(ReconstructorError::SegmentIndexOverflow)?;

            if segment_index != expected_segment_index {
                return Err(ReconstructorError::IncorrectSegmentOrder {
                    expected_segment_index,
                    actual_segment_index: segment_index,
                });
            }
        }

        self.last_segment_index.replace(segment_index);

        for segment_item in segment_items {
            #[expect(
                clippy::rest_pattern_accessible_field,
                reason = "Do not need other fields"
            )]
            match segment_item {
                SegmentItem::Padding => {
                    // Doesn't contain anything
                }
                SegmentItem::Block { bytes, .. } => {
                    if !self.partial_block.is_empty() {
                        reconstructed_contents
                            .blocks
                            .push((next_block_number, mem::take(&mut self.partial_block)));

                        next_block_number = next_block_number
                            .checked_add(BlockNumber::ONE)
                            .ok_or(ReconstructorError::BlockNumberOverflow)?;
                    }

                    reconstructed_contents
                        .blocks
                        .push((next_block_number, Vec::from(bytes)));

                    next_block_number = next_block_number
                        .checked_add(BlockNumber::ONE)
                        .ok_or(ReconstructorError::BlockNumberOverflow)?;
                }
                SegmentItem::BlockStart { bytes, .. } => {
                    if !self.partial_block.is_empty() {
                        reconstructed_contents
                            .blocks
                            .push((next_block_number, mem::take(&mut self.partial_block)));

                        next_block_number = next_block_number
                            .checked_add(BlockNumber::ONE)
                            .ok_or(ReconstructorError::BlockNumberOverflow)?;
                    }

                    self.partial_block = Vec::from(bytes);
                }
                SegmentItem::BlockContinuation { bytes, .. } => {
                    if self.partial_block.is_empty() {
                        // This is continuation from previous segment, we don't have the beginning
                        // of the block to continue.
                        continue;
                    }

                    self.partial_block.extend_from_slice(&bytes);
                }
                SegmentItem::ParentSegmentHeader(_) => {
                    return Err(ReconstructorError::UnexpectedParentSegmentHeader);
                }
            }
        }

        Ok((self, reconstructed_contents))
    }
}