ab_core_primitives/pieces/
flat_pieces.rs

1use crate::pieces::cow_bytes::CowBytes;
2use crate::pieces::{InnerPiece, Piece};
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 = [InnerPiece];
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; const { Piece::SIZE }]>(),
33                bytes.len() / Piece::SIZE,
34            )
35        };
36        InnerPiece::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; const { Piece::SIZE }]>(),
49                bytes.len() / Piece::SIZE,
50            )
51        };
52        InnerPiece::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                    .as_chunks::<const { Piece::SIZE }>()
76                    .0
77                    .iter()
78                    .map(|slice| Piece(CowBytes::Shared(bytes.slice_ref(slice)))),
79            ),
80            CowBytes::Owned(bytes) => Box::new(
81                bytes
82                    .as_chunks::<const { Piece::SIZE }>()
83                    .0
84                    .iter()
85                    .map(|slice| Piece(CowBytes::Shared(Bytes::copy_from_slice(slice)))),
86            ),
87        }
88    }
89
90    /// Iterator over source pieces (even indices)
91    #[inline]
92    pub fn source_pieces(&self) -> impl ExactSizeIterator<Item = Piece> + '_ {
93        self.pieces().take(RecordedHistorySegment::NUM_RAW_RECORDS)
94    }
95
96    /// Iterator over source pieces (even indices)
97    #[inline]
98    pub fn source(&self) -> impl ExactSizeIterator<Item = &'_ InnerPiece> + '_ {
99        self.iter().take(RecordedHistorySegment::NUM_RAW_RECORDS)
100    }
101
102    /// Mutable iterator over source pieces (even indices)
103    #[inline]
104    pub fn source_mut(&mut self) -> impl ExactSizeIterator<Item = &'_ mut InnerPiece> + '_ {
105        self.iter_mut()
106            .take(RecordedHistorySegment::NUM_RAW_RECORDS)
107    }
108
109    /// Iterator over parity pieces (odd indices)
110    #[inline]
111    pub fn parity_pieces(&self) -> impl ExactSizeIterator<Item = Piece> + '_ {
112        self.pieces().skip(RecordedHistorySegment::NUM_RAW_RECORDS)
113    }
114
115    /// Iterator over parity pieces (odd indices)
116    #[inline]
117    pub fn parity(&self) -> impl ExactSizeIterator<Item = &'_ InnerPiece> + '_ {
118        self.iter().skip(RecordedHistorySegment::NUM_RAW_RECORDS)
119    }
120
121    /// Mutable iterator over parity pieces (odd indices)
122    #[inline]
123    pub fn parity_mut(&mut self) -> impl ExactSizeIterator<Item = &'_ mut InnerPiece> + '_ {
124        self.iter_mut()
125            .skip(RecordedHistorySegment::NUM_RAW_RECORDS)
126    }
127
128    /// Ensure flat pieces contains cheaply cloneable shared data.
129    ///
130    /// Internally flat pieces uses CoW mechanism and can store either mutable owned data or data
131    /// that is cheap to clone, calling this method will ensure further clones and returned pieces
132    /// will not result in additional memory allocations.
133    pub fn to_shared(self) -> Self {
134        Self(match self.0 {
135            CowBytes::Shared(bytes) => CowBytes::Shared(bytes),
136            CowBytes::Owned(bytes) => CowBytes::Shared(bytes.freeze()),
137        })
138    }
139
140    /// Parallel iterator over source pieces (even indices)
141    #[inline]
142    #[cfg(feature = "parallel")]
143    pub fn par_source(&self) -> impl IndexedParallelIterator<Item = &'_ InnerPiece> + '_ {
144        self.par_iter()
145            .take(RecordedHistorySegment::NUM_RAW_RECORDS)
146    }
147
148    /// Mutable parallel iterator over source pieces (even indices)
149    #[inline]
150    #[cfg(feature = "parallel")]
151    pub fn par_source_mut(
152        &mut self,
153    ) -> impl IndexedParallelIterator<Item = &'_ mut InnerPiece> + '_ {
154        self.par_iter_mut()
155            .take(RecordedHistorySegment::NUM_RAW_RECORDS)
156    }
157
158    /// Parallel iterator over parity pieces (odd indices)
159    #[inline]
160    #[cfg(feature = "parallel")]
161    pub fn par_parity(&self) -> impl IndexedParallelIterator<Item = &'_ InnerPiece> + '_ {
162        self.par_iter()
163            .skip(RecordedHistorySegment::NUM_RAW_RECORDS)
164    }
165
166    /// Mutable parallel iterator over parity pieces (odd indices)
167    #[inline]
168    #[cfg(feature = "parallel")]
169    pub fn par_parity_mut(
170        &mut self,
171    ) -> impl IndexedParallelIterator<Item = &'_ mut InnerPiece> + '_ {
172        self.par_iter_mut()
173            .skip(RecordedHistorySegment::NUM_RAW_RECORDS)
174    }
175}
ab_core_primitives/pieces/flat_pieces.rs

ab_core_primitives/pieces/
flat_pieces.rs
