NDEV-2438: Simplify Rpc trait usage

evm_loader/lib/src/account_storage.rs

@@ -25,8 +25,7 @@ use serde::{Deserialize, Serialize};
 use solana_client::client_error;
 use solana_sdk::{account::Account, account_info::AccountInfo, pubkey, pubkey::Pubkey};
 
-use crate::commands::get_config::ChainInfo;
-use crate::rpc::RpcEnum;
+use crate::commands::get_config::{BuildConfigSimulator, ChainInfo};
 use crate::tracing::{AccountOverride, AccountOverrides, BlockOverrides};
 use serde_with::{serde_as, DisplayFromStr};
 
@@ -44,25 +43,25 @@ pub struct SolanaAccount {
 }
 
 #[allow(clippy::module_name_repetitions)]
-pub struct EmulatorAccountStorage<'rpc> {
+pub struct EmulatorAccountStorage<'rpc, T: Rpc> {
     pub accounts: RefCell<HashMap<Pubkey, SolanaAccount>>,
     pub gas: u64,
-    rpc: &'rpc RpcEnum,
+    rpc: &'rpc T,
     program_id: Pubkey,
     chains: Vec<ChainInfo>,
     block_number: u64,
     block_timestamp: i64,
     state_overrides: Option<AccountOverrides>,
 }
 
-impl<'rpc> EmulatorAccountStorage<'rpc> {
+impl<'rpc, T: Rpc + BuildConfigSimulator> EmulatorAccountStorage<'rpc, T> {
     pub async fn new(
-        rpc: &'rpc RpcEnum,
+        rpc: &'rpc T,
         program_id: Pubkey,
         chains: Option<Vec<ChainInfo>>,
         block_overrides: Option<BlockOverrides>,
         state_overrides: Option<AccountOverrides>,
-    ) -> Result<EmulatorAccountStorage<'rpc>, NeonError> {
+    ) -> Result<EmulatorAccountStorage<T>, NeonError> {
         trace!("backend::new");
 
         let block_number = match block_overrides.as_ref().and_then(|o| o.number) {
@@ -93,20 +92,22 @@ impl<'rpc> EmulatorAccountStorage<'rpc> {
     }
 
     pub async fn with_accounts(
-        rpc: &'rpc RpcEnum,
+        rpc: &'rpc T,
         program_id: Pubkey,
         accounts: &[Pubkey],
         chains: Option<Vec<ChainInfo>>,
         block_overrides: Option<BlockOverrides>,
         state_overrides: Option<AccountOverrides>,
-    ) -> Result<EmulatorAccountStorage<'rpc>, NeonError> {
+    ) -> Result<EmulatorAccountStorage<'rpc, T>, NeonError> {
         let storage = Self::new(rpc, program_id, chains, block_overrides, state_overrides).await?;
 
         storage.download_accounts(accounts).await?;
 
         Ok(storage)
     }
+}
 
+impl<T: Rpc> EmulatorAccountStorage<'_, T> {
     async fn download_accounts(&self, pubkeys: &[Pubkey]) -> Result<(), NeonError> {
         let accounts = self.rpc.get_multiple_accounts(pubkeys).await?;
 
@@ -503,8 +504,8 @@ impl<'rpc> EmulatorAccountStorage<'rpc> {
     }
 }
 
