ab_archiving/

archiver.rs

use crate::objects::{BlockObject, GlobalObject};
use ab_core_primitives::block::BlockNumber;
use ab_core_primitives::hashes::Blake3Hash;
use ab_core_primitives::pieces::{
    InnerPiece, PieceHeader, PiecePosition, Record, RecordChunksRoot, RecordProof,
};
use ab_core_primitives::segments::{
    ArchivedBlockProgress, ArchivedHistorySegment, LastArchivedBlock, LocalSegmentIndex,
    RecordedHistorySegment, SegmentHeader, SegmentRoot,
};
use ab_core_primitives::shard::ShardIndex;
use ab_erasure_coding::ErasureCoding;
use ab_merkle_tree::blake3::BalancedMerkleTree;
use alloc::collections::VecDeque;
use alloc::vec;
use alloc::vec::Vec;
use core::cmp::Ordering;
use core::num::NonZeroU32;
use core::ops::Deref;
use parity_scale_codec::{Decode, Encode, Input, Output};
#[cfg(feature = "parallel")]
use rayon::prelude::*;

struct ArchivedHistorySegmentOutput<'a> {
    segment: &'a mut ArchivedHistorySegment,
    offset: usize,
}

impl Output for ArchivedHistorySegmentOutput<'_> {
    #[inline]
    fn write(&mut self, mut bytes: &[u8]) {
        while !bytes.is_empty() {
            let piece = self
                .segment
                .get_mut(self.offset / Record::SIZE)
                .expect("Encoding never exceeds the segment size; qed");
            let output = &mut piece.record.as_flattened_mut()[self.offset % Record::SIZE..];
            let bytes_to_write = output.len().min(bytes.len());
            output[..bytes_to_write].copy_from_slice(&bytes[..bytes_to_write]);
            self.offset += bytes_to_write;
            bytes = &bytes[bytes_to_write..];
        }
    }
}

/// The segment represents a collection of items stored in archival history
#[derive(Debug, Default, Clone, Eq, PartialEq)]
pub struct Segment {
    /// Segment items
    pub items: Vec<SegmentItem>,
}

impl Encode for Segment {
    #[inline(always)]
    fn size_hint(&self) -> usize {
        RecordedHistorySegment::SIZE
    }

    #[inline]
    fn encode_to<O>(&self, dest: &mut O)
    where
        O: Output + ?Sized,
    {
        for item in &self.items {
            item.encode_to(dest);
        }
    }
}

impl Decode for Segment {
    #[inline]
    fn decode<I>(input: &mut I) -> Result<Self, parity_scale_codec::Error>
    where
        I: Input,
    {
        let mut items = Vec::new();
        loop {
            match input.remaining_len()? {
                Some(0) => {
                    break;
                }
                Some(_) => {
                    // Processing continues below
                }
                None => {
                    return Err(
                        "Source doesn't report remaining length, decoding not possible".into(),
                    );
                }
            }

            match SegmentItem::decode(input) {
                Ok(item) => {
                    items.push(item);
                }
                Err(error) => {
                    return Err(error.chain("Could not decode `Segment::items`"));
                }
            }
        }

        Ok(Self { items })
    }
}

/// Similar to `Vec<u8>`, but when encoded with SCALE codec uses fixed size length encoding (as
/// little-endian `u32`)
#[derive(Debug, Clone, Eq, PartialEq)]
pub struct BlockBytes(Vec<u8>);

impl Deref for BlockBytes {
    type Target = [u8];

    #[inline(always)]
    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl From<BlockBytes> for Vec<u8> {
    #[inline(always)]
    fn from(value: BlockBytes) -> Self {
        value.0
    }
}

impl Encode for BlockBytes {
    #[inline(always)]
    fn size_hint(&self) -> usize {
        size_of::<u32>() + self.0.len()
    }

