ab_core_primitives/pieces/
flat_pieces.rs

1use crate::pieces::cow_bytes::CowBytes;
2use crate::pieces::{Piece, PieceArray};
3use crate::segments::RecordedHistorySegment;
4use alloc::boxed::Box;
5use bytes::{Bytes, BytesMut};
6use core::ops::{Deref, DerefMut};
7use core::{fmt, slice};
8#[cfg(feature = "parallel")]
9use rayon::prelude::*;
10
11/// Flat representation of multiple pieces concatenated for more efficient for processing
12#[derive(Clone, PartialEq, Eq)]
13pub struct FlatPieces(CowBytes);
14
15impl fmt::Debug for FlatPieces {
16    #[inline]
17    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
18        f.debug_struct("FlatPieces").finish_non_exhaustive()
19    }
20}
21
22impl Deref for FlatPieces {
23    type Target = [PieceArray];
24
25    #[inline]
26    fn deref(&self) -> &Self::Target {
27        let bytes = self.0.as_ref();
28        // SAFETY: Bytes slice has length of multiples of piece size and lifetimes of returned data
29        // are preserved
30        let pieces = unsafe {
31            slice::from_raw_parts(
32                bytes.as_ptr().cast::<[u8; Piece::SIZE]>(),
33                bytes.len() / Piece::SIZE,
34            )
35        };
36        PieceArray::slice_from_repr(pieces)
37    }
38}
39
40impl DerefMut for FlatPieces {
41    #[inline]
42    fn deref_mut(&mut self) -> &mut Self::Target {
43        let bytes = self.0.as_mut();
44        // SAFETY: Bytes slice has length of multiples of piece size and lifetimes of returned data
45        // are preserved
46        let pieces = unsafe {
47            slice::from_raw_parts_mut(
48                bytes.as_mut_ptr().cast::<[u8; Piece::SIZE]>(),
49                bytes.len() / Piece::SIZE,
50            )
51        };
52        PieceArray::slice_mut_from_repr(pieces)
53    }
54}
55
56impl FlatPieces {
57    /// Allocate `FlatPieces` that will hold `piece_count` pieces filled with zeroes
58    #[inline]
59    pub fn new(piece_count: usize) -> Self {
60        Self(CowBytes::Owned(BytesMut::zeroed(piece_count * Piece::SIZE)))
61    }
62
63    /// Iterate over all pieces.
64    ///
65    /// NOTE: [`Piece`] returned by iterator is sharing memory with the original [`FlatPieces`]
66    /// allocation, which is cheap, but may result in additional memory usage if the piece is kept
67    /// around for a long time. But it only happens if [`Self::to_shared`] was called first,
68    /// otherwise iterator may have to allocate each piece from scratch, which is rarely a
69    /// desired behavior.
70    #[inline]
71    pub fn pieces(&self) -> Box<dyn ExactSizeIterator<Item = Piece> + '_> {
72        match &self.0 {
73            CowBytes::Shared(bytes) => Box::new(
74                bytes
75                    .chunks_exact(Piece::SIZE)
76                    .map(|slice| Piece(CowBytes::Shared(bytes.slice_ref(slice)))),
77            ),
78            CowBytes::Owned(bytes) => Box::new(
79                bytes
80                    .chunks_exact(Piece::SIZE)
81                    .map(|slice| Piece(CowBytes::Shared(Bytes::copy_from_slice(slice)))),
82            ),
83        }
84    }
85
86    /// Iterator over source pieces (even indices)
87    #[inline]
88    pub fn source_pieces(&self) -> impl ExactSizeIterator<Item = Piece> + '_ {
89        self.pieces().take(RecordedHistorySegment::NUM_RAW_RECORDS)
90    }
91
92    /// Iterator over source pieces (even indices)
93    #[inline]
94    pub fn source(&self) -> impl ExactSizeIterator<Item = &'_ PieceArray> + '_ {
95        self.iter().take(RecordedHistorySegment::NUM_RAW_RECORDS)
96    }
97
98    /// Mutable iterator over source pieces (even indices)
99    #[inline]
100    pub fn source_mut(&mut self) -> impl ExactSizeIterator<Item = &'_ mut PieceArray> + '_ {
101        self.iter_mut()
102            .take(RecordedHistorySegment::NUM_RAW_RECORDS)
103    }
104
105    /// Iterator over parity pieces (odd indices)
106    #[inline]
107    pub fn parity_pieces(&self) -> impl ExactSizeIterator<Item = Piece> + '_ {
108        self.pieces().skip(RecordedHistorySegment::NUM_RAW_RECORDS)
109    }
110
111    /// Iterator over parity pieces (odd indices)
112    #[inline]
113    pub fn parity(&self) -> impl ExactSizeIterator<Item = &'_ PieceArray> + '_ {
114        self.iter().skip(RecordedHistorySegment::NUM_RAW_RECORDS)
115    }
116
117    /// Mutable iterator over parity pieces (odd indices)
118    #[inline]
119    pub fn parity_mut(&mut self) -> impl ExactSizeIterator<Item = &'_ mut PieceArray> + '_ {
120        self.iter_mut()
121            .skip(RecordedHistorySegment::NUM_RAW_RECORDS)
122    }
123
124    /// Ensure flat pieces contains cheaply cloneable shared data.
125    ///
126    /// Internally flat pieces uses CoW mechanism and can store either mutable owned data or data
127    /// that is cheap to clone, calling this method will ensure further clones and returned pieces
128    /// will not result in additional memory allocations.
129    pub fn to_shared(self) -> Self {
130        Self(match self.0 {
131            CowBytes::Shared(bytes) => CowBytes::Shared(bytes),
132            CowBytes::Owned(bytes) => CowBytes::Shared(bytes.freeze()),
133        })
134    }
135}
136
137#[cfg(feature = "parallel")]
138impl FlatPieces {
139    /// Parallel iterator over source pieces (even indices)
140    #[inline]
141    pub fn par_source(&self) -> impl IndexedParallelIterator<Item = &'_ PieceArray> + '_ {
142        self.par_iter()
143            .take(RecordedHistorySegment::NUM_RAW_RECORDS)
144    }
145
146    /// Mutable parallel iterator over source pieces (even indices)
147    #[inline]
148    pub fn par_source_mut(
149        &mut self,
150    ) -> impl IndexedParallelIterator<Item = &'_ mut PieceArray> + '_ {
151        self.par_iter_mut()
152            .take(RecordedHistorySegment::NUM_RAW_RECORDS)
153    }
154
155    /// Parallel iterator over parity pieces (odd indices)
156    #[inline]
157    pub fn par_parity(&self) -> impl IndexedParallelIterator<Item = &'_ PieceArray> + '_ {
158        self.par_iter()
159            .skip(RecordedHistorySegment::NUM_RAW_RECORDS)
160    }
161
162    /// Mutable parallel iterator over parity pieces (odd indices)
163    #[inline]
164    pub fn par_parity_mut(
165        &mut self,
166    ) -> impl IndexedParallelIterator<Item = &'_ mut PieceArray> + '_ {
167        self.par_iter_mut()
168            .skip(RecordedHistorySegment::NUM_RAW_RECORDS)
169    }
170}
ab_core_primitives/pieces/flat_pieces.rs

ab_core_primitives/pieces/
flat_pieces.rs
