Allow for genesis to be exported by the tool

crates/config/src/lib.rs

@@ -34,6 +34,9 @@ pub enum Context {
 
     /// Exports the JSON schema of the MatterLabs test format used by the tool.
     ExportJsonSchema,
+
+    /// Exports the genesis file of the desired platform.
+    ExportGenesis(Box<ExportGenesisContext>),
 }
 
 impl Context {
@@ -51,7 +54,7 @@ impl AsRef<WorkingDirectoryConfiguration> for Context {
         match self {
             Self::Test(context) => context.as_ref().as_ref(),
             Self::Benchmark(context) => context.as_ref().as_ref(),
-            Self::ExportJsonSchema => unreachable!(),
+            Self::ExportJsonSchema | Self::ExportGenesis(..) => unreachable!(),
         }
     }
 }
@@ -61,7 +64,7 @@ impl AsRef<CorpusConfiguration> for Context {
         match self {
             Self::Test(context) => context.as_ref().as_ref(),
             Self::Benchmark(context) => context.as_ref().as_ref(),
-            Self::ExportJsonSchema => unreachable!(),
+            Self::ExportJsonSchema | Self::ExportGenesis(..) => unreachable!(),
         }
     }
 }
@@ -71,7 +74,7 @@ impl AsRef<SolcConfiguration> for Context {
         match self {
             Self::Test(context) => context.as_ref().as_ref(),
             Self::Benchmark(context) => context.as_ref().as_ref(),
-            Self::ExportJsonSchema => unreachable!(),
+            Self::ExportJsonSchema | Self::ExportGenesis(..) => unreachable!(),
         }
     }
 }
@@ -81,7 +84,7 @@ impl AsRef<ResolcConfiguration> for Context {
         match self {
             Self::Test(context) => context.as_ref().as_ref(),
             Self::Benchmark(context) => context.as_ref().as_ref(),
-            Self::ExportJsonSchema => unreachable!(),
+            Self::ExportJsonSchema | Self::ExportGenesis(..) => unreachable!(),
         }
     }
 }
@@ -91,6 +94,7 @@ impl AsRef<GethConfiguration> for Context {
         match self {
             Self::Test(context) => context.as_ref().as_ref(),
             Self::Benchmark(context) => context.as_ref().as_ref(),
+            Self::ExportGenesis(context) => context.as_ref().as_ref(),
             Self::ExportJsonSchema => unreachable!(),
         }
     }
@@ -101,6 +105,7 @@ impl AsRef<KurtosisConfiguration> for Context {
         match self {
             Self::Test(context) => context.as_ref().as_ref(),
             Self::Benchmark(context) => context.as_ref().as_ref(),
+            Self::ExportGenesis(context) => context.as_ref().as_ref(),
             Self::ExportJsonSchema => unreachable!(),
         }
     }
@@ -111,6 +116,7 @@ impl AsRef<PolkadotParachainConfiguration> for Context {
         match self {
             Self::Test(context) => context.as_ref().as_ref(),
             Self::Benchmark(context) => context.as_ref().as_ref(),
+            Self::ExportGenesis(context) => context.as_ref().as_ref(),
             Self::ExportJsonSchema => unreachable!(),
         }
     }
@@ -121,6 +127,7 @@ impl AsRef<KitchensinkConfiguration> for Context {
         match self {
             Self::Test(context) => context.as_ref().as_ref(),
             Self::Benchmark(context) => context.as_ref().as_ref(),
+            Self::ExportGenesis(context) => context.as_ref().as_ref(),
             Self::ExportJsonSchema => unreachable!(),
         }
     }
@@ -131,6 +138,7 @@ impl AsRef<ReviveDevNodeConfiguration> for Context {
         match self {
             Self::Test(context) => context.as_ref().as_ref(),
             Self::Benchmark(context) => context.as_ref().as_ref(),
+            Self::ExportGenesis(context) => context.as_ref().as_ref(),
             Self::ExportJsonSchema => unreachable!(),
         }
     }
@@ -141,7 +149,7 @@ impl AsRef<EthRpcConfiguration> for Context {
         match self {
             Self::Test(context) => context.as_ref().as_ref(),
             Self::Benchmark(context) => context.as_ref().as_ref(),
-            Self::ExportJsonSchema => unreachable!(),
+            Self::ExportJsonSchema | Self::ExportGenesis(..) => unreachable!(),
         }
     }
 }