    #[inline]
    fn encode_to<O>(&self, dest: &mut O)
    where
        O: Output + ?Sized,
    {
        let length = u32::try_from(self.0.len())
            .expect("All constructors guarantee the size doesn't exceed `u32`; qed");

        length.encode_to(dest);
        dest.write(&self.0);
    }
}

impl Decode for BlockBytes {
    #[inline]
    fn decode<I>(input: &mut I) -> Result<Self, parity_scale_codec::Error>
    where
        I: Input,
    {
        let length = u32::decode(input)?;
        if length as usize > (RecordedHistorySegment::SIZE - size_of::<u32>()) {
            return Err("Segment item size is impossibly large".into());
        }
        // TODO: It is inefficient to zero it, but there is no API for it right now and actually
        //  implementation in `parity-scale-codec` itself is unsound:
        //  https://github.com/paritytech/parity-scale-codec/pull/605#discussion_r2076151291
        let mut bytes = vec![0; length as usize];
        input.read(&mut bytes)?;
        Ok(Self(bytes))
    }
}

impl BlockBytes {
    #[inline(always)]
    fn truncate(&mut self, size: usize) {
        self.0.truncate(size);
    }
}

/// Kinds of items that are contained within a segment
#[derive(Debug, Clone, Eq, PartialEq, Encode, Decode)]
pub enum SegmentItem {
    /// Special dummy enum variant only used as an implementation detail for padding purposes
    #[codec(index = 0)]
    Padding,
    /// Contains a full block inside
    #[codec(index = 1)]
    Block {
        /// Block bytes
        bytes: BlockBytes,
        /// This is a convenience implementation detail and will not be available on decoding
        #[doc(hidden)]
        #[codec(skip)]
        block_objects: Vec<BlockObject>,
    },
    /// Contains the beginning of the block inside, the remainder will be found in the following
    /// segments
    #[codec(index = 2)]
    BlockStart {
        /// Block bytes
        bytes: BlockBytes,
        /// This is a convenience implementation detail and will not be available on decoding
        #[doc(hidden)]
        #[codec(skip)]
        block_objects: Vec<BlockObject>,
    },
    /// Continuation of the partial block spilled over into the next segment
    #[codec(index = 3)]
    BlockContinuation {
        /// Block bytes
        bytes: BlockBytes,
        /// This is a convenience implementation detail and will not be available on decoding
        #[doc(hidden)]
        #[codec(skip)]
        block_objects: Vec<BlockObject>,
    },
    /// Segment header of the parent
    #[codec(index = 4)]
    ParentSegmentHeader(SegmentHeader),
}

/// Newly archived segment as a combination of a segment header and corresponding archived history
/// segment containing pieces
#[derive(Debug, Clone, Eq, PartialEq)]
pub struct NewArchivedSegment {
    /// Segment header
    pub segment_header: SegmentHeader,
    /// Segment of archived history containing pieces
    pub pieces: ArchivedHistorySegment,
}

/// The outcome of adding a block to the archiver.
#[derive(Debug, Clone, Eq, PartialEq)]
pub struct ArchiveBlockOutcome {
    /// The new segments archived after adding the block.
    /// There can be zero or more segments created after each block.
    pub archived_segments: Vec<NewArchivedSegment>,

    // TODO: The API for global objects needs to change and be associated with segments rather than
    //  blocks
    /// The new object mappings for those segments.
    /// There can be zero or more mappings created after each block.
    pub global_objects: Vec<GlobalObject>,
}

/// Archiver instantiation error
#[derive(Debug, Clone, Eq, PartialEq, Ord, PartialOrd, Hash, thiserror::Error)]
pub enum ArchiverInstantiationError {
    /// Invalid last archived block, its size is the same as the encoded block
    /// (so it should have been completely archived, not partially archived)
    #[error("Invalid last archived block, its size {0} bytes is the same as the encoded block")]
    InvalidLastArchivedBlock(u32),
    /// Invalid block, its size is smaller than the already archived number of bytes
    #[error(
        "Invalid block, its size {block_bytes} bytes is smaller than the already archived block \
        {archived_block_bytes} bytes"
    )]
    InvalidBlockSmallSize {
        /// Full block size
        block_bytes: u32,
        /// Already archived portion of the block
        archived_block_bytes: u32,
    },
    /// Invalid block, its size is larger than [`u32::MAX`] bytes
    #[error("Invalid block, its size {block_bytes} bytes is larger than `u32::MAX`")]
    InvalidBlockLargeSize {
        /// Full block size
        block_bytes: usize,
    },
}

