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 have monotonically increasing segment index
    #[error(
        "Incorrect segment order, expected index {expected_segment_index}, actual \
        {actual_segment_index}"
    )]
    IncorrectSegmentOrder {
        expected_segment_index: LocalSegmentIndex,
        actual_segment_index: LocalSegmentIndex,
    },
}

/// 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
    partial_block: Option<Vec<u8>>,
}

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

    /// 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)
            });

            self.erasure_coding.recover_source_scattered(
                {
                    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.
    ///
    /// 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<ReconstructedContents, ReconstructorError> {
        let segment = self.reconstruct_segment(segment_pieces)?;

        let mut reconstructed_contents = ReconstructedContents::default();
        let mut next_block_number = BlockNumber::ZERO;
        let mut partial_block = self.partial_block.take().unwrap_or_default();

        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 !partial_block.is_empty() {
                        reconstructed_contents
                            .blocks
                            .push((next_block_number, mem::take(&mut partial_block)));

                        next_block_number += BlockNumber::ONE;
                    }

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

                    next_block_number += BlockNumber::ONE;
                }
                SegmentItem::BlockStart { bytes, .. } => {
                    if !partial_block.is_empty() {
                        reconstructed_contents
                            .blocks
                            .push((next_block_number, mem::take(&mut partial_block)));

                        next_block_number += BlockNumber::ONE;
                    }

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

                    partial_block.extend_from_slice(&bytes);
                }
                SegmentItem::ParentSegmentHeader(segment_header) => {
                    let segment_index = segment_header.index.as_inner();

                    if let Some(last_segment_index) = self.last_segment_index
                        && last_segment_index != segment_index
                    {
                        return Err(ReconstructorError::IncorrectSegmentOrder {
                            expected_segment_index: last_segment_index + LocalSegmentIndex::ONE,
                            actual_segment_index: segment_index + LocalSegmentIndex::ONE,
                        });
                    }

                    self.last_segment_index
                        .replace(segment_index + LocalSegmentIndex::ONE);

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

                    reconstructed_contents
                        .segment_header
                        .replace(segment_header);

                    match archived_progress.partial() {
                        None => {
                            reconstructed_contents
                                .blocks
                                .push((next_block_number, mem::take(&mut partial_block)));

                            next_block_number = number.as_inner() + BlockNumber::ONE;
                        }
                        Some(_bytes) => {
                            next_block_number = number.as_inner();

                            if partial_block.is_empty() {
                                // Will not be able to recover full block, bump right away.
                                next_block_number += BlockNumber::ONE;
                            }
                        }
                    }
                }
            }
        }

        if !partial_block.is_empty() {
            self.partial_block.replace(partial_block);
        }

        if self.last_segment_index.is_none() {
            self.last_segment_index.replace(LocalSegmentIndex::ZERO);
        }

        Ok(reconstructed_contents)
    }
}