@@ -150,7 +158,7 @@ impl AsRef<GenesisConfiguration> for Context {
     fn as_ref(&self) -> &GenesisConfiguration {
         match self {
             Self::Test(context) => context.as_ref().as_ref(),
-            Self::Benchmark(..) => {
+            Self::Benchmark(..) | Self::ExportGenesis(..) => {
                 static GENESIS: LazyLock<GenesisConfiguration> = LazyLock::new(Default::default);
                 &GENESIS
             }
@@ -164,6 +172,7 @@ impl AsRef<WalletConfiguration> for Context {
         match self {
             Self::Test(context) => context.as_ref().as_ref(),
             Self::Benchmark(context) => context.as_ref().as_ref(),
+            Self::ExportGenesis(context) => context.as_ref().as_ref(),
             Self::ExportJsonSchema => unreachable!(),
         }
     }
@@ -174,7 +183,7 @@ impl AsRef<ConcurrencyConfiguration> for Context {
         match self {
             Self::Test(context) => context.as_ref().as_ref(),
             Self::Benchmark(context) => context.as_ref().as_ref(),
-            Self::ExportJsonSchema => unreachable!(),
+            Self::ExportJsonSchema | Self::ExportGenesis(..) => unreachable!(),
         }
     }
 }
@@ -184,7 +193,7 @@ impl AsRef<CompilationConfiguration> for Context {
         match self {
             Self::Test(context) => context.as_ref().as_ref(),
             Self::Benchmark(context) => context.as_ref().as_ref(),
-            Self::ExportJsonSchema => unreachable!(),
+            Self::ExportJsonSchema | Self::ExportGenesis(..) => unreachable!(),
         }
     }
 }
@@ -194,7 +203,7 @@ impl AsRef<ReportConfiguration> for Context {
         match self {
             Self::Test(context) => context.as_ref().as_ref(),
             Self::Benchmark(context) => context.as_ref().as_ref(),
-            Self::ExportJsonSchema => unreachable!(),
+            Self::ExportJsonSchema | Self::ExportGenesis(..) => unreachable!(),
         }
     }
 }
@@ -204,7 +213,7 @@ impl AsRef<IgnoreSuccessConfiguration> for Context {
         match self {
             Self::Test(context) => context.as_ref().as_ref(),
             Self::Benchmark(..) => unreachable!(),
-            Self::ExportJsonSchema => unreachable!(),
+            Self::ExportJsonSchema | Self::ExportGenesis(..) => unreachable!(),
         }
     }
 }
@@ -378,6 +387,36 @@ pub struct BenchmarkingContext {
     pub report_configuration: ReportConfiguration,
 }
 
+#[derive(Clone, Debug, Parser, Serialize, Deserialize)]
+pub struct ExportGenesisContext {
+    /// The platform of choice to export the genesis for.
+    pub platform: PlatformIdentifier,
+
+    /// Configuration parameters for the geth node.
+    #[clap(flatten, next_help_heading = "Geth Configuration")]
+    pub geth_configuration: GethConfiguration,
+
+    /// Configuration parameters for the lighthouse node.
+    #[clap(flatten, next_help_heading = "Lighthouse Configuration")]
+    pub lighthouse_configuration: KurtosisConfiguration,
+
+    /// Configuration parameters for the Kitchensink.
+    #[clap(flatten, next_help_heading = "Kitchensink Configuration")]
+    pub kitchensink_configuration: KitchensinkConfiguration,
+
+    /// Configuration parameters for the Polkadot Parachain.
+    #[clap(flatten, next_help_heading = "Polkadot Parachain Configuration")]
+    pub polkadot_parachain_configuration: PolkadotParachainConfiguration,
+
+    /// Configuration parameters for the Revive Dev Node.
+    #[clap(flatten, next_help_heading = "Revive Dev Node Configuration")]
+    pub revive_dev_node_configuration: ReviveDevNodeConfiguration,
+
+    /// Configuration parameters for the wallet.
+    #[clap(flatten, next_help_heading = "Wallet Configuration")]
+    pub wallet_configuration: WalletConfiguration,
+}
+
 impl Default for TestExecutionContext {
     fn default() -> Self {
         Self::parse_from(["execution-context"])
@@ -482,7 +521,7 @@ impl AsRef<IgnoreSuccessConfiguration> for TestExecutionContext {
 
 impl Default for BenchmarkingContext {
     fn default() -> Self {
-        Self::parse_from(["execution-context"])
+        Self::parse_from(["benchmarking-context"])
     }
 }
 
@@ -570,6 +609,48 @@ impl AsRef<ReportConfiguration> for BenchmarkingContext {
     }
 }
 
+impl Default for ExportGenesisContext {
+    fn default() -> Self {
+        Self::parse_from(["export-genesis-context"])
+    }
+}
+
+impl AsRef<GethConfiguration> for ExportGenesisContext {
+    fn as_ref(&self) -> &GethConfiguration {
+        &self.geth_configuration
+    }
+}
+
+impl AsRef<KurtosisConfiguration> for ExportGenesisContext {
+    fn as_ref(&self) -> &KurtosisConfiguration {
+        &self.lighthouse_configuration
+    }
+}
+
+impl AsRef<KitchensinkConfiguration> for ExportGenesisContext {
+    fn as_ref(&self) -> &KitchensinkConfiguration {
+        &self.kitchensink_configuration
+    }
+}
+
+impl AsRef<PolkadotParachainConfiguration> for ExportGenesisContext {
+    fn as_ref(&self) -> &PolkadotParachainConfiguration {
+        &self.polkadot_parachain_configuration
+    }
+}
+
+impl AsRef<ReviveDevNodeConfiguration> for ExportGenesisContext {
+    fn as_ref(&self) -> &ReviveDevNodeConfiguration {
+        &self.revive_dev_node_configuration
+    }
+}
+
+impl AsRef<WalletConfiguration> for ExportGenesisContext {
+    fn as_ref(&self) -> &WalletConfiguration {
+        &self.wallet_configuration
+    }
+}
+
 /// A set of configuration parameters for the corpus files to use for the execution.
 #[derive(Clone, Debug, Parser, Serialize, Deserialize)]
 pub struct CorpusConfiguration {

crates/core/src/lib.rs

@@ -59,6 +59,9 @@ pub trait Platform {
         context: Context,
         version: Option<VersionOrRequirement>,
     ) -> Pin<Box<dyn Future<Output = anyhow::Result<Box<dyn SolidityCompiler>>>>>;
+
+    /// Exports the genesis/chainspec for the node.
+    fn export_genesis(&self, context: Context) -> anyhow::Result<serde_json::Value>;
 }
 
 #[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Default, Hash)]
@@ -104,6 +107,15 @@ impl Platform for GethEvmSolcPlatform {
             compiler.map(|compiler| Box::new(compiler) as Box<dyn SolidityCompiler>)
         })
     }
