shredder introduce shred builder / new api

ledger/src/shredder.rs

@@ -1,15 +1,19 @@
+#![allow(unused, dead_code)]
+
 use {
     crate::shred::{
         self, Error, ProcessShredsStats, Shred, ShredData, ShredFlags, DATA_SHREDS_PER_FEC_BLOCK,
     },
     lazy_lru::LruCache,
+    rand::Rng,
     rayon::ThreadPool,
     reed_solomon_erasure::{galois_8::ReedSolomon, Error::TooFewDataShards},
     solana_clock::Slot,
     solana_entry::entry::Entry,
     solana_hash::Hash,
     solana_keypair::Keypair,
     solana_rayon_threadlimit::get_thread_count,
+    solana_transaction::Transaction,
     std::{
         fmt::Debug,
         sync::{Arc, OnceLock, RwLock},
@@ -25,6 +29,197 @@ static PAR_THREAD_POOL: std::sync::LazyLock<ThreadPool> = std::sync::LazyLock::n
         .unwrap()
 });
 
+/// Initial state of [`ShredBuilder`]
+pub struct ShredBuilderInit;
+
+/// Generic data filler interface
+pub trait ShredBuilderDataFiller {
+    fn data(&mut self) -> Vec<u8>;
+}
+
+/// Represent state type of [`ShredBuilder`] with randomly generated bytes
+pub struct ShredBuilderRandomData {
+    data: Vec<u8>,
+}
+impl ShredBuilderDataFiller for ShredBuilderRandomData {
+    fn data(&mut self) -> Vec<u8> {
+        std::mem::take(&mut self.data)
+    }
+}
+
+/// Represent state type of [`ShredBuilder`] with provided bytes
+pub struct ShredBuilderBytesData {
+    data: Vec<u8>,
+}
+impl ShredBuilderDataFiller for ShredBuilderBytesData {
+    fn data(&mut self) -> Vec<u8> {
+        std::mem::take(&mut self.data)
+    }
+}
+
+/// Represent state type of [`ShredBuilder`] with serialized data from provided transactions
+pub struct ShredBuilderTransactionsData {
+    data: Vec<u8>,
+}
+impl ShredBuilderDataFiller for ShredBuilderTransactionsData {
+    fn data(&mut self) -> Vec<u8> {
+        std::mem::take(&mut self.data)
+    }
+}
+
+/// Generic type stated shred builder
+pub struct ShredBuilder<S> {
+    state: S,
+    invalid_index: Option<u32>,
+
+    slot: Slot,
+    chained_merkle_root: Hash,
+    parent_slot: Option<Slot>,
+    version: Option<u16>,
+    reference_tick: Option<u8>,
+    start_index: Option<u32>,
+}
+
+impl ShredBuilder<ShredBuilderInit> {
+    pub fn new(slot: Slot, hash: Hash) -> Self {
+        Self {
+            state: ShredBuilderInit,
+            invalid_index: None,
+
+            slot,
+            chained_merkle_root: hash,
+            parent_slot: None,
+            version: None,
+            reference_tick: None,
+            start_index: None,
+        }
+    }
+
+    /// Set state with, no data and returns specified variant
+    fn set_state<T>(&mut self, state: T) -> ShredBuilder<T> {
+        ShredBuilder {
+            state,
+            invalid_index: self.invalid_index,
+
+            slot: self.slot,
+            chained_merkle_root: self.chained_merkle_root,
+            parent_slot: self.parent_slot,
+            version: self.version,
+            reference_tick: self.reference_tick,
+            start_index: self.start_index,
+        }
+    }
+
+    /// Progress state into variant with shreds from randomly generated data
+    pub fn with_random_bytes(mut self, len: usize) -> ShredBuilder<ShredBuilderRandomData> {
+        let mut data = vec![0u8; len];
+        rand::thread_rng().fill(&mut data[..]);
+        self.set_state(ShredBuilderRandomData { data })
+    }
+
+    /// Progress state into variant with shreds from provided slice of bytes
+    pub fn with_bytes<D>(mut self, data: D) -> ShredBuilder<ShredBuilderBytesData>
+    where
+        D: AsRef<[u8]>,
+    {
+        self.set_state(ShredBuilderBytesData {
+            data: data.as_ref().to_vec(), // FIXME: maybe something more straight forward
+        })
+    }
+
+    /// Progress state into variant with shreds from provided transactions iterator
+    pub fn with_transactions<I>(
+        mut self,
+        transactions: I,
+    ) -> ShredBuilder<ShredBuilderTransactionsData>
+    where
+        I: IntoIterator<Item = Transaction>,
+    {
+        let transactions: Vec<Transaction> = transactions.into_iter().collect();
+        self.set_state(ShredBuilderTransactionsData {
+            data: bincode::serialize(&Entry::new(
+                &Hash::default(),
+                transactions.len().try_into().unwrap(),
+                transactions,
+            ))
+            .unwrap(),
+        })
+    }
+}
+
+impl<S> ShredBuilder<S>
+where
+    S: ShredBuilderDataFiller,
+{
+    /// Set parent slot (default is slot - 1)
+    pub fn with_parent_slot(mut self, parent_slot: Slot) -> Self {
+        self.parent_slot = Some(parent_slot);
+        self
+    }
+
+    /// Set version (default is 0)
+    pub fn with_version(mut self, version: u16) -> Self {
+        self.version = Some(version);
+        self
+    }
+
+    /// Set reference tick (default is 0)
+    pub fn with_reference_tick(mut self, reference_tick: u8) -> Self {
+        self.reference_tick = Some(reference_tick);
+        self
+    }
+
+    /// Set start index (default is 0)
+    pub fn with_start_index(mut self, index: u32) -> Self {
+        self.start_index = Some(index);
+        self
+    }
+
+    /// Set next index (default is the same as start index)
+    pub fn with_invalid_index(mut self, index: u32) -> Self {
+        self.invalid_index = Some(index);
+        self
+    }
+
+    /// Build shred iterator from state
+    pub fn build(
+        mut self,
+        is_last_in_slot: bool,
+    ) -> Result<impl Iterator<Item = crate::shred::merkle::Shred>, crate::shred::Error> {
+        let mut stats = ProcessShredsStats::default();
+        let reed_solomon_cache = ReedSolomonCache::default();
+        let data = self.state.data();
+        let parent_slot = self
+            .parent_slot
+            .unwrap_or_else(|| self.slot.saturating_sub(1));
+        let version = self.version.unwrap_or_default();
+        let reference_tick = self.reference_tick.unwrap_or_default();
+        let start_index = self.start_index.unwrap_or_default();
+        let next_index = match self.invalid_index {
+            Some(invalid_index) => invalid_index,
+            None => start_index,
+        };
+
+        let shreds = crate::shred::merkle::make_shreds_from_data(
+            &PAR_THREAD_POOL,
+            &Keypair::new(),
+            Some(self.chained_merkle_root),
+            &data,
+            self.slot,
+            parent_slot,
+            version,
+            reference_tick,
+            is_last_in_slot,
+            start_index,
+            next_index,
+            &reed_solomon_cache,
+            &mut stats,
+        )?;
+
+        Ok(shreds.into_iter())
+    }
+}
+
 // Arc<...> wrapper so that cache entries can be initialized without locking
 // the entire cache.
 type LruCacheOnce<K, V> = RwLock<LruCache<K, Arc<OnceLock<V>>>>;
@@ -258,14 +453,15 @@ mod tests {
         assert_matches::assert_matches,
         itertools::Itertools,
         rand::Rng,
+        solana_entry::entry::Entry,
         solana_hash::Hash,
         solana_pubkey::Pubkey,
         solana_sha256_hasher::hash,
         solana_shred_version as shred_version,
         solana_signer::Signer,
         solana_system_transaction as system_transaction,
         std::{collections::HashSet, sync::Arc},
-        test_case::test_matrix,
+        test_case::{test_case, test_matrix},
     };
 
     fn verify_test_code_shred(shred: &Shred, index: u32, slot: Slot, pk: &Pubkey, verify: bool) {