/// Size of the smallest segment item with data: 1 byte for enum variant, 4 bytes for length and 1
/// for the actual data.
///
/// A segment is complete when less than this is left in it. Segment header item is never the last
/// one, so a block that is split always has some of its data before the split point.
const MIN_DATA_ITEM_SIZE: usize = 6;

/// Buffer of segment items that keeps track of their total encoded size and which items had their
/// object mappings produced already
#[derive(Debug, Default, Clone)]
struct SegmentItemsBuffer {
    items: VecDeque<SegmentItem>,
    /// Total encoded size of all items
    encoded_size: usize,
    /// Number of items at the beginning of the buffer that had their object mappings produced
    num_mapped_items: usize,
    /// Total encoded size of items that had their object mappings produced
    mapped_encoded_size: usize,
}

impl SegmentItemsBuffer {
    fn len(&self) -> usize {
        self.items.len()
    }

    fn push_back(&mut self, item: SegmentItem) {
        self.encoded_size += item.encoded_size();
        self.items.push_back(item);
    }

    fn push_front(&mut self, item: SegmentItem) {
        self.encoded_size += item.encoded_size();
        self.items.push_front(item);
        // New item may have object mappings and shifts offsets of all other items, so all items are
        // considered unmapped (items that were mapped before have no block objects left anyway)
        self.num_mapped_items = 0;
        self.mapped_encoded_size = 0;
    }

    fn pop_front(&mut self) -> Option<SegmentItem> {
        let item = self.items.pop_front()?;
        let item_encoded_size = item.encoded_size();
        self.encoded_size = self
            .encoded_size
            .checked_sub(item_encoded_size)
            .expect("Item size was added to the total when the item was inserted; qed");
        if self.num_mapped_items > 0 {
            self.num_mapped_items -= 1;
            self.mapped_encoded_size = self
                .mapped_encoded_size
                .checked_sub(item_encoded_size)
                .expect("Mapped item size was added to the mapped total; qed");
        }
        Some(item)
    }

    /// Returns items that didn't have their object mappings produced yet along with the offset of
    /// the first of them from the beginning of the buffer.
    ///
    /// Returned items are considered mapped afterwards, so object mappings must be produced for
    /// them. Only fields that are not encoded (like block objects) may be modified, otherwise the
    /// total encoded size will be wrong.
    fn take_unmapped_items(&mut self) -> (usize, impl Iterator<Item = &mut SegmentItem>) {
        let offset = self.mapped_encoded_size;
        let first_unmapped_item = self.num_mapped_items;

        self.num_mapped_items = self.items.len();
        self.mapped_encoded_size = self.encoded_size;

        (offset, self.items.range_mut(first_unmapped_item..))
    }
}

/// Block archiver.
///
/// It takes new confirmed (at `K` depth) blocks and concatenates them into a buffer, buffer is
/// sliced into segments of [`RecordedHistorySegment::SIZE`] size, segments are sliced into source
/// records of [`Record::SIZE`], records are erasure coded, committed to, then roots with proofs are
/// appended and records become pieces that are returned alongside the corresponding segment header.
#[derive(Debug, Clone)]
pub struct Archiver {
    shard_index: ShardIndex,
    /// Buffer containing blocks and other buffered items that are pending to be included into the
    /// next segment
    buffer: SegmentItemsBuffer,
    /// Erasure coding data structure
    erasure_coding: ErasureCoding,
    /// An index of the current segment
    segment_index: LocalSegmentIndex,
    /// Hash of the segment header of the previous segment
    prev_segment_header_hash: Blake3Hash,
    /// Last archived block
    last_archived_block: Option<LastArchivedBlock>,
}

impl Archiver {
    /// Create a new instance
    pub fn new(shard_index: ShardIndex, erasure_coding: ErasureCoding) -> Self {
        Self {
            shard_index,
            buffer: SegmentItemsBuffer::default(),
            erasure_coding,
            segment_index: LocalSegmentIndex::ZERO,
            prev_segment_header_hash: Blake3Hash::default(),
            last_archived_block: None,
        }
    }

