refactor: extract GenesisAllocator type from primitives (#7932)

Co-authored-by: Matthias Seitz <matthias.seitz@outlook.de>

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

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::<Address>().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<Address, GenesisAccount>,
+    /// The rng to use for generating key pairs.
+    rng: Box<dyn RngDebug + 'a>,
+}
+
+impl<'a> GenesisAllocator<'a> {
+    /// Initialize a new alloc builder with the provided rng.
+    pub fn new_with_rng<R>(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<R>(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<B256, B256>,
+    ) -> (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<B256, B256>,
+    ) -> (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<Address, GenesisAccount> {
+        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<T> RngDebug for T where T: RngCore + std::fmt::Debug {}

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;