feat: add GenesisAllocator for easier genesis alloc creation (#6021)

bin/reth/src/builder/mod.rs

@@ -1452,6 +1452,7 @@ mod tests {
         // spawn_test_node takes roughly 1 second per node, so this test takes ~4 seconds
         let num_nodes = 4;
 
+        // this reserves instances 3-6
         let starting_instance = 3;
         let mut handles = Vec::new();
         for i in 0..num_nodes {

crates/consensus/beacon/src/engine/mod.rs

@@ -2345,11 +2345,8 @@ mod tests {
 
     mod new_payload {
         use super::*;
-        use reth_interfaces::test_utils::{
-            generators,
-            generators::{generate_keys, random_block},
-        };
-        use reth_primitives::{public_key_to_address, Genesis, GenesisAccount, Hardfork, U256};
+        use reth_interfaces::test_utils::{generators, generators::random_block};
+        use reth_primitives::{Genesis, GenesisAllocator, Hardfork, U256};
         use reth_provider::test_utils::blocks::BlockChainTestData;
 
         #[tokio::test]
@@ -2453,14 +2450,14 @@ mod tests {
         #[tokio::test]
         async fn simple_validate_block() {
             let mut rng = generators::rng();
-            let genesis_keys = generate_keys(&mut rng, 16);
-            let amount = 1000000000000000000u64;
-            let alloc = genesis_keys.iter().map(|pair| {
-                (
-                    public_key_to_address(pair.public_key()),
-                    GenesisAccount::default().with_balance(U256::from(amount)),
-                )
-            });
+            let amount = U256::from(1000000000000000000u64);
+            let mut allocator = GenesisAllocator::default().with_rng(&mut rng);
+            for _ in 0..16 {
+                // add 16 new accounts
+                allocator.new_funded_account(amount);
+            }
+
+            let alloc = allocator.build();
 
             let genesis = Genesis::default().extend_accounts(alloc);
 

crates/primitives/src/genesis.rs

@@ -1,5 +1,5 @@
 use crate::{
-    keccak256,
+    keccak256, public_key_to_address,
     serde_helper::{
         json_u256::{deserialize_json_ttd_opt, deserialize_json_u256},
         num::{u64_hex_or_decimal, u64_hex_or_decimal_opt},
@@ -7,8 +7,12 @@ use crate::{
     },
     Account, Address, Bytes, B256, U256,
 };
+use secp256k1::{
+    rand::{thread_rng, RngCore},
+    KeyPair, Secp256k1,
+};
 use serde::{Deserialize, Serialize};
-use std::collections::HashMap;
+use std::collections::{hash_map::Entry, HashMap};
 
 /// The genesis block specification.
 #[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq, Eq)]
@@ -185,6 +189,198 @@ impl From<GenesisAccount> for Account {
     }
 }
 
+/// 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 {}
+
+/// 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::{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::<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: HashMap<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: HashMap<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()) }
+    }
+}
+
 /// Defines core blockchain settings per block.
 ///
 /// Tailors unique settings for each network based on its genesis block.

crates/primitives/src/lib.rs

@@ -59,7 +59,7 @@ pub use constants::{
     KECCAK_EMPTY, MAINNET_GENESIS_HASH, SEPOLIA_GENESIS_HASH,
 };
 pub use error::{GotExpected, GotExpectedBoxed};
-pub use genesis::{ChainConfig, Genesis, GenesisAccount};
+pub use genesis::{ChainConfig, Genesis, GenesisAccount, GenesisAllocator};
 pub use header::{Header, HeadersDirection, SealedHeader};
 pub use integer_list::IntegerList;
 pub use log::{logs_bloom, Log};