+
+    fn export_genesis(&self, context: Context) -> anyhow::Result<serde_json::Value> {
+        let genesis = AsRef::<GenesisConfiguration>::as_ref(&context).genesis()?;
+        let wallet = AsRef::<WalletConfiguration>::as_ref(&context).wallet();
+
+        let node_genesis = GethNode::node_genesis(genesis.clone(), &wallet);
+        serde_json::to_value(node_genesis)
+            .context("Failed to convert node genesis to a serde_value")
+    }
 }
 
 #[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Default, Hash)]
@@ -149,6 +161,15 @@ impl Platform for LighthouseGethEvmSolcPlatform {
             compiler.map(|compiler| Box::new(compiler) as Box<dyn SolidityCompiler>)
         })
     }
+
+    fn export_genesis(&self, context: Context) -> anyhow::Result<serde_json::Value> {
+        let genesis = AsRef::<GenesisConfiguration>::as_ref(&context).genesis()?;
+        let wallet = AsRef::<WalletConfiguration>::as_ref(&context).wallet();
+
+        let node_genesis = LighthouseGethNode::node_genesis(genesis.clone(), &wallet);
+        serde_json::to_value(node_genesis)
+            .context("Failed to convert node genesis to a serde_value")
+    }
 }
 
 #[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Default, Hash)]
@@ -202,6 +223,16 @@ impl Platform for KitchensinkPolkavmResolcPlatform {
             compiler.map(|compiler| Box::new(compiler) as Box<dyn SolidityCompiler>)
         })
     }
+
+    fn export_genesis(&self, context: Context) -> anyhow::Result<serde_json::Value> {
+        let kitchensink_path = AsRef::<KitchensinkConfiguration>::as_ref(&context)
+            .path
+            .as_path();
+        let wallet = AsRef::<WalletConfiguration>::as_ref(&context).wallet();
+        let export_chainspec_command = SubstrateNode::KITCHENSINK_EXPORT_CHAINSPEC_COMMAND;
+
+        SubstrateNode::node_genesis(kitchensink_path, export_chainspec_command, &wallet)
+    }
 }
 
 #[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Default, Hash)]
@@ -255,6 +286,16 @@ impl Platform for KitchensinkRevmSolcPlatform {
             compiler.map(|compiler| Box::new(compiler) as Box<dyn SolidityCompiler>)
         })
     }
+
+    fn export_genesis(&self, context: Context) -> anyhow::Result<serde_json::Value> {
+        let kitchensink_path = AsRef::<KitchensinkConfiguration>::as_ref(&context)
+            .path
+            .as_path();
+        let wallet = AsRef::<WalletConfiguration>::as_ref(&context).wallet();
+        let export_chainspec_command = SubstrateNode::KITCHENSINK_EXPORT_CHAINSPEC_COMMAND;
+
+        SubstrateNode::node_genesis(kitchensink_path, export_chainspec_command, &wallet)
+    }
 }
 
 #[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Default, Hash)]
@@ -310,6 +351,16 @@ impl Platform for ReviveDevNodePolkavmResolcPlatform {
             compiler.map(|compiler| Box::new(compiler) as Box<dyn SolidityCompiler>)
         })
     }
