ab_archiving/
piece_reconstructor.rs

1use ab_core_primitives::pieces::{
2    Piece, PieceHeader, PiecePosition, Record, RecordChunksRoot, RecordProof, SegmentProof,
3};
4use ab_core_primitives::segments::{
5    ArchivedHistorySegment, LocalSegmentIndex, RecordedHistorySegment, SegmentPosition,
6    SegmentRoot, SuperSegmentIndex,
7};
8use ab_core_primitives::shard::ShardIndex;
9use ab_erasure_coding::{ErasureCoding, ErasureCodingError, ShardsPresent};
10use ab_merkle_tree::balanced::BalancedMerkleTree;
11use alloc::vec::Vec;
12#[cfg(feature = "parallel")]
13use rayon::prelude::*;
14
15/// Reconstructor-related instantiation error
16#[derive(Debug, Clone, PartialEq, thiserror::Error)]
17pub enum ReconstructorError {
18    /// Segment size is not bigger than record size
19    #[error("Error during data shards reconstruction: {0}")]
20    DataShardsReconstruction(#[from] ErasureCodingError),
21    /// Not enough shards
22    #[error("Not enough shards: {num_shards}")]
23    NotEnoughShards { num_shards: usize },
24}
25
26struct SharedPieceDetails {
27    shard_index: ShardIndex,
28    local_segment_index: LocalSegmentIndex,
29    super_segment_index: SuperSegmentIndex,
30    segment_position: SegmentPosition,
31    segment_root: SegmentRoot,
32    segment_proof: SegmentProof,
33}
34
35/// Piece reconstructor helps to reconstruct missing pieces.
36#[derive(Debug, Clone)]
37pub struct PiecesReconstructor {
38    /// Erasure coding data structure
39    erasure_coding: ErasureCoding,
40}
41
42impl PiecesReconstructor {
43    /// Create a new instance
44    pub fn new(erasure_coding: ErasureCoding) -> Self {
45        Self { erasure_coding }
46    }
47
48    fn reconstruct_shards(
49        &self,
50        input_pieces: &[Option<Piece>],
51    ) -> Result<ArchivedHistorySegment, ReconstructorError> {
52        if input_pieces.len() < ArchivedHistorySegment::NUM_PIECES {
53            return Err(ReconstructorError::NotEnoughShards {
54                num_shards: input_pieces.len(),
55            });
56        }
57        let mut reconstructed_pieces = ArchivedHistorySegment::default();
58
59        // TODO: Fix up piece metadata
60        let mut shared_piece_details = None;
61        {
62            let (source_input_pieces, parity_input_pieces) =
63                input_pieces.split_at(RecordedHistorySegment::NUM_RAW_RECORDS);
64            let (source_reconstructed_pieces, parity_reconstructed_pieces) =
65                reconstructed_pieces.split_at_mut(RecordedHistorySegment::NUM_RAW_RECORDS);
66
67            let mut present = ShardsPresent::none();
68
69            for (index, (maybe_input_piece, output_piece)) in source_input_pieces
70                .iter()
71                .zip(source_reconstructed_pieces.iter_mut())
72                .enumerate()
73            {
74                if let Some(input_piece) = maybe_input_piece {
75                    if shared_piece_details.is_none() {
76                        shared_piece_details.replace(SharedPieceDetails {
77                            shard_index: input_piece.header.shard_index.as_inner(),
78                            local_segment_index: input_piece.header.local_segment_index.as_inner(),
79                            super_segment_index: input_piece.header.super_segment_index.as_inner(),
80                            segment_position: input_piece.header.segment_position.as_inner(),
81                            segment_root: input_piece.header.segment_root,
82                            segment_proof: input_piece.header.segment_proof,
83                        });
84                    }
85                    // Fancy way to insert value to avoid going through stack (if naive
86                    // dereferencing is used) and potentially causing stack overflow as the result
87                    output_piece.record.copy_from_slice(&*input_piece.record);
88                    present.source.set(index);
89                }
90            }
91
92            for (index, (maybe_input_piece, output_piece)) in parity_input_pieces
93                .iter()
94                .zip(parity_reconstructed_pieces.iter_mut())
95                .enumerate()
96            {
97                if let Some(input_piece) = maybe_input_piece {
98                    output_piece.record.copy_from_slice(&*input_piece.record);
99                    present.parity.set(index);
100                }
101            }
102
103            let (source_records, parity_records) = reconstructed_pieces.split_records_mut();
104
105            self.erasure_coding
106                .recover_all_scattered(source_records, parity_records, &present)?;
107        }
108        let SharedPieceDetails {
109            shard_index,
110            local_segment_index,
111            super_segment_index,
112            segment_position,
113            segment_root,
114            segment_proof,
115        } = shared_piece_details.expect(
116            "Sucessful recovery means there was at least one piece to fill this Option; qed",
117        );
118
119        let record_roots = {
120            #[cfg(not(feature = "parallel"))]
121            let iter = reconstructed_pieces.iter_mut().zip(input_pieces);
122            #[cfg(feature = "parallel")]
123            let iter = reconstructed_pieces.par_iter_mut().zip_eq(input_pieces);
124
125            iter.map(|(piece, maybe_input_piece)| {
126                let (record_root, parity_chunks_root) = if let Some(input_piece) = maybe_input_piece
127                {
128                    (
129                        *input_piece.record_root(),
130                        *input_piece.header.parity_chunks_root,
131                    )
132                } else {
133                    // TODO: Reuse allocations between iterations
134                    let [source_chunks_root, parity_chunks_root] = {
135                        let mut parity_chunks = Record::new_boxed();
136
137                        self.erasure_coding
138                            .extend(&piece.record, &mut parity_chunks)?;
139
140                        let source_chunks_root = *piece.record.source_chunks_root();
141                        let parity_chunks_root =
142                            BalancedMerkleTree::compute_root_only(&parity_chunks);
143
144                        [source_chunks_root, parity_chunks_root]
145                    };
146
147                    let record_root =
148                        BalancedMerkleTree::new(&[source_chunks_root, parity_chunks_root]).root();
149
150                    (record_root, parity_chunks_root)
151                };
152
153                piece.header.parity_chunks_root = RecordChunksRoot::from(parity_chunks_root);
154
155                Ok::<_, ReconstructorError>(record_root)
156            })
157            .collect::<Result<Vec<_>, _>>()?
158        };
159
160        let segment_merkle_tree =
161            BalancedMerkleTree::<{ ArchivedHistorySegment::NUM_PIECES }>::new_boxed(
162                record_roots
163                    .as_slice()
164                    .try_into()
165                    .expect("Statically guaranteed to have correct length; qed"),
166            );
167
168        reconstructed_pieces
169            .iter_mut()
170            .zip(segment_merkle_tree.all_proofs())
171            .for_each(|(piece, record_proof)| {
172                piece.header = PieceHeader {
173                    shard_index: shard_index.into(),
174                    local_segment_index: local_segment_index.into(),
175                    super_segment_index: super_segment_index.into(),
176                    segment_position: segment_position.into(),
177                    segment_root,
178                    segment_proof,
179                    parity_chunks_root: piece.header.parity_chunks_root,
180                    record_proof: RecordProof::from(record_proof),
181                };
182            });
183
184        Ok(reconstructed_pieces)
185    }
186
187    /// Returns all the pieces for a segment using a given set of pieces of a segment of the
188    /// archived history.
189    ///
190    /// Any half of all pieces are required to be present, the rest will be recovered automatically
191    /// due to use of erasure coding if needed.
192    pub fn reconstruct_segment(
193        &self,
194        segment_pieces: &[Option<Piece>],
195    ) -> Result<ArchivedHistorySegment, ReconstructorError> {
196        let pieces = self.reconstruct_shards(segment_pieces)?;
197
198        Ok(pieces.to_shared())
199    }
200
201    /// Returns the missing piece for a segment using given set of pieces of a segment of the
202    /// archived history (any half of all pieces are required to be present).
203    pub fn reconstruct_piece(
204        &self,
205        segment_pieces: &[Option<Piece>],
206        piece_position: PiecePosition,
207    ) -> Result<Piece, ReconstructorError> {
208        // TODO: Early exit if position already exists and doesn't need reconstruction
209        // TODO: It is now inefficient to recover all shards if only one piece is needed, especially
210        //  source piece
211        let pieces = self.reconstruct_shards(segment_pieces)?;
212
213        let piece = Piece::from(&pieces[piece_position]);
214
215        Ok(piece.to_shared())
216    }
217}
ab_archiving/piece_reconstructor.rs

ab_archiving/
piece_reconstructor.rs
