diff --git a/bin/strata-client/src/extractor.rs b/bin/strata-client/src/extractor.rs index 09fe237dd..3450d9804 100644 --- a/bin/strata-client/src/extractor.rs +++ b/bin/strata-client/src/extractor.rs @@ -337,7 +337,7 @@ mod tests { max_txs_per_block: usize, needle: (Vec, ProtocolOperation), ) -> usize { - let arb = ArbitraryGenerator::new(); + let mut arb = ArbitraryGenerator::new(); assert!( num_blocks.gt(&0) && max_txs_per_block.gt(&0), "num_blocks and max_tx_per_block must be at least 1" @@ -397,7 +397,7 @@ mod tests { /// Create a known transaction that should be present in some block. fn get_needle() -> (Vec, ProtocolOperation, DepositInfo) { - let arb = ArbitraryGenerator::new(); + let mut arb = ArbitraryGenerator::new(); let network = Network::Regtest; let el_address: [u8; 20] = arb.generate(); @@ -469,20 +469,20 @@ mod tests { /// So, you can assume that this flag represents whether the pair is valid (true) or invalid /// (false). fn generate_mock_tx() -> (Vec, ProtocolOperation, bool) { - let arb = ArbitraryGenerator::new(); + let mut arb = ArbitraryGenerator::new(); let should_be_valid: bool = arb.generate(); if should_be_valid.not() { - let (invalid_tx, invalid_protocol_op) = generate_invalid_tx(&arb); + let (invalid_tx, invalid_protocol_op) = generate_invalid_tx(&mut arb); return (invalid_tx, invalid_protocol_op, should_be_valid); } - let (valid_tx, valid_protocol_op) = generate_valid_tx(&arb); + let (valid_tx, valid_protocol_op) = generate_valid_tx(&mut arb); (valid_tx, valid_protocol_op, should_be_valid) } - fn generate_invalid_tx(arb: &ArbitraryGenerator) -> (Vec, ProtocolOperation) { + fn generate_invalid_tx(arb: &mut ArbitraryGenerator) -> (Vec, ProtocolOperation) { let random_protocol_op: ProtocolOperation = arb.generate(); // true => tx invalid @@ -509,7 +509,7 @@ mod tests { (tx_with_invalid_script_pubkey, random_protocol_op) } - fn generate_valid_tx(arb: &ArbitraryGenerator) -> (Vec, ProtocolOperation) { + fn generate_valid_tx(arb: &mut ArbitraryGenerator) -> (Vec, ProtocolOperation) { let (tx, spend_info, script_to_spend) = generate_mock_unsigned_tx(); let random_hash = *spend_info @@ -565,7 +565,7 @@ mod tests { let empty_deposits = empty_chain_state.deposits_table_mut(); let mut deposits_table = DepositsTable::new_empty(); - let arb = ArbitraryGenerator::new(); + let mut arb = ArbitraryGenerator::new(); let mut operators: Vec = arb.generate(); loop { diff --git a/crates/bridge-relay/src/relayer.rs b/crates/bridge-relay/src/relayer.rs index c7f747840..5c36033f5 100644 --- a/crates/bridge-relay/src/relayer.rs +++ b/crates/bridge-relay/src/relayer.rs @@ -245,7 +245,7 @@ mod tests { let processed_msgs = RecentMessageTracker::new(); // Create valid BridgeMsgId instances for testing - let ag = ArbitraryGenerator::new(); + let mut ag = ArbitraryGenerator::new(); let cur_message_id: BridgeMsgId = ag.generate(); let old_message_id: BridgeMsgId = ag.generate(); assert_ne!(cur_message_id, old_message_id); diff --git a/crates/consensus-logic/src/duty/block_assembly.rs b/crates/consensus-logic/src/duty/block_assembly.rs index 3191229e8..30cc0435b 100644 --- a/crates/consensus-logic/src/duty/block_assembly.rs +++ b/crates/consensus-logic/src/duty/block_assembly.rs @@ -444,7 +444,7 @@ mod tests { #[test] fn test_find_pivot_noop() { - let ag = ArbitraryGenerator::new_with_size(1 << 12); + let mut ag = ArbitraryGenerator::new_with_size(1 << 12); let blkids: [L1BlockId; 10] = ag.generate(); eprintln!("{blkids:#?}"); @@ -467,7 +467,7 @@ mod tests { #[test] fn test_find_pivot_noop_offset() { - let ag = ArbitraryGenerator::new_with_size(1 << 12); + let mut ag = ArbitraryGenerator::new_with_size(1 << 12); let blkids: [L1BlockId; 10] = ag.generate(); eprintln!("{blkids:#?}"); @@ -490,7 +490,7 @@ mod tests { #[test] fn test_find_pivot_simple_extend() { - let ag = ArbitraryGenerator::new_with_size(1 << 12); + let mut ag = ArbitraryGenerator::new_with_size(1 << 12); let blkids1: [L1BlockId; 10] = ag.generate(); let mut blkids2 = Vec::from(blkids1); @@ -522,7 +522,7 @@ mod tests { #[test] fn test_find_pivot_typical_reorg() { - let ag = ArbitraryGenerator::new_with_size(1 << 16); + let mut ag = ArbitraryGenerator::new_with_size(1 << 16); let mut blkids1: Vec = Vec::new(); for _ in 0..10 { @@ -559,7 +559,7 @@ mod tests { #[test] fn test_find_pivot_cur_shorter_reorg() { - let ag = ArbitraryGenerator::new_with_size(1 << 16); + let mut ag = ArbitraryGenerator::new_with_size(1 << 16); let mut blkids1: Vec = Vec::new(); for _ in 0..10 { @@ -600,7 +600,7 @@ mod tests { #[test] fn test_find_pivot_disjoint() { - let ag = ArbitraryGenerator::new_with_size(1 << 16); + let mut ag = ArbitraryGenerator::new_with_size(1 << 16); let mut blkids1: Vec = Vec::new(); for _ in 0..10 { diff --git a/crates/evmexec/src/engine.rs b/crates/evmexec/src/engine.rs index e05f852a8..39aea240d 100644 --- a/crates/evmexec/src/engine.rs +++ b/crates/evmexec/src/engine.rs @@ -518,7 +518,7 @@ mod tests { let el_payload = random_el_payload(); - let arb = strata_test_utils::ArbitraryGenerator::new(); + let mut arb = strata_test_utils::ArbitraryGenerator::new(); let l2block: L2Block = arb.generate(); let accessory = L2BlockAccessory::new(borsh::to_vec(&el_payload).unwrap()); let l2block_bundle = L2BlockBundle::new(l2block, accessory); diff --git a/crates/primitives/src/l1.rs b/crates/primitives/src/l1.rs index ab8a2b7b9..bfefc770f 100644 --- a/crates/primitives/src/l1.rs +++ b/crates/primitives/src/l1.rs @@ -1314,7 +1314,7 @@ mod tests { #[test] fn test_bitcoin_txid_serialize_deserialize() { - let generator = ArbitraryGenerator::new(); + let mut generator = ArbitraryGenerator::new(); let txid: BitcoinTxid = generator.generate(); let serialized_txid = diff --git a/crates/primitives/src/relay/util.rs b/crates/primitives/src/relay/util.rs index 6c0b20909..fb7fa122d 100644 --- a/crates/primitives/src/relay/util.rs +++ b/crates/primitives/src/relay/util.rs @@ -187,7 +187,7 @@ mod tests { #[test] fn test_message_signer_serde() { - let generator = ArbitraryGenerator::new(); + let mut generator = ArbitraryGenerator::new(); let txid: BitcoinTxid = generator.generate(); let scope = Scope::V0PubNonce(txid); let payload: Musig2PubNonce = generator.generate(); diff --git a/crates/rocksdb-store/src/bridge/db.rs b/crates/rocksdb-store/src/bridge/db.rs index 881a19886..635614a48 100644 --- a/crates/rocksdb-store/src/bridge/db.rs +++ b/crates/rocksdb-store/src/bridge/db.rs @@ -233,7 +233,7 @@ mod tests { fn test_bridge_duty_status_db() { let db = setup_duty_db(); - let arb = ArbitraryGenerator::new(); + let mut arb = ArbitraryGenerator::new(); let duty_status: BridgeDutyStatus = arb.generate(); let txid: BitcoinTxid = arb.generate(); @@ -315,7 +315,7 @@ mod tests { fn test_bridge_duty_index_db() { let db = setup_bridge_duty_index_db(); - let arb = ArbitraryGenerator::new(); + let mut arb = ArbitraryGenerator::new(); let checkpoint: u64 = arb.generate(); diff --git a/crates/rocksdb-store/src/bridge_relay/db.rs b/crates/rocksdb-store/src/bridge_relay/db.rs index 011660998..577bf6755 100644 --- a/crates/rocksdb-store/src/bridge_relay/db.rs +++ b/crates/rocksdb-store/src/bridge_relay/db.rs @@ -118,7 +118,7 @@ mod tests { } fn make_bridge_msg() -> (u128, BridgeMessage) { - let arb = ArbitraryGenerator::new(); + let mut arb = ArbitraryGenerator::new(); let msg: BridgeMessage = arb.generate(); diff --git a/crates/rocksdb-store/src/l1/db.rs b/crates/rocksdb-store/src/l1/db.rs index 415ee9aab..d9928a8e7 100644 --- a/crates/rocksdb-store/src/l1/db.rs +++ b/crates/rocksdb-store/src/l1/db.rs @@ -239,7 +239,7 @@ mod tests { db: &L1Db, num_txs: usize, ) -> (L1BlockManifest, Vec, CompactMmr) { - let arb = ArbitraryGenerator::new(); + let mut arb = ArbitraryGenerator::new(); // TODO maybe tweak this to make it a bit more realistic? let mf: L1BlockManifest = arb.generate(); diff --git a/crates/rocksdb-store/src/l2/db.rs b/crates/rocksdb-store/src/l2/db.rs index 283081d64..0a70c7579 100644 --- a/crates/rocksdb-store/src/l2/db.rs +++ b/crates/rocksdb-store/src/l2/db.rs @@ -117,7 +117,7 @@ mod tests { use crate::test_utils::get_rocksdb_tmp_instance; fn get_mock_data() -> L2BlockBundle { - let arb = ArbitraryGenerator::new(); + let mut arb = ArbitraryGenerator::new(); let l2_block: L2BlockBundle = arb.generate(); l2_block diff --git a/crates/test-utils/src/l2.rs b/crates/test-utils/src/l2.rs index 0b84d92d3..fbda160f2 100644 --- a/crates/test-utils/src/l2.rs +++ b/crates/test-utils/src/l2.rs @@ -20,7 +20,7 @@ use strata_state::{ use crate::{bitcoin::get_btc_chain, ArbitraryGenerator}; pub fn gen_block(parent: Option<&SignedL2BlockHeader>) -> L2BlockBundle { - let arb = ArbitraryGenerator::new(); + let mut arb = ArbitraryGenerator::new(); let header: L2BlockHeader = arb.generate(); let body: L2BlockBody = arb.generate(); let accessory: L2BlockAccessory = arb.generate(); diff --git a/crates/test-utils/src/lib.rs b/crates/test-utils/src/lib.rs index 9823504a7..9494337d0 100644 --- a/crates/test-utils/src/lib.rs +++ b/crates/test-utils/src/lib.rs @@ -1,46 +1,40 @@ -use std::sync::atomic::{AtomicUsize, Ordering}; - use arbitrary::{Arbitrary, Unstructured}; -use rand::{rngs::OsRng, RngCore}; +use rand::{thread_rng, RngCore}; pub mod bitcoin; pub mod bridge; pub mod evm_ee; pub mod l2; -const ARB_GEN_LEN: usize = 1 << 24; // 16 MiB +// Smaller buffer size as compared to 2^24 +const ARB_GEN_LEN: usize = 128; pub struct ArbitraryGenerator { - buf: Vec, - off: AtomicUsize, + rng: rand::rngs::ThreadRng, // Thread-local RNG + buf: Vec, // Persistent buffer } - impl Default for ArbitraryGenerator { fn default() -> Self { Self::new() } } - impl ArbitraryGenerator { pub fn new() -> Self { - Self::new_with_size(ARB_GEN_LEN) + ArbitraryGenerator { + rng: thread_rng(), + buf: vec![0u8; ARB_GEN_LEN], + } } - - pub fn new_with_size(n: usize) -> Self { - let mut buf = vec![0; n]; - OsRng.fill_bytes(&mut buf); // 128 wasn't enough - let off = AtomicUsize::new(0); - ArbitraryGenerator { buf, off } + pub fn new_with_size(s: usize) -> Self { + ArbitraryGenerator { + rng: thread_rng(), + buf: vec![0u8; s], + } } - pub fn generate<'a, T: Arbitrary<'a> + Clone>(&'a self) -> T { - // Doing hacky atomics to make this actually be reusable, this is pretty bad. - let off = self.off.load(Ordering::Relaxed); - let mut u = Unstructured::new(&self.buf[off..]); - let prev_off = u.len(); - let inst = T::arbitrary(&mut u).expect("failed to generate arbitrary instance"); - let additional_off = prev_off - u.len(); - self.off.store(off + additional_off, Ordering::Relaxed); - inst + pub fn generate<'a, T: Arbitrary<'a> + Clone>(&'a mut self) -> T { + self.rng.fill_bytes(&mut self.buf); + let mut u = Unstructured::new(&self.buf); + T::arbitrary(&mut u).expect("Failed to generate arbitrary instance") } }