ab_core_primitives/block/body/
owned.rs

1//! Data structures related to the owned version of [`BlockBody`]
2
3use crate::block::body::{
4    BeaconChainBody, BlockBody, GenericBlockBody, IntermediateShardBlockInfo,
5    IntermediateShardBody, LeafShardBlockInfo, LeafShardBody,
6};
7use crate::block::header::owned::{
8    OwnedIntermediateShardHeader, OwnedIntermediateShardHeaderError, OwnedLeafShardHeader,
9};
10use crate::pot::PotCheckpoints;
11use crate::segments::SegmentRoot;
12use crate::shard::ShardKind;
13use crate::transaction::Transaction;
14use crate::transaction::owned::{OwnedTransaction, OwnedTransactionError};
15use ab_aligned_buffer::{OwnedAlignedBuffer, SharedAlignedBuffer};
16use ab_io_type::trivial_type::TrivialType;
17use core::fmt;
18use core::iter::TrustedLen;
19use derive_more::From;
20use rclite::Arc;
21use yoke::Yoke;
22
23/// Generic owned block body
24pub trait GenericOwnedBlockBody: Clone + fmt::Debug + 'static {
25    /// Block body
26    type Body<'a>: GenericBlockBody<'a>
27    where
28        Self: 'a;
29
30    /// Get regular block body out of the owned version
31    fn body(&self) -> &Self::Body<'_>;
32}
33
34/// Transaction addition error
35#[derive(Debug, thiserror::Error)]
36enum AddTransactionError {
37    /// Block body is too large
38    #[error("Block body is too large")]
39    BlockBodyIsTooLarge,
40    /// Too many transactions
41    #[error("Too many transactions")]
42    TooManyTransactions,
43    /// Failed to add transaction
44    #[error("Failed to add transaction: {error}")]
45    FailedToAddTransaction {
46        /// Inner error
47        #[from]
48        error: OwnedTransactionError,
49    },
50}
51
52/// Transaction that can be written into the body
53pub trait WritableBodyTransaction {
54    /// Write this transaction into the body
55    fn write_into(&self, buffer: &mut OwnedAlignedBuffer) -> Result<(), OwnedTransactionError>;
56}
57
58impl WritableBodyTransaction for Transaction<'_> {
59    fn write_into(&self, buffer: &mut OwnedAlignedBuffer) -> Result<(), OwnedTransactionError> {
60        OwnedTransaction::from_parts_into(
61            self.header,
62            self.read_slots,
63            self.write_slots,
64            self.payload,
65            self.seal,
66            buffer,
67        )
68    }
69}
70
71impl WritableBodyTransaction for &OwnedTransaction {
72    fn write_into(&self, buffer: &mut OwnedAlignedBuffer) -> Result<(), OwnedTransactionError> {
73        if buffer.append(self.buffer().as_slice()) {
74            Ok(())
75        } else {
76            Err(OwnedTransactionError::TransactionTooLarge)
77        }
78    }
79}
80
81#[derive(Debug, Clone)]
82struct TransactionBuilder {
83    num_transactions_offset: usize,
84    buffer: OwnedAlignedBuffer,
85}
86
87impl TransactionBuilder {
88    fn new(num_transactions_offset: usize, buffer: OwnedAlignedBuffer) -> Self {
89        Self {
90            num_transactions_offset,
91            buffer,
92        }
93    }
94
95    /// Add transaction to the body
96    fn add_transaction<T>(&mut self, transaction: T) -> Result<(), AddTransactionError>
97    where
98        T: WritableBodyTransaction,
99    {
100        // Transactions are aligned, but the very first might come after non-transaction fields that
101        // were not aligned
102        if self.inc_transaction_count()? == 1 && !align_to_16_bytes_with_padding(&mut self.buffer) {
103            self.dec_transaction_count();
104            return Err(AddTransactionError::BlockBodyIsTooLarge);
105        }
106
107        let old_buffer_len = self.buffer.len();
108
109        transaction
110            .write_into(&mut self.buffer)
111            .inspect_err(|_error| {
112                self.dec_transaction_count();
113            })?;
114
115        if !align_to_16_bytes_with_padding(&mut self.buffer) {
116            self.dec_transaction_count();
117            // Length was obtained from the same buffer before last write
118            unsafe {
119                self.buffer.set_len(old_buffer_len);
120            }
121            return Err(AddTransactionError::BlockBodyIsTooLarge);
122        }
123
124        Ok(())
125    }
126
127    /// Finish building block body
128    #[inline(always)]
129    fn finish(self) -> OwnedAlignedBuffer {
130        self.buffer
131    }
132
133    /// Increase the number of stored transactions and return the new value
134    #[inline(always)]
135    fn inc_transaction_count(&mut self) -> Result<u32, AddTransactionError> {
136        // SAFETY: Constructor ensures the offset is valid and has space for `u32` (but not
137        // necessarily aligned)
138        unsafe {
139            let num_transactions_ptr = self
140                .buffer
141                .as_mut_ptr()
142                .add(self.num_transactions_offset)
143                .cast::<u32>();
144            let num_transactions = num_transactions_ptr.read_unaligned();
145            let num_transactions = num_transactions
146                .checked_add(1)
147                .ok_or(AddTransactionError::TooManyTransactions)?;
148            num_transactions_ptr.write_unaligned(num_transactions);
149            Ok(num_transactions)
150        }
151    }
152
153    /// Decrease the number of stored transactions
154    #[inline(always)]
155    fn dec_transaction_count(&mut self) {
156        // SAFETY: Constructor ensures the offset is valid and has space for `u32` (but not
157        // necessarily aligned)
158        unsafe {
159            let num_transactions_ptr = self
160                .buffer
161                .as_mut_ptr()
162                .add(self.num_transactions_offset)
163                .cast::<u32>();
164            let num_transactions = num_transactions_ptr.read_unaligned();
165            let num_transactions = num_transactions.saturating_sub(1);
166            num_transactions_ptr.write_unaligned(num_transactions);
167        }
168    }
169}
170
171/// Errors for [`OwnedBeaconChainBody`]
172#[derive(Debug, thiserror::Error)]
173pub enum OwnedBeaconChainBodyError {
174    /// Too many PoT checkpoints
175    #[error("Too many PoT checkpoints: {actual}")]
176    TooManyPotCheckpoints {
177        /// Actual number of PoT checkpoints
178        actual: usize,
179    },
180    /// Too many own segment roots
181    #[error("Too many own segment roots: {actual}")]
182    TooManyOwnSegmentRoots {
183        /// Actual number of own segment roots
184        actual: usize,
185    },
186    /// Too many intermediate shard blocks
187    #[error("Too many intermediate shard blocks: {actual}")]
188    TooManyIntermediateShardBlocks {
189        /// Actual number of intermediate shard blocks
190        actual: usize,
191    },
192    /// Too many intermediate shard own segment roots
193    #[error("Too many intermediate shard own segment roots: {actual}")]
194    TooManyIntermediateShardOwnSegmentRoots {
195        /// Actual number of own segment roots
196        actual: usize,
197    },
198    /// Too many intermediate shard child segment roots
199    #[error("Too many intermediate shard child segment roots: {actual}")]
200    TooManyIntermediateShardChildSegmentRoots {
201        /// Actual number of child segment roots
202        actual: usize,
203    },
204    /// Failed to intermediate shard header
205    #[error("Failed to intermediate shard header: {error}")]
206    FailedToAddIntermediateShard {
207        /// Inner error
208        #[from]
209        error: OwnedIntermediateShardHeaderError,
210    },
211    /// Block body is too large
212    #[error("Block body is too large")]
213    BlockBodyIsTooLarge,
214}
215
216/// An owned version of [`BeaconChainBody`].
217///
218/// It is correctly aligned in memory and well suited for sending and receiving over the network
219/// efficiently or storing in memory or on disk.
220#[derive(Debug, Clone)]
221pub struct OwnedBeaconChainBody {
222    inner: Arc<Yoke<BeaconChainBody<'static>, SharedAlignedBuffer>>,
223}
224
225impl GenericOwnedBlockBody for OwnedBeaconChainBody {
226    type Body<'a> = BeaconChainBody<'a>;
227
228    #[inline(always)]
229    fn body(&self) -> &Self::Body<'_> {
230        self.body()
231    }
232}
233
234impl OwnedBeaconChainBody {
235    /// Create a new instance
236    pub fn new<'a, OSR, ISB>(
237        own_segment_roots: OSR,
238        intermediate_shard_blocks: ISB,
239        pot_checkpoints: &[PotCheckpoints],
240    ) -> Result<Self, OwnedBeaconChainBodyError>
241    where
242        OSR: TrustedLen<Item = SegmentRoot>,
243        ISB: TrustedLen<Item = IntermediateShardBlockInfo<'a>> + Clone + 'a,
244    {
245        let num_pot_checkpoints = pot_checkpoints.len();
246        let num_pot_checkpoints = u32::try_from(num_pot_checkpoints).map_err(|_error| {
247            OwnedBeaconChainBodyError::TooManyPotCheckpoints {
248                actual: num_pot_checkpoints,
249            }
250        })?;
251        let num_own_segment_roots = own_segment_roots.size_hint().0;
252        let num_own_segment_roots = u8::try_from(num_own_segment_roots).map_err(|_error| {
253            OwnedBeaconChainBodyError::TooManyOwnSegmentRoots {
254                actual: num_own_segment_roots,
255            }
256        })?;
257        let num_blocks = intermediate_shard_blocks.size_hint().0;
258        let num_blocks = u8::try_from(num_blocks).map_err(|_error| {
259            OwnedBeaconChainBodyError::TooManyIntermediateShardBlocks { actual: num_blocks }
260        })?;
261
262        let mut buffer = OwnedAlignedBuffer::with_capacity(
263            u8::SIZE
264                + u32::from(num_own_segment_roots) * SegmentRoot::SIZE as u32
265                // This is only an estimate to get in the ballpark where reallocation should not be
266                // necessary in many cases
267                + u32::from(num_blocks) * OwnedIntermediateShardHeader::max_allocation_for(&[]) * 2,
268        );
269
270        let true = buffer.append(&num_pot_checkpoints.to_le_bytes()) else {
271            unreachable!("Fixed size data structures that are guaranteed to fit; qed");
272        };
273
274        let true = buffer.append(&[num_own_segment_roots]) else {
275            unreachable!("Fixed size data structures that are guaranteed to fit; qed");
276        };
277        for own_segment_root in own_segment_roots {
278            let true = buffer.append(own_segment_root.as_ref()) else {
279                unreachable!("Checked size above; qed");
280            };
281        }
282        // TODO: Would be nice for `IntermediateShardBlocksInfo` to have API to write this by itself
283        {
284            let true = buffer.append(&num_blocks.to_le_bytes()) else {
285                unreachable!("Fixed size data structures that are guaranteed to fit; qed");
286            };
287            let mut segments_roots_num_cursor = buffer.len() as usize;
288            for _ in 0..num_blocks {
289                let true = buffer.append(&[0, 0, 0]) else {
290                    unreachable!("Checked size above; qed");
291                };
292            }
293            let true = align_to_8_with_padding(&mut buffer) else {
294                unreachable!("Checked size above; qed");
295            };
296            for intermediate_shard_block in intermediate_shard_blocks.clone() {
297                if !intermediate_shard_block.own_segment_roots.is_empty()
298                    || !intermediate_shard_block.child_segment_roots.is_empty()
299                {
300                    let num_own_segment_roots = intermediate_shard_block.own_segment_roots.len();
301                    let num_own_segment_roots =
302                        u8::try_from(num_own_segment_roots).map_err(|_error| {
303                            OwnedBeaconChainBodyError::TooManyIntermediateShardOwnSegmentRoots {
304                                actual: num_own_segment_roots,
305                            }
306                        })?;
307                    let num_child_segment_roots =
308                        intermediate_shard_block.child_segment_roots.len();
309                    let num_child_segment_roots =
310                        u16::try_from(num_child_segment_roots).map_err(|_error| {
311                            OwnedBeaconChainBodyError::TooManyIntermediateShardChildSegmentRoots {
312                                actual: num_child_segment_roots,
313                            }
314                        })?;
315                    let num_child_segment_roots = num_child_segment_roots.to_le_bytes();
316                    buffer.as_mut_slice()[segments_roots_num_cursor..][..3].copy_from_slice(&[
317                        num_own_segment_roots,
318                        num_child_segment_roots[0],
319                        num_child_segment_roots[1],
320                    ]);
321                }
322                segments_roots_num_cursor += 3;
323
324                OwnedIntermediateShardHeader::from_parts_into(
325                    intermediate_shard_block.header.prefix,
326                    intermediate_shard_block.header.result,
327                    intermediate_shard_block.header.consensus_info,
328                    intermediate_shard_block.header.beacon_chain_info(),
329                    intermediate_shard_block.header.child_shard_blocks(),
330                    &mut buffer,
331                )?;
332                if !align_to_8_with_padding(&mut buffer) {
333                    return Err(OwnedBeaconChainBodyError::BlockBodyIsTooLarge);
334                }
335                if let Some(segment_roots_proof) = intermediate_shard_block.segment_roots_proof
336                    && !buffer.append(segment_roots_proof)
337                {
338                    return Err(OwnedBeaconChainBodyError::BlockBodyIsTooLarge);
339                }
340                if !intermediate_shard_block.own_segment_roots.is_empty()
341                    && !buffer.append(
342                        SegmentRoot::repr_from_slice(intermediate_shard_block.own_segment_roots)
343                            .as_flattened(),
344                    )
345                {
346                    return Err(OwnedBeaconChainBodyError::BlockBodyIsTooLarge);
347                }
348                if !intermediate_shard_block.child_segment_roots.is_empty()
349                    && !buffer.append(
350                        SegmentRoot::repr_from_slice(intermediate_shard_block.child_segment_roots)
351                            .as_flattened(),
352                    )
353                {
354                    return Err(OwnedBeaconChainBodyError::BlockBodyIsTooLarge);
355                }
356            }
357        }
358
359        let true = buffer.append(PotCheckpoints::bytes_from_slice(pot_checkpoints).as_flattened())
360        else {
361            return Err(OwnedBeaconChainBodyError::BlockBodyIsTooLarge);
362        };
363
364        // TODO: Avoid extra parsing here or at least go through unchecked version
365        Ok(Self::from_buffer(buffer.into_shared()).expect("Known to be created correctly; qed"))
366    }
367
368    /// Create owned body from a buffer
369    #[inline]
370    pub fn from_buffer(buffer: SharedAlignedBuffer) -> Result<Self, SharedAlignedBuffer> {
371        // TODO: Cloning is cheap, but will not be necessary if/when this is resolved:
372        //  https://github.com/unicode-org/icu4x/issues/6665
373        let inner = Yoke::try_attach_to_cart(buffer.clone(), |buffer| {
374            let Some((body, extra_bytes)) = BeaconChainBody::try_from_bytes(buffer) else {
375                return Err(());
376            };
377            if !extra_bytes.is_empty() {
378                return Err(());
379            }
380
381            Ok(body)
382        })
383        .map_err(move |()| buffer)?;
384
385        Ok(Self {
386            inner: Arc::new(inner),
387        })
388    }
389
390    /// Inner buffer with block body contents
391    #[inline(always)]
392    pub fn buffer(&self) -> &SharedAlignedBuffer {
393        self.inner.backing_cart()
394    }
395
396    /// Get [`BeaconChainBody`] out of [`OwnedBeaconChainBody`]
397    #[inline(always)]
398    pub fn body(&self) -> &BeaconChainBody<'_> {
399        self.inner.get()
400    }
401}
402
403/// Errors for [`OwnedIntermediateShardBody`]
404#[derive(Debug, thiserror::Error)]
405pub enum OwnedIntermediateShardBodyError {
406    /// Too many own segment roots
407    #[error("Too many own segment roots: {actual}")]
408    TooManyOwnSegmentRoots {
409        /// Actual number of own segment roots
410        actual: usize,
411    },
412    /// Too many leaf shard blocks
413    #[error("Too many leaf shard blocks: {actual}")]
414    TooManyLeafShardBlocks {
415        /// Actual number of leaf shard blocks
416        actual: usize,
417    },
418    /// Too many leaf shard own segment roots
419    #[error("Too many leaf shard own segment roots: {actual}")]
420    TooManyLeafShardOwnSegmentRoots {
421        /// Actual number of own segment roots
422        actual: usize,
423    },
424}
425
426/// An owned version of [`IntermediateShardBody`].
427///
428/// It is correctly aligned in memory and well suited for sending and receiving over the network
429/// efficiently or storing in memory or on disk.
430#[derive(Debug, Clone)]
431pub struct OwnedIntermediateShardBody {
432    inner: Arc<Yoke<IntermediateShardBody<'static>, SharedAlignedBuffer>>,
433}
434
435impl GenericOwnedBlockBody for OwnedIntermediateShardBody {
436    type Body<'a> = IntermediateShardBody<'a>;
437
438    #[inline(always)]
439    fn body(&self) -> &Self::Body<'_> {
440        self.body()
441    }
442}
443
444impl OwnedIntermediateShardBody {
445    /// Create a new instance
446    pub fn new<'a, OSR, LSB>(
447        own_segment_roots: OSR,
448        leaf_shard_blocks: LSB,
449    ) -> Result<Self, OwnedIntermediateShardBodyError>
450    where
451        OSR: TrustedLen<Item = SegmentRoot>,
452        LSB: TrustedLen<Item = LeafShardBlockInfo<'a>> + Clone + 'a,
453    {
454        let num_own_segment_roots = own_segment_roots.size_hint().0;
455        let num_own_segment_roots = u8::try_from(num_own_segment_roots).map_err(|_error| {
456            OwnedIntermediateShardBodyError::TooManyOwnSegmentRoots {
457                actual: num_own_segment_roots,
458            }
459        })?;
460        let num_blocks = leaf_shard_blocks.size_hint().0;
461        let num_blocks = u8::try_from(num_blocks).map_err(|_error| {
462            OwnedIntermediateShardBodyError::TooManyLeafShardBlocks { actual: num_blocks }
463        })?;
464
465        let mut buffer = OwnedAlignedBuffer::with_capacity(
466            u8::SIZE
467                + u32::from(num_own_segment_roots) * SegmentRoot::SIZE as u32
468                // This is only an estimate to get in the ballpark where reallocation should not be
469                // necessary
470                + u32::from(num_blocks) * OwnedLeafShardHeader::MAX_ALLOCATION * 2,
471        );
472
473        let true = buffer.append(&[num_own_segment_roots]) else {
474            unreachable!("Fixed size data structures that are guaranteed to fit; qed");
475        };
476        for own_segment_root in own_segment_roots {
477            let true = buffer.append(own_segment_root.as_ref()) else {
478                unreachable!("Checked size above; qed");
479            };
480        }
481        // TODO: Would be nice for `LeafShardBlocksInfo` to have API to write this by itself
482        {
483            let true = buffer.append(&num_blocks.to_le_bytes()) else {
484                unreachable!("Fixed size data structures that are guaranteed to fit; qed");
485            };
486            let mut own_segments_roots_num_cursor = buffer.len() as usize;
487            for _ in 0..num_blocks {
488                let true = buffer.append(&[0]) else {
489                    unreachable!("Checked size above; qed");
490                };
491            }
492            let true = align_to_8_with_padding(&mut buffer) else {
493                unreachable!("Checked size above; qed");
494            };
495            for leaf_shard_block in leaf_shard_blocks.clone() {
496                if !leaf_shard_block.own_segment_roots.is_empty() {
497                    let num_own_segment_roots = leaf_shard_block.own_segment_roots.len();
498                    let num_own_segment_roots =
499                        u8::try_from(num_own_segment_roots).map_err(|_error| {
500                            OwnedIntermediateShardBodyError::TooManyLeafShardOwnSegmentRoots {
501                                actual: num_own_segment_roots,
502                            }
503                        })?;
504                    buffer.as_mut_slice()[own_segments_roots_num_cursor] = num_own_segment_roots;
505                }
506                own_segments_roots_num_cursor += 1;
507
508                OwnedLeafShardHeader::from_parts_into(
509                    leaf_shard_block.header.prefix,
510                    leaf_shard_block.header.result,
511                    leaf_shard_block.header.consensus_info,
512                    leaf_shard_block.header.beacon_chain_info(),
513                    &mut buffer,
514                );
515                let true = align_to_8_with_padding(&mut buffer) else {
516                    unreachable!("Checked size above; qed");
517                };
518                if let Some(segment_roots_proof) = leaf_shard_block.segment_roots_proof {
519                    let true = buffer.append(segment_roots_proof) else {
520                        unreachable!("Checked size above; qed");
521                    };
522                }
523                if !leaf_shard_block.own_segment_roots.is_empty() {
524                    let true = buffer.append(
525                        SegmentRoot::repr_from_slice(leaf_shard_block.own_segment_roots)
526                            .as_flattened(),
527                    ) else {
528                        unreachable!("Checked size above; qed");
529                    };
530                }
531            }
532        }
533
534        // TODO: Avoid extra parsing here or at least go through unchecked version
535        Ok(Self::from_buffer(buffer.into_shared()).expect("Known to be created correctly; qed"))
536    }
537
538    /// Create owned body from a buffer
539    #[inline]
540    pub fn from_buffer(buffer: SharedAlignedBuffer) -> Result<Self, SharedAlignedBuffer> {
541        // TODO: Cloning is cheap, but will not be necessary if/when this is resolved:
542        //  https://github.com/unicode-org/icu4x/issues/6665
543        let inner = Yoke::try_attach_to_cart(buffer.clone(), |buffer| {
544            let Some((body, extra_bytes)) = IntermediateShardBody::try_from_bytes(buffer) else {
545                return Err(());
546            };
547            if !extra_bytes.is_empty() {
548                return Err(());
549            }
550
551            Ok(body)
552        })
553        .map_err(move |()| buffer)?;
554
555        Ok(Self {
556            inner: Arc::new(inner),
557        })
558    }
559
560    /// Inner buffer with block body contents
561    #[inline(always)]
562    pub fn buffer(&self) -> &SharedAlignedBuffer {
563        self.inner.backing_cart()
564    }
565
566    /// Get [`IntermediateShardBody`] out of [`OwnedIntermediateShardBody`]
567    #[inline(always)]
568    pub fn body(&self) -> &IntermediateShardBody<'_> {
569        self.inner.get()
570    }
571}
572
573/// Errors for [`OwnedLeafShardBody`]
574#[derive(Debug, thiserror::Error)]
575pub enum OwnedLeafShardBodyError {
576    /// Too many own segment roots
577    #[error("Too many own segment roots: {actual}")]
578    TooManyOwnSegmentRoots {
579        /// Actual number of own segment roots
580        actual: usize,
581    },
582    /// Block body is too large
583    #[error("Block body is too large")]
584    BlockBodyIsTooLarge,
585    /// Too many transactions
586    #[error("Too many transactions")]
587    TooManyTransactions,
588    /// Failed to add transaction
589    #[error("Failed to add transaction: {error}")]
590    FailedToAddTransaction {
591        /// Inner error
592        error: OwnedTransactionError,
593    },
594}
595
596impl From<AddTransactionError> for OwnedLeafShardBodyError {
597    fn from(value: AddTransactionError) -> Self {
598        match value {
599            AddTransactionError::BlockBodyIsTooLarge => {
600                OwnedLeafShardBodyError::BlockBodyIsTooLarge
601            }
602            AddTransactionError::TooManyTransactions => {
603                OwnedLeafShardBodyError::TooManyTransactions
604            }
605            AddTransactionError::FailedToAddTransaction { error } => {
606                OwnedLeafShardBodyError::FailedToAddTransaction { error }
607            }
608        }
609    }
610}
611
612/// An owned version of [`LeafShardBody`].
613///
614/// It is correctly aligned in memory and well suited for sending and receiving over the network
615/// efficiently or storing in memory or on disk.
616#[derive(Debug, Clone)]
617pub struct OwnedLeafShardBody {
618    inner: Arc<Yoke<LeafShardBody<'static>, SharedAlignedBuffer>>,
619}
620
621impl GenericOwnedBlockBody for OwnedLeafShardBody {
622    type Body<'a> = LeafShardBody<'a>;
623
624    #[inline(always)]
625    fn body(&self) -> &Self::Body<'_> {
626        self.body()
627    }
628}
629
630impl OwnedLeafShardBody {
631    /// Initialize building of [`OwnedLeafShardBody`]
632    pub fn init<OSR>(
633        own_segment_roots: OSR,
634    ) -> Result<OwnedLeafShardBlockBodyBuilder, OwnedLeafShardBodyError>
635    where
636        OSR: TrustedLen<Item = SegmentRoot>,
637    {
638        let num_own_segment_roots = own_segment_roots.size_hint().0;
639        let num_own_segment_roots = u8::try_from(num_own_segment_roots).map_err(|_error| {
640            OwnedLeafShardBodyError::TooManyOwnSegmentRoots {
641                actual: num_own_segment_roots,
642            }
643        })?;
644
645        let mut buffer = OwnedAlignedBuffer::with_capacity(
646            u8::SIZE + u32::from(num_own_segment_roots) * SegmentRoot::SIZE as u32,
647        );
648
649        let true = buffer.append(&[num_own_segment_roots]) else {
650            unreachable!("Fixed size data structures that are guaranteed to fit; qed");
651        };
652        for own_segment_root in own_segment_roots {
653            let true = buffer.append(own_segment_root.as_ref()) else {
654                unreachable!("Checked size above; qed");
655            };
656        }
657
658        let num_transactions_offset = buffer.len() as usize;
659        let true = buffer.append(&0u32.to_le_bytes()) else {
660            unreachable!("Checked size above; qed");
661        };
662
663        Ok(OwnedLeafShardBlockBodyBuilder {
664            transaction_builder: TransactionBuilder::new(num_transactions_offset, buffer),
665        })
666    }
667
668    /// Create owned body from a buffer
669    #[inline]
670    pub fn from_buffer(buffer: SharedAlignedBuffer) -> Result<Self, SharedAlignedBuffer> {
671        // TODO: Cloning is cheap, but will not be necessary if/when this is resolved:
672        //  https://github.com/unicode-org/icu4x/issues/6665
673        let inner = Yoke::try_attach_to_cart(buffer.clone(), |buffer| {
674            let Some((body, extra_bytes)) = LeafShardBody::try_from_bytes(buffer) else {
675                return Err(());
676            };
677            if !extra_bytes.is_empty() {
678                return Err(());
679            }
680
681            Ok(body)
682        })
683        .map_err(move |()| buffer)?;
684
685        Ok(Self {
686            inner: Arc::new(inner),
687        })
688    }
689
690    /// Inner buffer with block body contents
691    #[inline(always)]
692    pub fn buffer(&self) -> &SharedAlignedBuffer {
693        self.inner.backing_cart()
694    }
695
696    /// Get [`LeafShardBody`] out of [`OwnedLeafShardBody`]
697    #[inline(always)]
698    pub fn body(&self) -> &LeafShardBody<'_> {
699        self.inner.get()
700    }
701}
702
703/// Builder for [`OwnedLeafShardBody`] that allows to add more transactions
704#[derive(Debug, Clone)]
705pub struct OwnedLeafShardBlockBodyBuilder {
706    transaction_builder: TransactionBuilder,
707}
708
709impl OwnedLeafShardBlockBodyBuilder {
710    /// Add transaction to the body
711    #[inline(always)]
712    pub fn add_transaction<T>(&mut self, transaction: T) -> Result<(), OwnedLeafShardBodyError>
713    where
714        T: WritableBodyTransaction,
715    {
716        self.transaction_builder.add_transaction(transaction)?;
717
718        Ok(())
719    }
720
721    /// Finish building block body
722    pub fn finish(self) -> OwnedLeafShardBody {
723        // TODO: Avoid extra parsing here or at least go through unchecked version
724        OwnedLeafShardBody::from_buffer(self.transaction_builder.finish().into_shared())
725            .expect("Known to be created correctly; qed")
726    }
727}
728
729/// An owned version of [`BlockBody`].
730///
731/// It is correctly aligned in memory and well suited for sending and receiving over the network
732/// efficiently or storing in memory or on disk.
733#[derive(Debug, Clone, From)]
734pub enum OwnedBlockBody {
735    /// Block body corresponds to the beacon chain
736    BeaconChain(OwnedBeaconChainBody),
737    /// Block body corresponds to an intermediate shard
738    IntermediateShard(OwnedIntermediateShardBody),
739    /// Block body corresponds to a leaf shard
740    LeafShard(OwnedLeafShardBody),
741}
742
743impl OwnedBlockBody {
744    /// Create owned body from a buffer
745    #[inline]
746    pub fn from_buffer(
747        buffer: SharedAlignedBuffer,
748        shard_kind: ShardKind,
749    ) -> Result<Self, SharedAlignedBuffer> {
750        Ok(match shard_kind {
751            ShardKind::BeaconChain => Self::BeaconChain(OwnedBeaconChainBody::from_buffer(buffer)?),
752            ShardKind::IntermediateShard => {
753                Self::IntermediateShard(OwnedIntermediateShardBody::from_buffer(buffer)?)
754            }
755            ShardKind::LeafShard => Self::LeafShard(OwnedLeafShardBody::from_buffer(buffer)?),
756            ShardKind::Phantom | ShardKind::Invalid => {
757                // Blocks for such shards do not exist
758                return Err(buffer);
759            }
760        })
761    }
762
763    /// Inner buffer block body contents
764    #[inline]
765    pub fn buffer(&self) -> &SharedAlignedBuffer {
766        match self {
767            Self::BeaconChain(owned_body) => owned_body.buffer(),
768            Self::IntermediateShard(owned_body) => owned_body.buffer(),
769            Self::LeafShard(owned_body) => owned_body.buffer(),
770        }
771    }
772
773    /// Get [`BlockBody`] out of [`OwnedBlockBody`]
774    #[inline]
775    pub fn body(&self) -> BlockBody<'_> {
776        match self {
777            Self::BeaconChain(owned_body) => BlockBody::BeaconChain(*owned_body.body()),
778            Self::IntermediateShard(owned_body) => BlockBody::IntermediateShard(*owned_body.body()),
779            Self::LeafShard(owned_body) => BlockBody::LeafShard(*owned_body.body()),
780        }
781    }
782}
783
784/// Aligns buffer to 8 bytes by adding necessary padding zero bytes.
785///
786/// Returns `false` if buffer becomes too long.
787#[inline(always)]
788#[must_use]
789fn align_to_8_with_padding(buffer: &mut OwnedAlignedBuffer) -> bool {
790    let alignment = align_of::<u64>();
791    // Optimized version of the following due to alignment being a power of 2:
792    // let unaligned_by = self.payload.as_ptr().addr() % alignment;
793    let unaligned_by = buffer.as_ptr().addr() & (alignment - 1);
794    if unaligned_by > 0 {
795        // SAFETY: Subtracted value is always smaller than alignment
796        let padding_bytes = unsafe { alignment.unchecked_sub(unaligned_by) };
797
798        if !buffer.append(&0u64.to_le_bytes()[..padding_bytes]) {
799            return false;
800        }
801    }
802
803    true
804}
805
806/// Aligns buffer to 16 bytes by adding necessary padding zero bytes.
807///
808/// Returns `false` if buffer becomes too long.
809#[inline(always)]
810#[must_use]
811fn align_to_16_bytes_with_padding(buffer: &mut OwnedAlignedBuffer) -> bool {
812    let alignment = align_of::<u128>();
813    // Optimized version of the following due to alignment being a power of 2:
814    // let unaligned_by = self.payload.as_ptr().addr() % alignment;
815    let unaligned_by = buffer.as_ptr().addr() & (alignment - 1);
816    if unaligned_by > 0 {
817        // SAFETY: Subtracted value is always smaller than alignment
818        let padding_bytes = unsafe { alignment.unchecked_sub(unaligned_by) };
819
820        if !buffer.append(&0u128.to_le_bytes()[..padding_bytes]) {
821            return false;
822        }
823    }
824
825    true
826}
ab_core_primitives/block/body/owned.rs

ab_core_primitives/block/body/
owned.rs