+
+    fn export_genesis(&self, context: Context) -> anyhow::Result<serde_json::Value> {
+        let revive_dev_node_path = AsRef::<ReviveDevNodeConfiguration>::as_ref(&context)
+            .path
+            .as_path();
+        let wallet = AsRef::<WalletConfiguration>::as_ref(&context).wallet();
+        let export_chainspec_command = SubstrateNode::REVIVE_DEV_NODE_EXPORT_CHAINSPEC_COMMAND;
+
+        SubstrateNode::node_genesis(revive_dev_node_path, export_chainspec_command, &wallet)
+    }
 }
 
 #[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Default, Hash)]
@@ -365,6 +416,16 @@ impl Platform for ReviveDevNodeRevmSolcPlatform {
             compiler.map(|compiler| Box::new(compiler) as Box<dyn SolidityCompiler>)
         })
     }
+
+    fn export_genesis(&self, context: Context) -> anyhow::Result<serde_json::Value> {
+        let revive_dev_node_path = AsRef::<ReviveDevNodeConfiguration>::as_ref(&context)
+            .path
+            .as_path();
+        let wallet = AsRef::<WalletConfiguration>::as_ref(&context).wallet();
+        let export_chainspec_command = SubstrateNode::REVIVE_DEV_NODE_EXPORT_CHAINSPEC_COMMAND;
+
+        SubstrateNode::node_genesis(revive_dev_node_path, export_chainspec_command, &wallet)
+    }
 }
 
 #[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Default, Hash)]
@@ -413,6 +474,15 @@ impl Platform for ZombienetPolkavmResolcPlatform {
             compiler.map(|compiler| Box::new(compiler) as Box<dyn SolidityCompiler>)
         })
     }
+
+    fn export_genesis(&self, context: Context) -> anyhow::Result<serde_json::Value> {
+        let polkadot_parachain_path = AsRef::<PolkadotParachainConfiguration>::as_ref(&context)
+            .path
+            .as_path();
+        let wallet = AsRef::<WalletConfiguration>::as_ref(&context).wallet();
+
+        ZombienetNode::node_genesis(polkadot_parachain_path, &wallet)
+    }
 }
 
 #[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Default, Hash)]
@@ -461,6 +531,15 @@ impl Platform for ZombienetRevmSolcPlatform {
             compiler.map(|compiler| Box::new(compiler) as Box<dyn SolidityCompiler>)
         })
     }
+
+    fn export_genesis(&self, context: Context) -> anyhow::Result<serde_json::Value> {
+        let polkadot_parachain_path = AsRef::<PolkadotParachainConfiguration>::as_ref(&context)
+            .path
+            .as_path();
+        let wallet = AsRef::<WalletConfiguration>::as_ref(&context).wallet();
+
+        ZombienetNode::node_genesis(polkadot_parachain_path, &wallet)
+    }
 }
 
 impl From<PlatformIdentifier> for Box<dyn Platform> {

crates/core/src/main.rs

@@ -2,6 +2,7 @@ mod differential_benchmarks;
 mod differential_tests;
 mod helpers;
 
+use anyhow::Context as _;
 use clap::Parser;
 use revive_dt_report::ReportAggregator;
 use schemars::schema_for;
@@ -72,6 +73,14 @@ fn main() -> anyhow::Result<()> {
 
                 Ok(())
             }),
