From c535d59c8dada6029c9036560eadb5274e8fcd2e Mon Sep 17 00:00:00 2001 From: Qiwei Yang Date: Sun, 28 Apr 2024 21:15:35 +0800 Subject: [PATCH] refactor: extract GenesisAllocator type from primitives (#7932) Co-authored-by: Matthias Seitz --- Cargo.lock | 10 + Cargo.toml | 2 + crates/consensus/beacon/Cargo.toml | 1 + crates/consensus/beacon/src/engine/mod.rs | 6 +- crates/primitives/src/genesis.rs | 209 ------------------ testing/testing-utils/Cargo.toml | 17 ++ .../testing-utils/src/genesis_allocator.rs | 204 +++++++++++++++++ testing/testing-utils/src/lib.rs | 12 + 8 files changed, 248 insertions(+), 213 deletions(-) create mode 100644 testing/testing-utils/Cargo.toml create mode 100644 testing/testing-utils/src/genesis_allocator.rs create mode 100644 testing/testing-utils/src/lib.rs diff --git a/Cargo.lock b/Cargo.lock index 50fa59f01..54db17110 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6395,6 +6395,7 @@ dependencies = [ "reth-stages-api", "reth-static-file", "reth-tasks", + "reth-testing-utils", "reth-tokio-util", "reth-tracing", "schnellru", @@ -7780,6 +7781,15 @@ dependencies = [ "tracing-futures", ] +[[package]] +name = "reth-testing-utils" +version = "0.2.0-beta.6" +dependencies = [ + "alloy-genesis", + "reth-primitives", + "secp256k1", +] + [[package]] name = "reth-tokio-util" version = "0.2.0-beta.6" diff --git a/Cargo.toml b/Cargo.toml index 7edc96312..c970bf6a5 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -91,6 +91,7 @@ members = [ "examples/exex/op-bridge/", "examples/db-access", "testing/ef-tests/", + "testing/testing-utils", ] default-members = ["bin/reth"] @@ -275,6 +276,7 @@ reth-trie = { path = "crates/trie" } reth-trie-parallel = { path = "crates/trie-parallel" } reth-optimism-consensus = { path = "crates/optimism/consensus" } reth-node-events = { path = "crates/node/events" } +reth-testing-utils = { path = "testing/testing-utils" } # revm revm = { version = "8.0.0", features = ["std", "secp256k1"], default-features = false } diff --git a/crates/consensus/beacon/Cargo.toml b/crates/consensus/beacon/Cargo.toml index 38dd772af..4e35d06f0 100644 --- a/crates/consensus/beacon/Cargo.toml +++ b/crates/consensus/beacon/Cargo.toml @@ -59,6 +59,7 @@ reth-downloaders.workspace = true reth-evm-ethereum.workspace = true reth-ethereum-engine-primitives.workspace = true reth-config.workspace = true +reth-testing-utils.workspace = true assert_matches.workspace = true diff --git a/crates/consensus/beacon/src/engine/mod.rs b/crates/consensus/beacon/src/engine/mod.rs index 72fc97297..428b95c0b 100644 --- a/crates/consensus/beacon/src/engine/mod.rs +++ b/crates/consensus/beacon/src/engine/mod.rs @@ -2350,11 +2350,9 @@ mod tests { use super::*; use reth_db::test_utils::create_test_static_files_dir; use reth_interfaces::test_utils::generators::random_block; - use reth_primitives::{ - genesis::{Genesis, GenesisAllocator}, - Hardfork, U256, - }; + use reth_primitives::{genesis::Genesis, Hardfork, U256}; use reth_provider::test_utils::blocks::BlockchainTestData; + use reth_testing_utils::GenesisAllocator; #[tokio::test] async fn new_payload_before_forkchoice() { diff --git a/crates/primitives/src/genesis.rs b/crates/primitives/src/genesis.rs index 991b01bd7..c81d9488d 100644 --- a/crates/primitives/src/genesis.rs +++ b/crates/primitives/src/genesis.rs @@ -3,212 +3,3 @@ // re-export genesis types #[doc(inline)] pub use alloy_genesis::*; - -#[cfg(any(test, feature = "test-utils"))] -pub use allocator::GenesisAllocator; - -#[cfg(any(test, feature = "test-utils"))] -mod allocator { - use crate::{public_key_to_address, Address, Bytes, B256, U256}; - use alloy_genesis::GenesisAccount; - use secp256k1::{ - rand::{thread_rng, RngCore}, - Keypair, Secp256k1, - }; - use std::collections::{hash_map::Entry, BTreeMap, HashMap}; - - /// This helps create a custom genesis alloc by making it easy to add funded accounts with known - /// signers to the genesis block. - /// - /// # Example - /// ``` - /// # use reth_primitives::{ genesis::GenesisAllocator, Address, U256, hex, Bytes}; - /// # use std::str::FromStr; - /// let mut allocator = GenesisAllocator::default(); - /// - /// // This will add a genesis account to the alloc builder, with the provided balance. The - /// // signer for the account will be returned. - /// let (_signer, _addr) = allocator.new_funded_account(U256::from(100_000_000_000_000_000u128)); - /// - /// // You can also provide code for the account. - /// let code = Bytes::from_str("0x1234").unwrap(); - /// let (_second_signer, _second_addr) = - /// allocator.new_funded_account_with_code(U256::from(100_000_000_000_000_000u128), code); - /// - /// // You can also add an account with a specific address. - /// // This will not return a signer, since the address is provided by the user and the signer - /// // may be unknown. - /// let addr = "0Ac1dF02185025F65202660F8167210A80dD5086".parse::
().unwrap(); - /// allocator.add_funded_account_with_address(addr, U256::from(100_000_000_000_000_000u128)); - /// - /// // Once you're done adding accounts, you can build the alloc. - /// let alloc = allocator.build(); - /// ``` - #[derive(Debug)] - pub struct GenesisAllocator<'a> { - /// The genesis alloc to be built. - alloc: HashMap, - /// The rng to use for generating key pairs. - rng: Box, - } - - impl<'a> GenesisAllocator<'a> { - /// Initialize a new alloc builder with the provided rng. - pub fn new_with_rng(rng: &'a mut R) -> Self - where - R: RngCore + std::fmt::Debug, - { - Self { alloc: HashMap::default(), rng: Box::new(rng) } - } - - /// Use the provided rng for generating key pairs. - pub fn with_rng(mut self, rng: &'a mut R) -> Self - where - R: RngCore + std::fmt::Debug, - { - self.rng = Box::new(rng); - self - } - - /// Add a funded account to the genesis alloc. - /// - /// Returns the key pair for the account and the account's address. - pub fn new_funded_account(&mut self, balance: U256) -> (Keypair, Address) { - let secp = Secp256k1::new(); - let pair = Keypair::new(&secp, &mut self.rng); - let address = public_key_to_address(pair.public_key()); - - self.alloc.insert(address, GenesisAccount::default().with_balance(balance)); - - (pair, address) - } - - /// Add a funded account to the genesis alloc with the provided code. - /// - /// Returns the key pair for the account and the account's address. - pub fn new_funded_account_with_code( - &mut self, - balance: U256, - code: Bytes, - ) -> (Keypair, Address) { - let secp = Secp256k1::new(); - let pair = Keypair::new(&secp, &mut self.rng); - let address = public_key_to_address(pair.public_key()); - - self.alloc.insert( - address, - GenesisAccount::default().with_balance(balance).with_code(Some(code)), - ); - - (pair, address) - } - - /// Adds a funded account to the genesis alloc with the provided storage. - /// - /// Returns the key pair for the account and the account's address. - pub fn new_funded_account_with_storage( - &mut self, - balance: U256, - storage: BTreeMap, - ) -> (Keypair, Address) { - let secp = Secp256k1::new(); - let pair = Keypair::new(&secp, &mut self.rng); - let address = public_key_to_address(pair.public_key()); - - self.alloc.insert( - address, - GenesisAccount::default().with_balance(balance).with_storage(Some(storage)), - ); - - (pair, address) - } - - /// Adds an account with code and storage to the genesis alloc. - /// - /// Returns the key pair for the account and the account's address. - pub fn new_account_with_code_and_storage( - &mut self, - code: Bytes, - storage: BTreeMap, - ) -> (Keypair, Address) { - let secp = Secp256k1::new(); - let pair = Keypair::new(&secp, &mut self.rng); - let address = public_key_to_address(pair.public_key()); - - self.alloc.insert( - address, - GenesisAccount::default().with_code(Some(code)).with_storage(Some(storage)), - ); - - (pair, address) - } - - /// Adds an account with code to the genesis alloc. - /// - /// Returns the key pair for the account and the account's address. - pub fn new_account_with_code(&mut self, code: Bytes) -> (Keypair, Address) { - let secp = Secp256k1::new(); - let pair = Keypair::new(&secp, &mut self.rng); - let address = public_key_to_address(pair.public_key()); - - self.alloc.insert(address, GenesisAccount::default().with_code(Some(code))); - - (pair, address) - } - - /// Add a funded account to the genesis alloc with the provided address. - /// - /// Neither the key pair nor the account will be returned, since the address is provided by - /// the user and the signer may be unknown. - pub fn add_funded_account_with_address(&mut self, address: Address, balance: U256) { - self.alloc.insert(address, GenesisAccount::default().with_balance(balance)); - } - - /// Adds the given [GenesisAccount] to the genesis alloc. - /// - /// Returns the key pair for the account and the account's address. - pub fn add_account(&mut self, account: GenesisAccount) -> Address { - let secp = Secp256k1::new(); - let pair = Keypair::new(&secp, &mut self.rng); - let address = public_key_to_address(pair.public_key()); - - self.alloc.insert(address, account); - - address - } - - /// Gets the account for the provided address. - /// - /// If it does not exist, this returns `None`. - pub fn get_account(&self, address: &Address) -> Option<&GenesisAccount> { - self.alloc.get(address) - } - - /// Gets a mutable version of the account for the provided address, if it exists. - pub fn get_account_mut(&mut self, address: &Address) -> Option<&mut GenesisAccount> { - self.alloc.get_mut(address) - } - - /// Gets an [Entry] for the provided address. - pub fn account_entry(&mut self, address: Address) -> Entry<'_, Address, GenesisAccount> { - self.alloc.entry(address) - } - - /// Build the genesis alloc. - pub fn build(self) -> HashMap { - self.alloc - } - } - - impl Default for GenesisAllocator<'_> { - fn default() -> Self { - Self { alloc: HashMap::default(), rng: Box::new(thread_rng()) } - } - } - - /// Helper trait that encapsulates [RngCore], and [Debug](std::fmt::Debug) to get around rules - /// for auto traits (Opt-in built-in traits). - trait RngDebug: RngCore + std::fmt::Debug {} - - impl RngDebug for T where T: RngCore + std::fmt::Debug {} -} diff --git a/testing/testing-utils/Cargo.toml b/testing/testing-utils/Cargo.toml new file mode 100644 index 000000000..97a4c78df --- /dev/null +++ b/testing/testing-utils/Cargo.toml @@ -0,0 +1,17 @@ +[package] +name = "reth-testing-utils" +description = "Testing utils for reth." +version.workspace = true +edition.workspace = true +rust-version.workspace = true +license.workspace = true +homepage.workspace = true +repository.workspace = true + +[lints] +workspace = true + +[dependencies] +secp256k1.workspace = true +alloy-genesis.workspace = true +reth-primitives.workspace = true diff --git a/testing/testing-utils/src/genesis_allocator.rs b/testing/testing-utils/src/genesis_allocator.rs new file mode 100644 index 000000000..33b7188d1 --- /dev/null +++ b/testing/testing-utils/src/genesis_allocator.rs @@ -0,0 +1,204 @@ +//! Helps create a custom genesis alloc by making it easy to add funded accounts with known +//! signers to the genesis block. + +use alloy_genesis::GenesisAccount; +use reth_primitives::{public_key_to_address, Address, Bytes, B256, U256}; +use secp256k1::{ + rand::{thread_rng, RngCore}, + Keypair, Secp256k1, +}; +use std::collections::{hash_map::Entry, BTreeMap, HashMap}; + +/// This helps create a custom genesis alloc by making it easy to add funded accounts with known +/// signers to the genesis block. +/// +/// # Example +/// ``` +/// # use reth_primitives::{Address, U256, hex, Bytes}; +/// # use reth_testing_utils::GenesisAllocator; +/// # use std::str::FromStr; +/// let mut allocator = GenesisAllocator::default(); +/// +/// // This will add a genesis account to the alloc builder, with the provided balance. The +/// // signer for the account will be returned. +/// let (_signer, _addr) = allocator.new_funded_account(U256::from(100_000_000_000_000_000u128)); +/// +/// // You can also provide code for the account. +/// let code = Bytes::from_str("0x1234").unwrap(); +/// let (_second_signer, _second_addr) = +/// allocator.new_funded_account_with_code(U256::from(100_000_000_000_000_000u128), code); +/// +/// // You can also add an account with a specific address. +/// // This will not return a signer, since the address is provided by the user and the signer +/// // may be unknown. +/// let addr = "0Ac1dF02185025F65202660F8167210A80dD5086".parse::
().unwrap(); +/// allocator.add_funded_account_with_address(addr, U256::from(100_000_000_000_000_000u128)); +/// +/// // Once you're done adding accounts, you can build the alloc. +/// let alloc = allocator.build(); +/// ``` +#[derive(Debug)] +pub struct GenesisAllocator<'a> { + /// The genesis alloc to be built. + alloc: HashMap, + /// The rng to use for generating key pairs. + rng: Box, +} + +impl<'a> GenesisAllocator<'a> { + /// Initialize a new alloc builder with the provided rng. + pub fn new_with_rng(rng: &'a mut R) -> Self + where + R: RngCore + std::fmt::Debug, + { + Self { alloc: HashMap::default(), rng: Box::new(rng) } + } + + /// Use the provided rng for generating key pairs. + pub fn with_rng(mut self, rng: &'a mut R) -> Self + where + R: RngCore + std::fmt::Debug, + { + self.rng = Box::new(rng); + self + } + + /// Add a funded account to the genesis alloc. + /// + /// Returns the key pair for the account and the account's address. + pub fn new_funded_account(&mut self, balance: U256) -> (Keypair, Address) { + let secp = Secp256k1::new(); + let pair = Keypair::new(&secp, &mut self.rng); + let address = public_key_to_address(pair.public_key()); + + self.alloc.insert(address, GenesisAccount::default().with_balance(balance)); + + (pair, address) + } + + /// Add a funded account to the genesis alloc with the provided code. + /// + /// Returns the key pair for the account and the account's address. + pub fn new_funded_account_with_code( + &mut self, + balance: U256, + code: Bytes, + ) -> (Keypair, Address) { + let secp = Secp256k1::new(); + let pair = Keypair::new(&secp, &mut self.rng); + let address = public_key_to_address(pair.public_key()); + + self.alloc + .insert(address, GenesisAccount::default().with_balance(balance).with_code(Some(code))); + + (pair, address) + } + + /// Adds a funded account to the genesis alloc with the provided storage. + /// + /// Returns the key pair for the account and the account's address. + pub fn new_funded_account_with_storage( + &mut self, + balance: U256, + storage: BTreeMap, + ) -> (Keypair, Address) { + let secp = Secp256k1::new(); + let pair = Keypair::new(&secp, &mut self.rng); + let address = public_key_to_address(pair.public_key()); + + self.alloc.insert( + address, + GenesisAccount::default().with_balance(balance).with_storage(Some(storage)), + ); + + (pair, address) + } + + /// Adds an account with code and storage to the genesis alloc. + /// + /// Returns the key pair for the account and the account's address. + pub fn new_account_with_code_and_storage( + &mut self, + code: Bytes, + storage: BTreeMap, + ) -> (Keypair, Address) { + let secp = Secp256k1::new(); + let pair = Keypair::new(&secp, &mut self.rng); + let address = public_key_to_address(pair.public_key()); + + self.alloc.insert( + address, + GenesisAccount::default().with_code(Some(code)).with_storage(Some(storage)), + ); + + (pair, address) + } + + /// Adds an account with code to the genesis alloc. + /// + /// Returns the key pair for the account and the account's address. + pub fn new_account_with_code(&mut self, code: Bytes) -> (Keypair, Address) { + let secp = Secp256k1::new(); + let pair = Keypair::new(&secp, &mut self.rng); + let address = public_key_to_address(pair.public_key()); + + self.alloc.insert(address, GenesisAccount::default().with_code(Some(code))); + + (pair, address) + } + + /// Add a funded account to the genesis alloc with the provided address. + /// + /// Neither the key pair nor the account will be returned, since the address is provided by + /// the user and the signer may be unknown. + pub fn add_funded_account_with_address(&mut self, address: Address, balance: U256) { + self.alloc.insert(address, GenesisAccount::default().with_balance(balance)); + } + + /// Adds the given [GenesisAccount] to the genesis alloc. + /// + /// Returns the key pair for the account and the account's address. + pub fn add_account(&mut self, account: GenesisAccount) -> Address { + let secp = Secp256k1::new(); + let pair = Keypair::new(&secp, &mut self.rng); + let address = public_key_to_address(pair.public_key()); + + self.alloc.insert(address, account); + + address + } + + /// Gets the account for the provided address. + /// + /// If it does not exist, this returns `None`. + pub fn get_account(&self, address: &Address) -> Option<&GenesisAccount> { + self.alloc.get(address) + } + + /// Gets a mutable version of the account for the provided address, if it exists. + pub fn get_account_mut(&mut self, address: &Address) -> Option<&mut GenesisAccount> { + self.alloc.get_mut(address) + } + + /// Gets an [Entry] for the provided address. + pub fn account_entry(&mut self, address: Address) -> Entry<'_, Address, GenesisAccount> { + self.alloc.entry(address) + } + + /// Build the genesis alloc. + pub fn build(self) -> HashMap { + self.alloc + } +} + +impl Default for GenesisAllocator<'_> { + fn default() -> Self { + Self { alloc: HashMap::default(), rng: Box::new(thread_rng()) } + } +} + +/// Helper trait that encapsulates [RngCore], and [Debug](std::fmt::Debug) to get around rules +/// for auto traits (Opt-in built-in traits). +trait RngDebug: RngCore + std::fmt::Debug {} + +impl RngDebug for T where T: RngCore + std::fmt::Debug {} diff --git a/testing/testing-utils/src/lib.rs b/testing/testing-utils/src/lib.rs new file mode 100644 index 000000000..0cf98c6ff --- /dev/null +++ b/testing/testing-utils/src/lib.rs @@ -0,0 +1,12 @@ +//! Testing utilities. + +#![doc( + html_logo_url = "https://raw.githubusercontent.com/paradigmxyz/reth/main/assets/reth-docs.png", + html_favicon_url = "https://avatars0.githubusercontent.com/u/97369466?s=256", + issue_tracker_base_url = "https://github.com/paradigmxyz/reth/issues/" +)] +#![cfg_attr(docsrs, feature(doc_cfg, doc_auto_cfg))] + +pub mod genesis_allocator; + +pub use genesis_allocator::GenesisAllocator;