    /// Create a new instance of the archiver with the initial state in case of restart.
    ///
    /// `block` corresponds to `last_archived_block` and will be processed according to its state.
    pub fn with_initial_state(
        shard_index: ShardIndex,
        erasure_coding: ErasureCoding,
        segment_header: SegmentHeader,
        encoded_block: &[u8],
        mut block_objects: Vec<BlockObject>,
    ) -> Result<Self, ArchiverInstantiationError> {
        let mut archiver = Self::new(shard_index, erasure_coding);

        archiver.segment_index = segment_header.index.as_inner() + LocalSegmentIndex::ONE;
        archiver.prev_segment_header_hash = segment_header.hash();
        archiver.last_archived_block = Some(segment_header.last_archived_block);

        // The first thing in the buffer should be segment header
        archiver
            .buffer
            .push_back(SegmentItem::ParentSegmentHeader(segment_header));

        if let Some(archived_block_bytes) = archiver
            .last_archived_block
            .expect("Just inserted; qed")
            .partial_archived()
        {
            let archived_block_bytes = archived_block_bytes.get();
            let encoded_block_bytes = u32::try_from(encoded_block.len()).map_err(|_error| {
                ArchiverInstantiationError::InvalidBlockLargeSize {
                    block_bytes: encoded_block.len(),
                }
            })?;

            match encoded_block_bytes.cmp(&archived_block_bytes) {
                Ordering::Less => {
                    return Err(ArchiverInstantiationError::InvalidBlockSmallSize {
                        block_bytes: encoded_block_bytes,
                        archived_block_bytes,
                    });
                }
                Ordering::Equal => {
                    return Err(ArchiverInstantiationError::InvalidLastArchivedBlock(
                        encoded_block_bytes,
                    ));
                }
                Ordering::Greater => {
                    // Take part of the encoded block that wasn't archived yet and push to the
                    // buffer as a block continuation
                    block_objects.retain_mut(|block_object: &mut BlockObject| {
                        if block_object.offset >= archived_block_bytes {
                            block_object.offset -= archived_block_bytes;
                            true
                        } else {
                            false
                        }
                    });
                    archiver.buffer.push_back(SegmentItem::BlockContinuation {
                        bytes: BlockBytes(
                            encoded_block[(archived_block_bytes as usize)..].to_vec(),
                        ),
                        block_objects,
                    });
                }
            }
        }

        Ok(archiver)
    }

    /// Get the last archived block if there was any
    pub fn last_archived_block_number(&self) -> Option<BlockNumber> {
        self.last_archived_block
            .map(|last_archived_block| last_archived_block.number())
    }

    /// Adds a new block to the internal buffer, potentially producing pieces, segment headers, and
    /// object mappings.
    ///
    /// NOTE:
    /// * pieces inside [`NewArchivedSegment`] are shared initially for efficient memory reuse when
    ///   working with pieces
    /// * pieces do not have a segment position and a segment proof filled in because it is not yet
    ///   known at the time of local segment creation
    ///
    /// Returns `None` if the block is larger than [`u32::MAX`] bytes.
    pub fn add_block(
        &mut self,
        bytes: Vec<u8>,
        block_objects: Vec<BlockObject>,
    ) -> Option<ArchiveBlockOutcome> {
        if u32::try_from(bytes.len()).is_err() {
            return None;
        }

        // Append new block to the buffer
        self.buffer.push_back(SegmentItem::Block {
            bytes: BlockBytes(bytes),
            block_objects,
        });

        let mut archived_segments = Vec::new();
        let mut object_mapping = Vec::new();

        // Add completed segments and their mappings for this block
        while let Some(mut segment) = self.produce_segment() {
            // Produce any segment mappings that haven't already been produced.
            object_mapping.extend(Self::produce_object_mappings(
                Segment::default().encoded_size(),
                segment.items.iter_mut(),
            ));
            archived_segments.push(self.produce_archived_segment(segment));
        }

        // Produce any next segment buffer mappings that haven't already been produced.
        object_mapping.extend(self.produce_next_segment_mappings());

        Some(ArchiveBlockOutcome {
            archived_segments,
            global_objects: object_mapping,
        })
    }