+        Context::ExportGenesis(ref export_genesis_context) => {
+            let platform = Into::<&dyn Platform>::into(export_genesis_context.platform);
+            let genesis = platform.export_genesis(context)?;
+            let genesis_json = serde_json::to_string_pretty(&genesis)
+                .context("Failed to serialize the genesis to JSON")?;
+            println!("{genesis_json}");
+            Ok(())
+        }
         Context::ExportJsonSchema => {
             let schema = schema_for!(Metadata);
             println!("{}", serde_json::to_string_pretty(&schema).unwrap());

crates/node/src/node_implementations/geth.rs

@@ -130,7 +130,7 @@ impl GethNode {
 
     /// Create the node directory and call `geth init` to configure the genesis.
     #[instrument(level = "info", skip_all, fields(geth_node_id = self.id))]
-    fn init(&mut self, mut genesis: Genesis) -> anyhow::Result<&mut Self> {
+    fn init(&mut self, genesis: Genesis) -> anyhow::Result<&mut Self> {
         let _ = clear_directory(&self.base_directory);
         let _ = clear_directory(&self.logs_directory);
 
@@ -139,16 +139,7 @@ impl GethNode {
         create_dir_all(&self.logs_directory)
             .context("Failed to create logs directory for geth node")?;
 
-        for signer_address in
-            <EthereumWallet as NetworkWallet<Ethereum>>::signer_addresses(&self.wallet)
-        {
-            // Note, the use of the entry API here means that we only modify the entries for any
-            // account that is not in the `alloc` field of the genesis state.
-            genesis
-                .alloc
-                .entry(signer_address)
-                .or_insert(GenesisAccount::default().with_balance(U256::from(INITIAL_BALANCE)));
-        }
+        let genesis = Self::node_genesis(genesis, self.wallet.as_ref());
         let genesis_path = self.base_directory.join(Self::GENESIS_JSON_FILE);
         serde_json::to_writer(
             File::create(&genesis_path).context("Failed to create geth genesis file")?,
@@ -265,6 +256,16 @@ impl GethNode {
             .await
             .cloned()
     }
+
+    pub fn node_genesis(mut genesis: Genesis, wallet: &EthereumWallet) -> Genesis {
+        for signer_address in NetworkWallet::<Ethereum>::signer_addresses(&wallet) {
+            genesis
+                .alloc
+                .entry(signer_address)
+                .or_insert(GenesisAccount::default().with_balance(U256::from(INITIAL_BALANCE)));
+        }
+        genesis
+    }
 }
 
 impl EthereumNode for GethNode {

crates/node/src/node_implementations/lighthouse_geth.rs

@@ -541,6 +541,16 @@ impl LighthouseGethNode {
             .await
         })
     }
+
+    pub fn node_genesis(mut genesis: Genesis, wallet: &EthereumWallet) -> Genesis {
+        for signer_address in NetworkWallet::<Ethereum>::signer_addresses(&wallet) {
+            genesis
+                .alloc
+                .entry(signer_address)
+                .or_insert(GenesisAccount::default().with_balance(U256::from(INITIAL_BALANCE)));
+        }
+        genesis
+    }
 }
 
 impl EthereumNode for LighthouseGethNode {

crates/node/src/node_implementations/substrate.rs

@@ -1,6 +1,6 @@
 use std::{
     fs::{create_dir_all, remove_dir_all},
-    path::PathBuf,
+    path::{Path, PathBuf},
     pin::Pin,
     process::{Command, Stdio},
     sync::{
@@ -12,7 +12,7 @@ use std::{
 
 use alloy::{
     eips::BlockNumberOrTag,
-    genesis::{Genesis, GenesisAccount},
+    genesis::Genesis,
     network::{Ethereum, EthereumWallet, NetworkWallet},
     primitives::{Address, BlockHash, BlockNumber, BlockTimestamp, StorageKey, TxHash, U256},
     providers::{
@@ -32,7 +32,7 @@ use futures::{FutureExt, Stream, StreamExt};
 use revive_common::EVMVersion;
 use revive_dt_common::fs::clear_directory;
 use revive_dt_format::traits::ResolverApi;
-use serde_json::{Value as JsonValue, json};
+use serde_json::json;
 use sp_core::crypto::Ss58Codec;
 use sp_runtime::AccountId32;
 
@@ -129,7 +129,7 @@ impl SubstrateNode {
         }
     }
 
-    fn init(&mut self, mut genesis: Genesis) -> anyhow::Result<&mut Self> {
+    fn init(&mut self, _: Genesis) -> anyhow::Result<&mut Self> {
         let _ = remove_dir_all(self.base_directory.as_path());
         let _ = clear_directory(&self.base_directory);
         let _ = clear_directory(&self.logs_directory);
@@ -141,65 +141,12 @@ impl SubstrateNode {
 
         let template_chainspec_path = self.base_directory.join(Self::CHAIN_SPEC_JSON_FILE);
 
-        // Note: we do not pipe the logs of this process to a separate file since this is just a
-        // once-off export of the default chain spec and not part of the long-running node process.
-        let output = Command::new(&self.node_binary)
-            .arg(self.export_chainspec_command.as_str())
-            .arg("--chain")
-            .arg("dev")
-            .env_remove("RUST_LOG")
-            .output()
-            .context("Failed to export the chain-spec")?;
-
-        if !output.status.success() {
-            anyhow::bail!(
-                "Substrate-node export-chain-spec failed: {}",
-                String::from_utf8_lossy(&output.stderr)
-            );
-        }
-
-        let content = String::from_utf8(output.stdout)
-            .context("Failed to decode Substrate export-chain-spec output as UTF-8")?;
-        let mut chainspec_json: JsonValue =
-            serde_json::from_str(&content).context("Failed to parse Substrate chain spec JSON")?;
-
-        let existing_chainspec_balances =
-            chainspec_json["genesis"]["runtimeGenesis"]["patch"]["balances"]["balances"]
-                .as_array()
-                .cloned()
-                .unwrap_or_default();
-
-        let mut merged_balances: Vec<(String, u128)> = existing_chainspec_balances
-            .into_iter()
-            .filter_map(|val| {
-                if let Some(arr) = val.as_array() {
-                    if arr.len() == 2 {
-                        let account = arr[0].as_str()?.to_string();
-                        let balance = arr[1].as_f64()? as u128;
-                        return Some((account, balance));
-                    }
-                }
-                None
-            })
-            .collect();
-        let mut eth_balances = {
-            for signer_address in
-                <EthereumWallet as NetworkWallet<Ethereum>>::signer_addresses(&self.wallet)
-            {
-                // Note, the use of the entry API here means that we only modify the entries for any
-                // account that is not in the `alloc` field of the genesis state.
-                genesis
-                    .alloc
-                    .entry(signer_address)
-                    .or_insert(GenesisAccount::default().with_balance(U256::from(INITIAL_BALANCE)));
-            }
-            self.extract_balance_from_genesis_file(&genesis)
-                .context("Failed to extract balances from EVM genesis JSON")?
-        };
-        merged_balances.append(&mut eth_balances);
-
-        chainspec_json["genesis"]["runtimeGenesis"]["patch"]["balances"]["balances"] =
-            json!(merged_balances);
+        let chainspec_json = Self::node_genesis(
+            &self.node_binary,
+            &self.export_chainspec_command,
+            &self.wallet,
+        )
+        .context("Failed to prepare the chainspec command")?;
 
         serde_json::to_writer_pretty(
             std::fs::File::create(&template_chainspec_path)
@@ -307,21 +254,6 @@ impl SubstrateNode {
         Ok(())
     }
 
-    fn extract_balance_from_genesis_file(
-        &self,
-        genesis: &Genesis,
-    ) -> anyhow::Result<Vec<(String, u128)>> {
-        genesis
-            .alloc
-            .iter()
-            .try_fold(Vec::new(), |mut vec, (address, acc)| {
-                let substrate_address = Self::eth_to_substrate_address(address);
-                let balance = acc.balance.try_into()?;
-                vec.push((substrate_address, balance));
-                Ok(vec)
-            })
-    }
-
     fn eth_to_substrate_address(address: &Address) -> String {
         let eth_bytes = address.0.0;
 
@@ -360,6 +292,45 @@ impl SubstrateNode {
             .await
             .cloned()
     }
+
+    pub fn node_genesis(
+        node_path: &Path,
+        export_chainspec_command: &str,
+        wallet: &EthereumWallet,
+    ) -> anyhow::Result<serde_json::Value> {
+        let output = Command::new(node_path)
+            .arg(export_chainspec_command)
+            .arg("--chain")
+            .arg("dev")
+            .env_remove("RUST_LOG")
+            .output()
+            .context("Failed to export the chain-spec")?;
+
+        if !output.status.success() {
+            anyhow::bail!(
+                "Substrate-node export-chain-spec failed: {}",
+                String::from_utf8_lossy(&output.stderr)
+            );
+        }
+
+        let content = String::from_utf8(output.stdout)
+            .context("Failed to decode Substrate export-chain-spec output as UTF-8")?;
+        let mut chainspec_json = serde_json::from_str::<serde_json::Value>(&content)
+            .context("Failed to parse Substrate chain spec JSON")?;
+
+        let existing_chainspec_balances =
+            chainspec_json["genesis"]["runtimeGenesis"]["patch"]["balances"]["balances"]
+                .as_array_mut()
+                .expect("Can't fail");
+
+        for address in NetworkWallet::<Ethereum>::signer_addresses(wallet) {
+            let substrate_address = Self::eth_to_substrate_address(&address);
+            let balance = INITIAL_BALANCE;
+            existing_chainspec_balances.push(json!((substrate_address, balance)));
+        }
+
+        Ok(chainspec_json)
+    }
 }
 
 impl EthereumNode for SubstrateNode {
@@ -900,50 +871,6 @@ mod tests {
         );
     }
 
-    #[test]
-    #[ignore = "Ignored since they take a long time to run"]
-    fn test_parse_genesis_alloc() {
-        // Create test genesis file
-        let genesis_json = r#"
-        {
-          "alloc": {
-            "0x90F8bf6A479f320ead074411a4B0e7944Ea8c9C1": { "balance": "1000000000000000000" },
-            "0x0000000000000000000000000000000000000000": { "balance": "0xDE0B6B3A7640000" },
-            "0xffffffffffffffffffffffffffffffffffffffff": { "balance": "123456789" }
-          }
-        }
-        "#;
-
-        let context = test_config();
-        let node = SubstrateNode::new(
-            context.kitchensink_configuration.path.clone(),
-            SubstrateNode::KITCHENSINK_EXPORT_CHAINSPEC_COMMAND,
-            None,
-            &context,
-        );
-
-        let result = node
-            .extract_balance_from_genesis_file(&serde_json::from_str(genesis_json).unwrap())
-            .unwrap();
-
-        let result_map: std::collections::HashMap<_, _> = result.into_iter().collect();
-
-        assert_eq!(
-            result_map.get("5FLneRcWAfk3X3tg6PuGyLNGAquPAZez5gpqvyuf3yUK8VaV"),
-            Some(&1_000_000_000_000_000_000u128)
-        );
-
-        assert_eq!(
-            result_map.get("5C4hrfjw9DjXZTzV3MwzrrAr9P1MLDHajjSidz9bR544LEq1"),
-            Some(&1_000_000_000_000_000_000u128)
-        );
-
-        assert_eq!(
-            result_map.get("5HrN7fHLXWcFiXPwwtq2EkSGns9eMmoUQnbVKweNz3VVr6N4"),
-            Some(&123_456_789u128)
-        );
-    }
-
     #[test]
     #[ignore = "Ignored since they take a long time to run"]
     fn print_eth_to_substrate_mappings() {

crates/node/src/node_implementations/zombienet.rs

@@ -28,7 +28,7 @@
 
 use std::{
     fs::{create_dir_all, remove_dir_all},
-    path::PathBuf,
+    path::{Path, PathBuf},
     pin::Pin,
     process::{Command, Stdio},
     sync::{
@@ -40,7 +40,7 @@ use std::{
 
 use alloy::{
     eips::BlockNumberOrTag,
-    genesis::{Genesis, GenesisAccount},
+    genesis::Genesis,
     network::{Ethereum, EthereumWallet, NetworkWallet},
     primitives::{Address, BlockHash, BlockNumber, BlockTimestamp, StorageKey, TxHash, U256},
     providers::{
@@ -61,7 +61,7 @@ use revive_dt_common::fs::clear_directory;
 use revive_dt_config::*;
 use revive_dt_format::traits::ResolverApi;
 use revive_dt_node_interaction::{EthereumNode, MinedBlockInformation};
-use serde_json::{Value as JsonValue, json};
+use serde_json::json;
 use sp_core::crypto::Ss58Codec;
 use sp_runtime::AccountId32;
 use subxt::{OnlineClient, SubstrateConfig};
@@ -164,7 +164,7 @@ impl ZombienetNode {
         }
     }
 
-    fn init(&mut self, genesis: Genesis) -> anyhow::Result<&mut Self> {
+    fn init(&mut self, _: Genesis) -> anyhow::Result<&mut Self> {
         let _ = clear_directory(&self.base_directory);
         let _ = clear_directory(&self.logs_directory);
 
@@ -174,7 +174,7 @@ impl ZombienetNode {
             .context("Failed to create logs directory for zombie node")?;
 
         let template_chainspec_path = self.base_directory.join(Self::CHAIN_SPEC_JSON_FILE);
-        self.prepare_chainspec(template_chainspec_path.clone(), genesis)?;
+        self.prepare_chainspec(template_chainspec_path.clone())?;
         let polkadot_parachain_path = self
             .polkadot_parachain_path
             .to_str()
@@ -286,71 +286,9 @@ impl ZombienetNode {
         Ok(())
     }
 
-    fn prepare_chainspec(
-        &mut self,
-        template_chainspec_path: PathBuf,
-        mut genesis: Genesis,
-    ) -> anyhow::Result<()> {
-        let output = Command::new(self.polkadot_parachain_path.as_path())
-            .arg(Self::EXPORT_CHAINSPEC_COMMAND)
-            .arg("--chain")
-            .arg("asset-hub-westend-local")
-            .env_remove("RUST_LOG")
-            .output()
-            .context("Failed to export the chainspec of the chain")?;
-
-        if !output.status.success() {
-            anyhow::bail!(
-                "Build chain-spec failed: {}",
-                String::from_utf8_lossy(&output.stderr)
-            );
-        }
-
-        let content = String::from_utf8(output.stdout)
-            .context("Failed to decode collators chain-spec output as UTF-8")?;
-        let mut chainspec_json: JsonValue =
-            serde_json::from_str(&content).context("Failed to parse collators chain spec JSON")?;
-
-        let existing_chainspec_balances =
-            chainspec_json["genesis"]["runtimeGenesis"]["patch"]["balances"]["balances"]
-                .as_array()
-                .cloned()
-                .unwrap_or_default();
-
-        let mut merged_balances: Vec<(String, u128)> = existing_chainspec_balances
-            .into_iter()
-            .filter_map(|val| {
-                if let Some(arr) = val.as_array() {
-                    if arr.len() == 2 {
-                        let account = arr[0].as_str()?.to_string();
-                        let balance = arr[1].as_f64()? as u128;
-                        return Some((account, balance));
-                    }
-                }
-                None
-            })
-            .collect();
-
-        let mut eth_balances = {
-            for signer_address in
-                <EthereumWallet as NetworkWallet<Ethereum>>::signer_addresses(&self.wallet)
-            {
-                // Note, the use of the entry API here means that we only modify the entries for any
-                // account that is not in the `alloc` field of the genesis state.
-                genesis
-                    .alloc
-                    .entry(signer_address)
-                    .or_insert(GenesisAccount::default().with_balance(U256::from(INITIAL_BALANCE)));
-            }
-            self.extract_balance_from_genesis_file(&genesis)
-                .context("Failed to extract balances from EVM genesis JSON")?
-        };
-
-        merged_balances.append(&mut eth_balances);
-
-        chainspec_json["genesis"]["runtimeGenesis"]["patch"]["balances"]["balances"] =
-            json!(merged_balances);
-
+    fn prepare_chainspec(&mut self, template_chainspec_path: PathBuf) -> anyhow::Result<()> {
+        let chainspec_json = Self::node_genesis(&self.polkadot_parachain_path, &self.wallet)
+            .context("Failed to prepare the zombienet chainspec file")?;
         let writer = std::fs::File::create(&template_chainspec_path)
             .context("Failed to create template chainspec file")?;
 
@@ -360,21 +298,6 @@ impl ZombienetNode {
         Ok(())
     }
 
-    fn extract_balance_from_genesis_file(
-        &self,
-        genesis: &Genesis,
-    ) -> anyhow::Result<Vec<(String, u128)>> {
-        genesis
-            .alloc
-            .iter()
-            .try_fold(Vec::new(), |mut vec, (address, acc)| {
-                let polkadot_address = Self::eth_to_polkadot_address(address);
-                let balance = acc.balance.try_into()?;
-                vec.push((polkadot_address, balance));
-                Ok(vec)
-            })
-    }
-
     fn eth_to_polkadot_address(address: &Address) -> String {
         let eth_bytes = address.0.0;
 
@@ -414,6 +337,44 @@ impl ZombienetNode {
             .await
             .cloned()
     }
+
+    pub fn node_genesis(
+        node_path: &Path,
+        wallet: &EthereumWallet,
+    ) -> anyhow::Result<serde_json::Value> {
+        let output = Command::new(node_path)
+            .arg(Self::EXPORT_CHAINSPEC_COMMAND)
+            .arg("--chain")
+            .arg("asset-hub-westend-local")
+            .env_remove("RUST_LOG")
+            .output()
+            .context("Failed to export the chainspec of the chain")?;
+
+        if !output.status.success() {
+            anyhow::bail!(
+                "Substrate-node export-chain-spec failed: {}",
+                String::from_utf8_lossy(&output.stderr)
+            );
+        }
+
+        let content = String::from_utf8(output.stdout)
+            .context("Failed to decode Substrate export-chain-spec output as UTF-8")?;
+        let mut chainspec_json = serde_json::from_str::<serde_json::Value>(&content)
+            .context("Failed to parse Substrate chain spec JSON")?;
+
+        let existing_chainspec_balances =
+            chainspec_json["genesis"]["runtimeGenesis"]["patch"]["balances"]["balances"]
+                .as_array_mut()
+                .expect("Can't fail");
+
+        for address in NetworkWallet::<Ethereum>::signer_addresses(wallet) {
+            let substrate_address = Self::eth_to_polkadot_address(&address);
+            let balance = INITIAL_BALANCE;
+            existing_chainspec_balances.push(json!((substrate_address, balance)));
+        }
+
+        Ok(chainspec_json)
+    }
 }
 
 impl EthereumNode for ZombienetNode {
@@ -956,47 +917,6 @@ mod tests {
         );
     }
 
-    #[tokio::test]
-    async fn test_parse_genesis_alloc() {
-        // Create test genesis file
-        let genesis_json = r#"
-        {
-          "alloc": {
-            "0x90F8bf6A479f320ead074411a4B0e7944Ea8c9C1": { "balance": "1000000000000000000" },
-            "0x0000000000000000000000000000000000000000": { "balance": "0xDE0B6B3A7640000" },
-            "0xffffffffffffffffffffffffffffffffffffffff": { "balance": "123456789" }
-          }
-        }
-        "#;
-
-        let context = test_config();
-        let node = ZombienetNode::new(
-            context.polkadot_parachain_configuration.path.clone(),
-            &context,
-        );
-
-        let result = node
-            .extract_balance_from_genesis_file(&serde_json::from_str(genesis_json).unwrap())
-            .unwrap();
-
-        let result_map: std::collections::HashMap<_, _> = result.into_iter().collect();
-
-        assert_eq!(
-            result_map.get("5FLneRcWAfk3X3tg6PuGyLNGAquPAZez5gpqvyuf3yUK8VaV"),
-            Some(&1_000_000_000_000_000_000u128)
-        );
-
-        assert_eq!(
-            result_map.get("5C4hrfjw9DjXZTzV3MwzrrAr9P1MLDHajjSidz9bR544LEq1"),
-            Some(&1_000_000_000_000_000_000u128)
-        );
-
-        assert_eq!(
-            result_map.get("5HrN7fHLXWcFiXPwwtq2EkSGns9eMmoUQnbVKweNz3VVr6N4"),
-            Some(&123_456_789u128)
-        );
-    }
-
     #[test]
     fn print_eth_to_polkadot_mappings() {
         let eth_addresses = vec![