remove NetworkArg type

yetanotherco · Dec 10, 2024 · f6caa0d · f6caa0d
1 parent eadd8ae
commit f6caa0d
Show file tree

Hide file tree

Showing 4 changed files with 44 additions and 60 deletions.
diff --git a/batcher/aligned-sdk/Cargo.toml b/batcher/aligned-sdk/Cargo.toml
@@ -25,3 +25,4 @@ hex = "0.4.3"
 ciborium = "=0.2.2"
 serde_repr = "0.1.19"
 dialoguer = "0.11.0"
+clap = { version = "4.5.4", features = ["derive"] }
diff --git a/batcher/aligned-sdk/src/core/types.rs b/batcher/aligned-sdk/src/core/types.rs
@@ -3,6 +3,7 @@ use std::fmt::Display;
 use std::fmt::Formatter;
 use std::str::FromStr;
 
+use clap::ValueEnum;
 use ethers::core::k256::ecdsa::SigningKey;
 use ethers::signers::Signer;
 use ethers::signers::Wallet;
@@ -11,6 +12,7 @@ use ethers::types::transaction::eip712::Eip712;
 use ethers::types::transaction::eip712::Eip712Error;
 use ethers::types::Address;
 use ethers::types::Signature;
+use ethers::types::H160;
 use ethers::types::U256;
 use lambdaworks_crypto::merkle_tree::{
     merkle::MerkleTree, proof::Proof, traits::IsMerkleTreeBackend,
@@ -396,14 +398,38 @@ pub enum GetNonceResponseMessage {
     InvalidRequest(String),
 }
 
-#[derive(Debug, Clone, Copy)]
+#[derive(Debug, Clone, Copy, ValueEnum)]
 pub enum Network {
     Devnet,
     Holesky,
     HoleskyStage,
     Mainnet,
 }
 
+impl Network {
+    pub fn get_batcher_payment_service_address(&self) -> ethers::types::H160 {
+        match self {
+            Self::Devnet => H160::from_str("0x7bc06c482DEAd17c0e297aFbC32f6e63d3846650").unwrap(),
+            Self::Holesky => H160::from_str("0x815aeCA64a974297942D2Bbf034ABEe22a38A003").unwrap(),
+            Self::HoleskyStage => {
+                H160::from_str("0x7577Ec4ccC1E6C529162ec8019A49C13F6DAd98b").unwrap()
+            }
+            Self::Mainnet => H160::from_str("0xb0567184A52cB40956df6333510d6eF35B89C8de").unwrap(),
+        }
+    }
+
+    pub fn get_aligned_service_manager_address(&self) -> ethers::types::H160 {
+        match self {
+            Self::Devnet => H160::from_str("0x1613beB3B2C4f22Ee086B2b38C1476A3cE7f78E8").unwrap(),
+            Self::Holesky => H160::from_str("0x58F280BeBE9B34c9939C3C39e0890C81f163B623").unwrap(),
+            Self::HoleskyStage => {
+                H160::from_str("0x9C5231FC88059C086Ea95712d105A2026048c39B").unwrap()
+            }
+            Self::Mainnet => H160::from_str("0xeF2A435e5EE44B2041100EF8cbC8ae035166606c").unwrap(),
+        }
+    }
+}
+
 impl FromStr for Network {
     type Err = String;
 

diff --git a/batcher/aligned-sdk/src/sdk.rs b/batcher/aligned-sdk/src/sdk.rs
@@ -28,10 +28,10 @@ use ethers::{
     prelude::k256::ecdsa::SigningKey,
     providers::{Http, Middleware, Provider},
     signers::{LocalWallet, Wallet},
-    types::{Address, H160, U256},
+    types::{Address, U256},
 };
 use sha3::{Digest, Keccak256};
-use std::{str::FromStr, sync::Arc};
+use std::sync::Arc;
 use tokio::{net::TcpStream, sync::Mutex};
 use tokio_tungstenite::{connect_async, tungstenite::Message, MaybeTlsStream, WebSocketStream};
 
@@ -271,28 +271,6 @@ pub async fn submit_multiple(
     .await
 }
 
-pub fn get_payment_service_address(network: Network) -> ethers::types::H160 {
-    match network {
-        Network::Devnet => H160::from_str("0x7bc06c482DEAd17c0e297aFbC32f6e63d3846650").unwrap(),
-        Network::Holesky => H160::from_str("0x815aeCA64a974297942D2Bbf034ABEe22a38A003").unwrap(),
-        Network::HoleskyStage => {
-            H160::from_str("0x7577Ec4ccC1E6C529162ec8019A49C13F6DAd98b").unwrap()
-        }
-        Network::Mainnet => H160::from_str("0xb0567184A52cB40956df6333510d6eF35B89C8de").unwrap(),
-    }
-}
-
-pub fn get_aligned_service_manager_address(network: Network) -> ethers::types::H160 {
-    match network {
-        Network::Devnet => H160::from_str("0x1613beB3B2C4f22Ee086B2b38C1476A3cE7f78E8").unwrap(),
-        Network::Holesky => H160::from_str("0x58F280BeBE9B34c9939C3C39e0890C81f163B623").unwrap(),
-        Network::HoleskyStage => {
-            H160::from_str("0x9C5231FC88059C086Ea95712d105A2026048c39B").unwrap()
-        }
-        Network::Mainnet => H160::from_str("0xeF2A435e5EE44B2041100EF8cbC8ae035166606c").unwrap(),
-    }
-}
-
 // Will submit the proofs to the batcher and wait for their responses
 // Will return once all proofs are responded, or up to a proof that is responded with an error
 async fn _submit_multiple(
@@ -328,7 +306,7 @@ async fn _submit_multiple(
 
     let response_stream = Arc::new(Mutex::new(response_stream));
 
-    let payment_service_addr = get_payment_service_address(network);
+    let payment_service_addr = network.get_batcher_payment_service_address();
 
     let result = async {
         let sent_verification_data_rev = send_messages(
@@ -498,8 +476,8 @@ async fn _is_proof_verified(
     network: Network,
     eth_rpc_provider: Provider<Http>,
 ) -> Result<bool, errors::VerificationError> {
-    let contract_address = get_aligned_service_manager_address(network);
-    let payment_service_addr = get_payment_service_address(network);
+    let contract_address = network.get_aligned_service_manager_address();
+    let payment_service_addr = network.get_batcher_payment_service_address();
 
     // All the elements from the merkle proof have to be concatenated
     let merkle_proof: Vec<u8> = aligned_verification_data
@@ -640,7 +618,7 @@ pub async fn get_nonce_from_ethereum(
     let eth_rpc_provider = Provider::<Http>::try_from(eth_rpc_url)
         .map_err(|e| GetNonceError::EthRpcError(e.to_string()))?;
 
-    let payment_service_address = get_payment_service_address(network);
+    let payment_service_address = network.get_batcher_payment_service_address();
 
     match batcher_payment_service(eth_rpc_provider, payment_service_address).await {
         Ok(contract) => {
@@ -691,7 +669,7 @@ pub async fn deposit_to_aligned(
     signer: SignerMiddleware<Provider<Http>, LocalWallet>,
     network: Network,
 ) -> Result<ethers::types::TransactionReceipt, errors::PaymentError> {
-    let payment_service_address = get_payment_service_address(network);
+    let payment_service_address = network.get_batcher_payment_service_address();
     let from = signer.address();
 
     let tx = TransactionRequest::new()
@@ -729,7 +707,7 @@ pub async fn get_balance_in_aligned(
     let eth_rpc_provider = Provider::<Http>::try_from(eth_rpc_url)
         .map_err(|e| errors::BalanceError::EthereumProviderError(e.to_string()))?;
 
-    let payment_service_address = get_payment_service_address(network);
+    let payment_service_address = network.get_batcher_payment_service_address();
 
     match batcher_payment_service(eth_rpc_provider, payment_service_address).await {
         Ok(batcher_payment_service) => {

diff --git a/batcher/aligned/src/main.rs b/batcher/aligned/src/main.rs
@@ -120,7 +120,7 @@ pub struct SubmitArgs {
         long = "network",
         default_value = "devnet"
     )]
-    network: NetworkArg,
+    network: Network,
 }
 
 #[derive(Parser, Debug)]
@@ -143,7 +143,7 @@ pub struct DepositToBatcherArgs {
         long = "network",
         default_value = "devnet"
     )]
-    network: NetworkArg,
+    network: Network,
     #[arg(name = "Amount to deposit", long = "amount", required = true)]
     amount: String,
 }
@@ -164,7 +164,7 @@ pub struct VerifyProofOnchainArgs {
         long = "network",
         default_value = "devnet"
     )]
-    network: NetworkArg,
+    network: Network,
 }
 
 #[derive(Parser, Debug)]
@@ -186,7 +186,7 @@ pub struct GetUserBalanceArgs {
         long = "network",
         default_value = "devnet"
     )]
-    network: NetworkArg,
+    network: Network,
     #[arg(
         name = "Ethereum RPC provider address",
         long = "rpc_url",
@@ -218,25 +218,6 @@ pub struct GetUserNonceArgs {
     address: String,
 }
 
-#[derive(Debug, Clone, ValueEnum, Copy)]
-enum NetworkArg {
-    Devnet,
-    Holesky,
-    HoleskyStage,
-    Mainnet,
-}
-
-impl From<NetworkArg> for Network {
-    fn from(env_arg: NetworkArg) -> Self {
-        match env_arg {
-            NetworkArg::Devnet => Network::Devnet,
-            NetworkArg::Holesky => Network::Holesky,
-            NetworkArg::HoleskyStage => Network::HoleskyStage,
-            NetworkArg::Mainnet => Network::Mainnet,
-        }
-    }
-}
-
 #[derive(Debug, Clone, ValueEnum)]
 pub enum ProvingSystemArg {
     #[clap(name = "GnarkPlonkBls12_381")]
@@ -366,7 +347,7 @@ async fn main() -> Result<(), AlignedError> {
 
             let aligned_verification_data_vec = submit_multiple(
                 &connect_addr,
-                submit_args.network.into(),
+                submit_args.network,
                 &verification_data_arr,
                 max_fee_wei,
                 wallet.clone(),
@@ -425,7 +406,7 @@ async fn main() -> Result<(), AlignedError> {
             info!("Verifying response data matches sent proof data...");
             let response = is_proof_verified(
                 &aligned_verification_data,
-                verify_inclusion_args.network.into(),
+                verify_inclusion_args.network,
                 &verify_inclusion_args.eth_rpc_url,
             )
             .await?;
@@ -490,9 +471,7 @@ async fn main() -> Result<(), AlignedError> {
 
             let client = SignerMiddleware::new(eth_rpc_provider.clone(), wallet.clone());
 
-            match deposit_to_aligned(amount_wei, client, deposit_to_batcher_args.network.into())
-                .await
-            {
+            match deposit_to_aligned(amount_wei, client, deposit_to_batcher_args.network).await {
                 Ok(receipt) => {
                     info!(
                         "Payment sent to the batcher successfully. Tx: 0x{:x}",
@@ -509,7 +488,7 @@ async fn main() -> Result<(), AlignedError> {
             match get_balance_in_aligned(
                 user_address,
                 &get_user_balance_args.eth_rpc_url,
-                get_user_balance_args.network.into(),
+                get_user_balance_args.network,
             )
             .await
             {