    /// Try to slice buffer contents into segments if there is enough data, producing one segment at
    /// a time
    fn produce_segment(&mut self) -> Option<Segment> {
        let mut segment_size = Segment::default().encoded_size();

        // Not enough data for a segment yet
        if RecordedHistorySegment::SIZE.saturating_sub(segment_size + self.buffer.encoded_size)
            >= MIN_DATA_ITEM_SIZE
        {
            return None;
        }

        let mut segment = Segment {
            items: Vec::with_capacity(self.buffer.len()),
        };

        let mut last_archived_block = self.last_archived_block;

        while RecordedHistorySegment::SIZE.saturating_sub(segment_size) >= MIN_DATA_ITEM_SIZE {
            let segment_item = self.buffer.pop_front().expect(
                "Buffered items are large enough to fill a segment, which is checked above; qed",
            );

            let segment_item_encoded_size = segment_item.encoded_size();
            segment_size += segment_item_encoded_size;

            // Check if there is an excess of data that should be spilled over into the next segment
            let spill_over = segment_size.saturating_sub(RecordedHistorySegment::SIZE);

            let segment_item = match segment_item {
                SegmentItem::Padding => {
                    unreachable!("Buffer never contains SegmentItem::Padding; qed");
                }
                SegmentItem::Block {
                    mut bytes,
                    mut block_objects,
                } => {
                    let last_archived_block =
                        if let Some(last_archived_block) = &mut last_archived_block {
                            // Increase the archived block number and assume the whole block was
                            // archived (spill over checked below)
                            last_archived_block
                                .number
                                .replace(last_archived_block.number() + BlockNumber::ONE);
                            last_archived_block.set_complete();
                            last_archived_block
                        } else {
                            // Genesis block
                            last_archived_block.insert(LastArchivedBlock {
                                number: BlockNumber::ZERO.into(),
                                archived_progress: ArchivedBlockProgress::new_complete(),
                            })
                        };

                    if spill_over == 0 {
                        SegmentItem::Block {
                            bytes,
                            block_objects,
                        }
                    } else {
                        let split_point = bytes.len() - spill_over;

                        {
                            let continuation_bytes = bytes[split_point..].to_vec();
                            let continuation_block_objects = block_objects
                                .extract_if(.., |block_object: &mut BlockObject| {
                                    if block_object.offset >= split_point as u32 {
                                        block_object.offset -= split_point as u32;
                                        true
                                    } else {
                                        false
                                    }
                                })
                                .collect();

                            // Push a continuation element back into the buffer where the removed
                            // segment item was
                            self.buffer.push_front(SegmentItem::BlockContinuation {
                                bytes: BlockBytes(continuation_bytes),
                                block_objects: continuation_block_objects,
                            });
                        }

                        bytes.truncate(split_point);
                        // Update last archived block to include partial archiving info
                        let archived_bytes = u32::try_from(split_point)
                            .ok()
                            .and_then(NonZeroU32::new)
                            .expect(
                                "The loop condition leaves space for at least one byte of the \
                                block before the split point, which is below the block size \
                                that `::add_block()` ensures fits into `u32`; qed",
                            );
                        last_archived_block.set_partial_archived(archived_bytes);

                        SegmentItem::BlockStart {
                            bytes,
                            block_objects,
                        }
                    }
                }
                #[expect(
                    clippy::rest_pattern_accessible_field,
                    reason = "Do not care about fields"
                )]
                SegmentItem::BlockStart { .. } => {
                    unreachable!("Buffer never contains SegmentItem::BlockStart; qed");
                }
                SegmentItem::BlockContinuation {
                    mut bytes,
                    mut block_objects,
                } => {
                    let last_archived_block = last_archived_block.as_mut().expect(
                        "Block continuation implies that there are some bytes archived \
                        already; qed",
                    );

                    let previously_archived_bytes = last_archived_block.partial_archived().expect(
                        "Block continuation implies that there are some bytes archived \
                        already; qed",
                    );

                    if spill_over == 0 {
                        last_archived_block.set_complete();

                        SegmentItem::BlockContinuation {
                            bytes,
                            block_objects,
                        }
                    } else {
                        let split_point = bytes.len() - spill_over;

                        {
                            let continuation_bytes = bytes[split_point..].to_vec();
                            let continuation_block_objects = block_objects
                                .extract_if(.., |block_object: &mut BlockObject| {
                                    if block_object.offset >= split_point as u32 {
                                        block_object.offset -= split_point as u32;
                                        true
                                    } else {
                                        false
                                    }
                                })
                                .collect();
                            // Push a continuation element back into the buffer where the removed
                            // segment item was
                            self.buffer.push_front(SegmentItem::BlockContinuation {
                                bytes: BlockBytes(continuation_bytes),
                                block_objects: continuation_block_objects,
                            });
                        }

                        bytes.truncate(split_point);
                        // Update last archived block to include partial archiving info
                        let archived_bytes = previously_archived_bytes.get()
                            + u32::try_from(split_point).expect(
                                "The split point is below the size of the block, which \
                                `::add_block()` and `::with_initial_state()` ensure fits into \
                                `u32`; qed",
                            );
                        let archived_bytes = NonZeroU32::new(archived_bytes).expect(
                            "Spillover means non-zero length of the block was archived; qed",
                        );
                        last_archived_block.set_partial_archived(archived_bytes);

                        SegmentItem::BlockContinuation {
                            bytes,
                            block_objects,
                        }
                    }
                }
                SegmentItem::ParentSegmentHeader(parent_segment_header) => {
                    // We are not interested in segment header here
                    SegmentItem::ParentSegmentHeader(parent_segment_header)
                }
            };