-#[async_trait(? Send)]
-impl<'a> AccountStorage for EmulatorAccountStorage<'a> {
+#[async_trait(?Send)]
+impl<T: Rpc> AccountStorage for EmulatorAccountStorage<'_, T> {
     fn program_id(&self) -> &Pubkey {
         debug!("program_id");
         &self.program_id

evm_loader/lib/src/commands/collect_treasury.rs

@@ -26,7 +26,7 @@ pub async fn execute(
     rpc_client: &CloneRpcClient,
     signer: &dyn Signer,
 ) -> NeonResult<CollectTreasuryReturn> {
-    let neon_params = read_elf_parameters_from_account(config, &rpc_client.clone().into()).await?;
+    let neon_params = read_elf_parameters_from_account(config, rpc_client).await?;
 
     let pool_count: u32 = neon_params
         .get("NEON_TREASURY_POOL_COUNT")

evm_loader/lib/src/commands/emulate.rs

@@ -5,7 +5,8 @@ use serde::{Deserialize, Serialize};
 use solana_sdk::entrypoint::MAX_PERMITTED_DATA_INCREASE;
 use solana_sdk::pubkey::Pubkey;
 
-use crate::rpc::RpcEnum;
+use crate::commands::get_config::BuildConfigSimulator;
+use crate::rpc::Rpc;
 use crate::syscall_stubs::setup_emulator_syscall_stubs;
 use crate::types::{EmulateRequest, TxParams};
 use crate::{
@@ -50,39 +51,39 @@ impl EmulateResponse {
 }
 
 pub async fn execute(
-    rpc: &RpcEnum,
+    rpc: &(impl Rpc + BuildConfigSimulator),
     program_id: Pubkey,
-    config: EmulateRequest,
+    emulate_request: EmulateRequest,
     tracer: Option<TracerType>,
 ) -> NeonResult<EmulateResponse> {
-    let block_overrides = config
+    let block_overrides = emulate_request
         .trace_config
         .as_ref()
         .and_then(|t| t.block_overrides.clone());
-    let state_overrides = config
+    let state_overrides = emulate_request
         .trace_config
         .as_ref()
         .and_then(|t| t.state_overrides.clone());
 
     let mut storage = EmulatorAccountStorage::with_accounts(
         rpc,
         program_id,
-        &config.accounts,
-        config.chains,
+        &emulate_request.accounts,
+        emulate_request.chains,
         block_overrides,
         state_overrides,
     )
     .await?;
 
-    let step_limit = config.step_limit.unwrap_or(100000);
+    let step_limit = emulate_request.step_limit.unwrap_or(100000);
 
     setup_emulator_syscall_stubs(rpc).await?;
-    emulate_trx(config.tx, &mut storage, step_limit, tracer).await
+    emulate_trx(emulate_request.tx, &mut storage, step_limit, tracer).await
 }
 
 async fn emulate_trx(
     tx_params: TxParams,
-    storage: &mut EmulatorAccountStorage<'_>,
+    storage: &mut EmulatorAccountStorage<'_, impl Rpc>,
     step_limit: u64,
     tracer: Option<TracerType>,
 ) -> NeonResult<EmulateResponse> {

evm_loader/lib/src/commands/get_balance.rs

@@ -6,10 +6,9 @@ use solana_sdk::{account::Account, pubkey::Pubkey};
 
 use crate::{account_storage::account_info, rpc::Rpc, types::BalanceAddress, NeonResult};
 
-use crate::rpc::RpcEnum;
 use serde_with::{serde_as, DisplayFromStr};
 
-use super::get_config::ChainInfo;
+use super::get_config::{BuildConfigSimulator, ChainInfo};
 
 #[derive(Debug, Serialize, Deserialize)]
 pub enum BalanceStatus {
@@ -92,7 +91,7 @@ fn is_legacy_chain_id(id: u64, chains: &[ChainInfo]) -> bool {
 }
 
 pub async fn execute(
-    rpc: &RpcEnum,
+    rpc: &(impl Rpc + BuildConfigSimulator),
     program_id: &Pubkey,
     address: &[BalanceAddress],
 ) -> NeonResult<Vec<GetBalanceResponse>> {

evm_loader/lib/src/commands/get_config.rs

@@ -1,4 +1,6 @@
+use async_trait::async_trait;
 use base64::Engine;
+use enum_dispatch::enum_dispatch;
 use std::collections::BTreeMap;
 use tokio::sync::{Mutex, MutexGuard, OnceCell};
 
@@ -18,7 +20,7 @@ use solana_sdk::{
 
 use crate::{rpc::Rpc, NeonError, NeonResult};
 
-use crate::rpc::{CallDbClient, RpcEnum, SolanaRpc};
+use crate::rpc::{CallDbClient, CloneRpcClient, SolanaRpc};
 use serde_with::{serde_as, DisplayFromStr};
 use solana_client::nonblocking::rpc_client::RpcClient;
 
@@ -114,30 +116,39 @@ async fn lock_program_test(
     context
 }
 
-enum ConfigSimulator<'r> {
+pub enum ConfigSimulator<'r> {
     Rpc(Pubkey, &'r RpcClient),
     ProgramTest(MutexGuard<'static, ProgramTestContext>),
 }
 
-impl<'r> ConfigSimulator<'r> {
-    pub async fn new(rpc: &'r RpcEnum, program_id: Pubkey) -> NeonResult<ConfigSimulator<'r>> {
-        let simulator = match rpc {
-            RpcEnum::CloneRpcClient(clone_rpc_client) => Self::Rpc(
-                clone_rpc_client.get_account_with_sol().await?,
-                clone_rpc_client,
-            ),
-            RpcEnum::CallDbClient(call_db_client) => {
-                let program_data =
-                    read_program_data_from_account(call_db_client, program_id).await?;
-                let mut program_test = lock_program_test(program_id, program_data).await;
-                program_test.get_new_latest_blockhash().await?;
-
-                Self::ProgramTest(program_test)
-            }
-        };
-        Ok(simulator)
+#[async_trait(?Send)]
+#[enum_dispatch]
+pub trait BuildConfigSimulator {
+    async fn build_config_simulator(&self, program_id: Pubkey) -> NeonResult<ConfigSimulator>;
+}
+
+#[async_trait(?Send)]
+impl BuildConfigSimulator for CloneRpcClient {
+    async fn build_config_simulator(&self, _program_id: Pubkey) -> NeonResult<ConfigSimulator> {
+        Ok(ConfigSimulator::Rpc(
+            self.get_account_with_sol().await?,
+            self,
+        ))
     }
+}
+
+#[async_trait(?Send)]
+impl BuildConfigSimulator for CallDbClient {
+    async fn build_config_simulator(&self, program_id: Pubkey) -> NeonResult<ConfigSimulator> {
+        let program_data = read_program_data_from_account(self, program_id).await?;
+        let mut program_test = lock_program_test(program_id, program_data).await;
+        program_test.get_new_latest_blockhash().await?;
 
+        Ok(ConfigSimulator::ProgramTest(program_test))
+    }
+}
+
+impl ConfigSimulator<'_> {
     async fn simulate_config(
         &mut self,
         program_id: Pubkey,
@@ -195,97 +206,90 @@ impl<'r> ConfigSimulator<'r> {
 
         Ok(return_data)
     }
-}
 
-async fn get_version(
-    context: &mut ConfigSimulator<'_>,
-    program_id: Pubkey,
-) -> NeonResult<(String, String)> {
-    let return_data = context.simulate_config(program_id, 0xA7, &[]).await?;
-    let (version, revision) = bincode::deserialize(&return_data)?;
+    async fn get_version(&mut self, program_id: Pubkey) -> NeonResult<(String, String)> {
+        let return_data = self.simulate_config(program_id, 0xA7, &[]).await?;
+        let (version, revision) = bincode::deserialize(&return_data)?;
 
-    Ok((version, revision))
-}
+        Ok((version, revision))
+    }
 
-async fn get_status(context: &mut ConfigSimulator<'_>, program_id: Pubkey) -> NeonResult<Status> {
-    let return_data = context.simulate_config(program_id, 0xA6, &[]).await?;
-    match return_data[0] {
-        0 => Ok(Status::Emergency),
-        1 => Ok(Status::Ok),
-        _ => Ok(Status::Unknown),
+    async fn get_status(&mut self, program_id: Pubkey) -> NeonResult<Status> {
+        let return_data = self.simulate_config(program_id, 0xA6, &[]).await?;
+        match return_data[0] {
+            0 => Ok(Status::Emergency),
+            1 => Ok(Status::Ok),
+            _ => Ok(Status::Unknown),
+        }
     }
-}
 
-async fn get_environment(
-    context: &mut ConfigSimulator<'_>,
-    program_id: Pubkey,
-) -> NeonResult<String> {
-    let return_data = context.simulate_config(program_id, 0xA2, &[]).await?;
-    let environment = String::from_utf8(return_data)?;
+    async fn get_environment(&mut self, program_id: Pubkey) -> NeonResult<String> {
+        let return_data = self.simulate_config(program_id, 0xA2, &[]).await?;
+        let environment = String::from_utf8(return_data)?;
 
-    Ok(environment)
-}
+        Ok(environment)
+    }
 
-async fn get_chains(
-    context: &mut ConfigSimulator<'_>,
-    program_id: Pubkey,
-) -> NeonResult<Vec<ChainInfo>> {
-    let mut result = Vec::new();
+    async fn get_chains(&mut self, program_id: Pubkey) -> NeonResult<Vec<ChainInfo>> {
+        let mut result = Vec::new();
+
+        let return_data = self.simulate_config(program_id, 0xA0, &[]).await?;
+        let chain_count = return_data.as_slice().try_into()?;
+        let chain_count = usize::from_le_bytes(chain_count);
 
-    let return_data = context.simulate_config(program_id, 0xA0, &[]).await?;
-    let chain_count = return_data.as_slice().try_into()?;
-    let chain_count = usize::from_le_bytes(chain_count);
+        for i in 0..chain_count {
+            let index = i.to_le_bytes();
+            let return_data = self.simulate_config(program_id, 0xA1, &index).await?;
 
-    for i in 0..chain_count {
-        let index = i.to_le_bytes();
-        let return_data = context.simulate_config(program_id, 0xA1, &index).await?;
+            let (id, name, token) = bincode::deserialize(&return_data)?;
+            result.push(ChainInfo { id, name, token });
+        }
 
-        let (id, name, token) = bincode::deserialize(&return_data)?;
-        result.push(ChainInfo { id, name, token });
+        Ok(result)
     }
 
-    Ok(result)
-}
+    async fn get_properties(&mut self, program_id: Pubkey) -> NeonResult<BTreeMap<String, String>> {
+        let mut result = BTreeMap::new();
 
-async fn get_properties(
-    context: &mut ConfigSimulator<'_>,
-    program_id: Pubkey,
-) -> NeonResult<BTreeMap<String, String>> {
-    let mut result = BTreeMap::new();
+        let return_data = self.simulate_config(program_id, 0xA3, &[]).await?;
+        let count = return_data.as_slice().try_into()?;
+        let count = usize::from_le_bytes(count);
 
-    let return_data = context.simulate_config(program_id, 0xA3, &[]).await?;
-    let count = return_data.as_slice().try_into()?;
-    let count = usize::from_le_bytes(count);
+        for i in 0..count {
+            let index = i.to_le_bytes();
+            let return_data = self.simulate_config(program_id, 0xA4, &index).await?;
 
-    for i in 0..count {
-        let index = i.to_le_bytes();
-        let return_data = context.simulate_config(program_id, 0xA4, &index).await?;
+            let (name, value) = bincode::deserialize(&return_data)?;
+            result.insert(name, value);
+        }
 
-        let (name, value) = bincode::deserialize(&return_data)?;
-        result.insert(name, value);
+        Ok(result)
     }
-
-    Ok(result)
 }
 
-pub async fn execute(rpc: &RpcEnum, program_id: Pubkey) -> NeonResult<GetConfigResponse> {
-    let mut simulator = ConfigSimulator::new(rpc, program_id).await?;
+pub async fn execute(
+    rpc: &impl BuildConfigSimulator,
+    program_id: Pubkey,
+) -> NeonResult<GetConfigResponse> {
+    let mut simulator = rpc.build_config_simulator(program_id).await?;
 
-    let (version, revision) = get_version(&mut simulator, program_id).await?;
+    let (version, revision) = simulator.get_version(program_id).await?;
 
     Ok(GetConfigResponse {
         version,
         revision,
-        status: get_status(&mut simulator, program_id).await?,
-        environment: get_environment(&mut simulator, program_id).await?,
-        chains: get_chains(&mut simulator, program_id).await?,
-        config: get_properties(&mut simulator, program_id).await?,
+        status: simulator.get_status(program_id).await?,
+        environment: simulator.get_environment(program_id).await?,
+        chains: simulator.get_chains(program_id).await?,
+        config: simulator.get_properties(program_id).await?,
     })
 }
 
-pub async fn read_chains(rpc: &RpcEnum, program_id: Pubkey) -> NeonResult<Vec<ChainInfo>> {
-    let mut simulator = ConfigSimulator::new(rpc, program_id).await?;
+pub async fn read_chains(
+    rpc: &impl BuildConfigSimulator,
+    program_id: Pubkey,
+) -> NeonResult<Vec<ChainInfo>> {
+    let mut simulator = rpc.build_config_simulator(program_id).await?;
 
-    let chains = get_chains(&mut simulator, program_id).await?;
-    Ok(chains)
+    simulator.get_chains(program_id).await
 }