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

ab_core_primitives/block/body/
owned.rs