            segment.items.push(segment_item);
        }

        self.last_archived_block = last_archived_block;

        Some(segment)
    }

    /// Produce object mappings for the buffered items for the next segment. Then remove the
    /// mappings in those items.
    ///
    /// Must only be called after all complete segments for a block have been produced. Before
    /// that, the buffer can contain a `BlockContinuation` which spans multiple segments.
    fn produce_next_segment_mappings(&mut self) -> Vec<GlobalObject> {
        // Usually this is just the block that was just added, avoiding iteration over the whole
        // buffer is important for performance when there are many small blocks
        let (offset, unmapped_items) = self.buffer.take_unmapped_items();

        Self::produce_object_mappings(Segment::default().encoded_size() + offset, unmapped_items)
    }

    /// Produce object mappings for `items` in `segment_index`, where the first item starts at
    /// `base_offset_in_segment`. Then remove the mappings from those items.
    ///
    /// This method can be called on a `Segment`’s items, or on the `Archiver`'s internal buffer.
    fn produce_object_mappings<'a, Items>(
        mut base_offset_in_segment: usize,
        items: Items,
    ) -> Vec<GlobalObject>
    where
        Items: Iterator<Item = &'a mut SegmentItem>,
    {
        let mut corrected_object_mapping = Vec::new();
        for segment_item in items {
            match segment_item {
                SegmentItem::Padding => {
                    unreachable!(
                        "Segment during archiving never contains SegmentItem::Padding; qed"
                    );
                }
                SegmentItem::Block {
                    bytes: _,
                    block_objects,
                }
                | SegmentItem::BlockStart {
                    bytes: _,
                    block_objects,
                }
                | SegmentItem::BlockContinuation {
                    bytes: _,
                    block_objects,
                } => {
                    for block_object in block_objects.drain(..) {
                        // `+1` corresponds to `SegmentItem::X {}` enum variant encoding
                        let offset_in_segment = base_offset_in_segment
                            + 1
                            + u32::encoded_fixed_size().expect("Fixed size; qed")
                            + block_object.offset as usize;
                        let raw_piece_offset = (offset_in_segment % Record::SIZE)
                            .try_into()
                            .expect("Offset within piece should always fit in 32-bit integer; qed");
                        corrected_object_mapping.push(GlobalObject {
                            hash: block_object.hash,
                            piece_position: PiecePosition::from(
                                (offset_in_segment / Record::SIZE) as u8,
                            ),
                            offset: raw_piece_offset,
                        });
                    }
                }
                SegmentItem::ParentSegmentHeader(_) => {
                    // Ignore, no object mappings here
                }
            }

            base_offset_in_segment += segment_item.encoded_size();
        }

        corrected_object_mapping
    }

    /// Take a segment as an input, apply necessary transformations, and produce an archived segment
    fn produce_archived_segment(&mut self, segment: Segment) -> NewArchivedSegment {
        let mut pieces = {
            let mut pieces = ArchivedHistorySegment::default();

            segment.encode_to(&mut ArchivedHistorySegmentOutput {
                segment: &mut pieces,
                offset: 0,
            });
            // Segment is quite big and no longer necessary
            drop(segment);

            let (source_records, parity_records) = pieces.split_records_mut();

            #[cfg(not(feature = "parallel"))]
            let extend = ErasureCoding::extend_scattered;
            #[cfg(feature = "parallel")]
            let extend = ErasureCoding::par_extend_scattered;

            extend(
                &self.erasure_coding,
                source_records.map(|record| &*record),
                parity_records,
            )
            .expect("Statically correct parameters; qed");

            pieces
        };

        // Collect hashes to roots from all records
        let mut record_roots = [[0; _]; ArchivedHistorySegment::NUM_PIECES];
        {
            // Here we build a tree of record chunks, with the first half being source chunks as
            // they are originally and the second half being parity chunks. While we build tree
            // threes here (for source chunks, parity chunks and combined for the whole record), it
            // could have been a single tree, and it would end up with the same root. Building them
            // separately requires less RAM and allows capturing parity chunks root more easily.
            //
            // `parity_chunks` is only used for deriving a root and is fully overwritten by erasure
            // coding, so it is reused between pieces.
            let process_piece =
                |parity_chunks: &mut Record, piece: &mut InnerPiece, record_root: &mut [u8; _]| {
                    self.erasure_coding
                        .extend(&piece.record, parity_chunks)
                        .expect(
                            "Erasure coding instance is deliberately configured to support this \
                            input; qed",
                        );

                    let source_chunks_root = *piece.record.source_chunks_root();
                    let parity_chunks_root = BalancedMerkleTree::compute_root_only(parity_chunks);

                    *record_root = BalancedMerkleTree::compute_root_only(&[
                        source_chunks_root,
                        parity_chunks_root,
                    ]);

                    piece.header.parity_chunks_root = RecordChunksRoot::from(parity_chunks_root);
                };

            #[cfg(not(feature = "parallel"))]
            {
                let mut parity_chunks = Record::new_boxed();
                pieces
                    .iter_mut()
                    .zip(&mut record_roots)
                    .for_each(|(piece, record_root)| {
                        process_piece(&mut parity_chunks, piece, record_root);
                    });
            }
            #[cfg(feature = "parallel")]
            pieces
                .par_iter_mut()
                .zip_eq(&mut record_roots)
                .for_each_init(Record::new_boxed, |parity_chunks, (piece, record_root)| {
                    process_piece(parity_chunks, piece, record_root);
                });
        }

        let segment_merkle_tree = BalancedMerkleTree::new_boxed(&record_roots);

        let segment_root = SegmentRoot::from(segment_merkle_tree.root());

        // Now produce segment header
        let segment_header = SegmentHeader {
            index: self.segment_index.into(),
            root: segment_root,
            prev_segment_header_hash: self.prev_segment_header_hash,
            last_archived_block: self
                .last_archived_block
                .expect("Never empty by the time segment is produced; qed"),
        };

        // Create proof for every record and write it to corresponding piece.
        pieces
            .iter_mut()
            .zip(segment_merkle_tree.all_proofs())
            .for_each(|(piece, record_proof)| {
                piece.header = PieceHeader {
                    shard_index: self.shard_index.into(),
                    local_segment_index: segment_header.index,
                    segment_root: segment_header.root,
                    record_proof: RecordProof::from(record_proof),
                    ..piece.header
                };
            });

        // Update state
        self.segment_index += LocalSegmentIndex::ONE;
        self.prev_segment_header_hash = segment_header.hash();

        // Add segment header to the beginning of the buffer to be the first thing included in the
        // next segment
        self.buffer
            .push_front(SegmentItem::ParentSegmentHeader(segment_header));

        NewArchivedSegment {
            segment_header,
            pieces: pieces.to_shared(),
        }
    }
}