From 33eeb372e781241cfdb55f9b176866eed202be89 Mon Sep 17 00:00:00 2001
From: Bushstar <bushsolo@gmail.com>
Date: Wed, 5 Apr 2023 09:46:27 +0100
Subject: [PATCH 1/8] Return error on sub balance failure.

---
 src/rust/crates/ain-evm-ffi/src/lib.rs       | 15 +++++++++------
 src/rust/crates/ain-evm-state/src/handler.rs | 18 +++++++++++++++---
 2 files changed, 24 insertions(+), 9 deletions(-)

diff --git a/src/rust/crates/ain-evm-ffi/src/lib.rs b/src/rust/crates/ain-evm-ffi/src/lib.rs
index a171e53429a..f719076e83e 100644
--- a/src/rust/crates/ain-evm-ffi/src/lib.rs
+++ b/src/rust/crates/ain-evm-ffi/src/lib.rs
@@ -7,8 +7,8 @@ use std::error::Error;
 #[cxx::bridge]
 mod ffi {
     extern "Rust" {
-        fn evm_add_balance(address: &str, amount: [u8; 32]) -> Result<()>;
-        fn evm_sub_balance(address: &str, amount: [u8; 32]) -> Result<()>;
+        fn evm_add_balance(context: u64, address: &str, amount: [u8; 32]) -> Result<()>;
+        fn evm_sub_balance(context: u64, address: &str, amount: [u8; 32]) -> Result<bool>;
         fn evm_validate_raw_tx(tx: &str) -> Result<bool>;
 
         fn evm_get_context() -> u64;
@@ -33,12 +33,15 @@ mod ffi {
     }
 }
 
-pub fn evm_add_balance(address: &str, amount: [u8; 32]) -> Result<(), Box<dyn Error>> {
-    RUNTIME.evm.add_balance(address, amount.into())
+pub fn evm_add_balance(context: u64, address: &str, amount: [u8; 32]) -> Result<(), Box<dyn Error>> {
+    RUNTIME.evm.add_balance(context, address, amount.into())
 }
 
-pub fn evm_sub_balance(address: &str, amount: [u8; 32]) -> Result<(), Box<dyn Error>> {
-    RUNTIME.evm.sub_balance(address, amount.into())
+pub fn evm_sub_balance(context: u64, address: &str, amount: [u8; 32]) -> Result<bool, Box<dyn Error>> {
+    match RUNTIME.evm.sub_balance(context, address, amount.into()) {
+        Ok(_) => Ok(true),
+        Err(_) => Ok(false),
+    }
 }
 
 pub fn evm_validate_raw_tx(tx: &str) -> Result<bool, Box<dyn Error>> {
diff --git a/src/rust/crates/ain-evm-state/src/handler.rs b/src/rust/crates/ain-evm-state/src/handler.rs
index 1a4fbcc5e39..aca4bec24e2 100644
--- a/src/rust/crates/ain-evm-state/src/handler.rs
+++ b/src/rust/crates/ain-evm-state/src/handler.rs
@@ -13,6 +13,17 @@ use std::sync::{Arc, RwLock};
 use crate::tx_queue::TransactionQueueMap;
 use crate::EVMState;
 
+#[derive(Debug)]
+struct ErrorStr(String);
+
+impl std::fmt::Display for ErrorStr {
+    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+        write!(f, "Error: {}", self.0)
+    }
+}
+
+impl Error for ErrorStr {}
+
 #[derive(Clone, Debug)]
 pub struct EVMHandler {
     pub state: Arc<RwLock<EVMState>>,
@@ -47,7 +58,7 @@ impl EVMHandler {
     }
 
     // TODO wrap in EVM transaction and dryrun with evm_call
-    pub fn add_balance(&self, address: &str, value: U256) -> Result<(), Box<dyn Error>> {
+    pub fn add_balance(&self, context: u64, address: &str, value: U256) -> Result<(), Box<dyn Error>> {
         let to = address.parse()?;
         let mut state = self.state.write().unwrap();
         let mut account = state.entry(to).or_default();
@@ -55,14 +66,15 @@ impl EVMHandler {
         Ok(())
     }
 
-    pub fn sub_balance(&self, address: &str, value: U256) -> Result<(), Box<dyn Error>> {
+    pub fn sub_balance(&self, context: u64, address: &str, value: U256) -> Result<(), Box<dyn Error>> {
         let address = address.parse()?;
         let mut state = self.state.write().unwrap();
         let mut account = state.get_mut(&address).unwrap();
         if account.balance > value {
             account.balance = account.balance - value;
+            return Ok(());
         }
-        Ok(())
+        Err(Box::new(ErrorStr("Sub balance failed".into())))
     }
 
     pub fn validate_raw_tx(&self, tx: &str) -> Result<SignedTx, Box<dyn Error>> {

From 4971d726a2ebeb95e8df90b2aaeea1f898277f85 Mon Sep 17 00:00:00 2001
From: Bushstar <bushsolo@gmail.com>
Date: Wed, 5 Apr 2023 09:55:24 +0100
Subject: [PATCH 2/8] Change to expected return type

---
 src/wallet/wallet.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/wallet/wallet.cpp b/src/wallet/wallet.cpp
index 67a80432a75..85afc6cbe60 100644
--- a/src/wallet/wallet.cpp
+++ b/src/wallet/wallet.cpp
@@ -167,7 +167,7 @@ std::array<uint8_t, 32> GetKeyFromWallets(rust::Vec<uint8_t> input) {
         }
     }
 
-    std::array<uint32_t, 32> result{};
+    std::array<uint8_t, 32> result{};
     std::copy(key.begin(), key.end(), result.begin());
 
     return result;

From 518ab4dae34f04131879223cb1597b1131462768 Mon Sep 17 00:00:00 2001
From: Bushstar <bushsolo@gmail.com>
Date: Wed, 5 Apr 2023 09:55:46 +0100
Subject: [PATCH 3/8] Update EVM in/out TX

---
 src/amount.h                  |  2 ++
 src/masternodes/mn_checks.cpp | 50 ++++++++++++++++++++++++++++-------
 src/masternodes/rpc_evm.cpp   |  3 ---
 3 files changed, 42 insertions(+), 13 deletions(-)

diff --git a/src/amount.h b/src/amount.h
index dd26bdc686b..67db644848f 100644
--- a/src/amount.h
+++ b/src/amount.h
@@ -77,6 +77,8 @@ struct DCT_ID {
 
 static constexpr CAmount COIN = 100000000;
 static constexpr CAmount CENT = 1000000;
+static constexpr int64_t WEI_IN_GWEI    = 1000000000;
+static constexpr int64_t CAMOUNT_TO_WEI = 10;
 
 //Converts the given value to decimal format string with COIN precision.
 inline std::string GetDecimalString(CAmount nValue)
diff --git a/src/masternodes/mn_checks.cpp b/src/masternodes/mn_checks.cpp
index 7afd8c7faaa..8f79d2a4aa7 100644
--- a/src/masternodes/mn_checks.cpp
+++ b/src/masternodes/mn_checks.cpp
@@ -3834,12 +3834,6 @@ class CCustomTxApplyVisitor : public CCustomTxVisitor {
         const auto sumFrom = SumAllTransfers(obj.from);
         const auto sumTo   = SumAllTransfers(obj.to);
 
-        if (obj.type == CTransferBalanceType::EvmIn || obj.type == CTransferBalanceType::EvmOut) {
-            for (const auto& [id, _] : sumFrom.balances)
-                if (id != DCT_ID{0})
-                    return Res::Err("For EVM in/out transfers only DFI token is currently supported");
-        }
-
         if (sumFrom != sumTo)
             return Res::Err("sum of inputs (from) != sum of outputs (to)");
 
@@ -3851,18 +3845,54 @@ class CCustomTxApplyVisitor : public CCustomTxVisitor {
             if (!res)
                 return res;
         } else if (obj.type == CTransferBalanceType::EvmIn) {
-            for (const auto& [addr, _] : obj.to) {
+            res = SubBalancesDelShares(obj.from);
+            if (!res)
+                return res;
+
+            for (const auto& [addr, balances] : obj.to) {
                 CTxDestination dest;
                 if (ExtractDestination(addr, dest)) {
                     if (dest.index() != WitV16KeyEthHashType) {
                         return Res::Err("To address must be an ETH address in case of \"evmin\" transfertype");
                     }
                 }
+
+                const auto toAddress = std::get<WitnessV16EthHash>(dest);
+
+                for (const auto& [id, amount] : balances.balances) {
+                    if (id != DCT_ID{0}) {
+                        return Res::Err("For EVM out transfers, only DFI token is currently supported");
+                    }
+
+                    arith_uint256 balanceIn = amount;
+                    balanceIn *= CAMOUNT_TO_WEI * WEI_IN_GWEI;
+                    evm_add_balance(evmContext, HexStr(toAddress.begin(), toAddress.end()), ArithToUint256(balanceIn).ToArrayReversed());
+                }
             }
-            res = SubBalancesDelShares(obj.from);
-            if (!res)
-                return res;
         } else if (obj.type == CTransferBalanceType::EvmOut) {
+            for (const auto& [addr, balances] : obj.from) {
+                CTxDestination dest;
+                if (ExtractDestination(addr, dest)) {
+                    if (dest.index() != WitV16KeyEthHashType) {
+                        return Res::Err("Invalid destination");
+                    }
+                }
+
+                const auto fromAddress = std::get<WitnessV16EthHash>(dest);
+
+                for (const auto& [id, amount] : balances.balances) {
+                    if (id != DCT_ID{0}) {
+                        return Res::Err("For EVM out transfers, only DFI token is currently supported");
+                    }
+
+                    arith_uint256 balanceIn = amount;
+                    balanceIn *= CAMOUNT_TO_WEI * WEI_IN_GWEI;
+                    if (!evm_sub_balance(evmContext, HexStr(fromAddress.begin(), fromAddress.end()), ArithToUint256(balanceIn).ToArrayReversed())) {
+                        return Res::Err("Not enough balance in %s to cover EVM out", EncodeDestination(dest));
+                    }
+                }
+            }
+
             res = AddBalancesSetShares(obj.to);
             if (!res)
                 return res;
diff --git a/src/masternodes/rpc_evm.cpp b/src/masternodes/rpc_evm.cpp
index 1bdd020873a..fed5bc3c220 100644
--- a/src/masternodes/rpc_evm.cpp
+++ b/src/masternodes/rpc_evm.cpp
@@ -3,9 +3,6 @@
 #include <libain_evm.h>
 #include <key_io.h>
 
-const int64_t WEI_IN_GWEI    = 1000000000;
-const int64_t CAMOUNT_TO_WEI = 10;
-
 UniValue evmtx(const JSONRPCRequest& request) {
     auto pwallet = GetWallet(request);
 

From 51ee3dc80d6d6288f92b68296674d0d67253fc53 Mon Sep 17 00:00:00 2001
From: Bushstar <bushsolo@gmail.com>
Date: Wed, 5 Apr 2023 10:07:06 +0100
Subject: [PATCH 4/8] Use array instead of rust::Vec

---
 src/wallet/wallet.cpp | 2 +-
 src/wallet/wallet.h   | 3 +--
 2 files changed, 2 insertions(+), 3 deletions(-)

diff --git a/src/wallet/wallet.cpp b/src/wallet/wallet.cpp
index 85afc6cbe60..1bb27f4d52d 100644
--- a/src/wallet/wallet.cpp
+++ b/src/wallet/wallet.cpp
@@ -156,7 +156,7 @@ std::shared_ptr<CWallet> LoadWallet(interfaces::Chain& chain, const WalletLocati
     return wallet;
 }
 
-std::array<uint8_t, 32> GetKeyFromWallets(rust::Vec<uint8_t> input) {
+std::array<uint8_t, 32> GetKeyFromWallets(std::array<uint8_t, 20> input) {
     CKey key;
     CKeyID keyID;
     std::copy(input.begin(), input.end(), keyID.begin());
diff --git a/src/wallet/wallet.h b/src/wallet/wallet.h
index d976bded457..4a297c7689a 100644
--- a/src/wallet/wallet.h
+++ b/src/wallet/wallet.h
@@ -9,7 +9,6 @@
 #include <amount.h>
 #include <interfaces/chain.h>
 #include <interfaces/handler.h>
-#include <libain_evm.h>
 #include <outputtype.h>
 #include <policy/feerate.h>
 #include <script/sign.h>
@@ -59,7 +58,7 @@ bool HasWallets();
 std::vector<std::shared_ptr<CWallet>> GetWallets();
 std::shared_ptr<CWallet> GetWallet(const std::string& name);
 std::shared_ptr<CWallet> LoadWallet(interfaces::Chain& chain, const WalletLocation& location, std::string& error, std::string& warning);
-std::array<uint8_t, 32> GetKeyFromWallets(rust::Vec<uint8_t> input);
+std::array<uint8_t, 32> GetKeyFromWallets(std::array<uint8_t, 20> input);
 
 CKey GetWalletsKey(CKeyID const & keyid);
 

From d0c35e789b2791dc9efa858d91b43a2eaef897be Mon Sep 17 00:00:00 2001
From: Bushstar <bushsolo@gmail.com>
Date: Wed, 5 Apr 2023 13:33:55 +0100
Subject: [PATCH 5/8] Pass context to CustomTxVisit

---
 src/masternodes/mn_checks.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/masternodes/mn_checks.cpp b/src/masternodes/mn_checks.cpp
index 8f79d2a4aa7..960f04fb705 100644
--- a/src/masternodes/mn_checks.cpp
+++ b/src/masternodes/mn_checks.cpp
@@ -4115,7 +4115,7 @@ Res ApplyCustomTx(CCustomCSView &mnview,
             PopulateVaultHistoryData(mnview.GetHistoryWriters(), view, txMessage, txType, height, txn, tx.GetHash());
         }
 
-        res = CustomTxVisit(view, coins, tx, height, consensus, txMessage, time, txn);
+        res = CustomTxVisit(view, coins, tx, height, consensus, txMessage, time, txn, evmContext);
 
         if (res) {
             if (canSpend && txType == CustomTxType::UpdateMasternode) {

From 943dbccefe15128da99b06fefde1e179a6d80f3c Mon Sep 17 00:00:00 2001
From: jouzo <jdesclercs@gmail.com>
Date: Wed, 5 Apr 2023 13:59:15 +0100
Subject: [PATCH 6/8] Add temporary state to TransactionQueue

---
 src/rust/crates/ain-evm-ffi/src/lib.rs        | 16 +++-
 src/rust/crates/ain-evm-state/src/handler.rs  | 81 ++++++++-----------
 src/rust/crates/ain-evm-state/src/tx_queue.rs | 79 +++++++++++++++++-
 3 files changed, 124 insertions(+), 52 deletions(-)

diff --git a/src/rust/crates/ain-evm-ffi/src/lib.rs b/src/rust/crates/ain-evm-ffi/src/lib.rs
index f719076e83e..9f4b7db452f 100644
--- a/src/rust/crates/ain-evm-ffi/src/lib.rs
+++ b/src/rust/crates/ain-evm-ffi/src/lib.rs
@@ -33,11 +33,21 @@ mod ffi {
     }
 }
 
-pub fn evm_add_balance(context: u64, address: &str, amount: [u8; 32]) -> Result<(), Box<dyn Error>> {
-    RUNTIME.evm.add_balance(context, address, amount.into())
+pub fn evm_add_balance(
+    context: u64,
+    address: &str,
+    amount: [u8; 32],
+) -> Result<(), Box<dyn Error>> {
+    let address = address.parse()?;
+    Ok(RUNTIME.evm.add_balance(context, address, amount.into()))
 }
 
-pub fn evm_sub_balance(context: u64, address: &str, amount: [u8; 32]) -> Result<bool, Box<dyn Error>> {
+pub fn evm_sub_balance(
+    context: u64,
+    address: &str,
+    amount: [u8; 32],
+) -> Result<bool, Box<dyn Error>> {
+    let address = address.parse()?;
     match RUNTIME.evm.sub_balance(context, address, amount.into()) {
         Ok(_) => Ok(true),
         Err(_) => Ok(false),
diff --git a/src/rust/crates/ain-evm-state/src/handler.rs b/src/rust/crates/ain-evm-state/src/handler.rs
index aca4bec24e2..d351ef4ebfe 100644
--- a/src/rust/crates/ain-evm-state/src/handler.rs
+++ b/src/rust/crates/ain-evm-state/src/handler.rs
@@ -13,17 +13,6 @@ use std::sync::{Arc, RwLock};
 use crate::tx_queue::TransactionQueueMap;
 use crate::EVMState;
 
-#[derive(Debug)]
-struct ErrorStr(String);
-
-impl std::fmt::Display for ErrorStr {
-    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
-        write!(f, "Error: {}", self.0)
-    }
-}
-
-impl Error for ErrorStr {}
-
 #[derive(Clone, Debug)]
 pub struct EVMHandler {
     pub state: Arc<RwLock<EVMState>>,
@@ -58,23 +47,17 @@ impl EVMHandler {
     }
 
     // TODO wrap in EVM transaction and dryrun with evm_call
-    pub fn add_balance(&self, context: u64, address: &str, value: U256) -> Result<(), Box<dyn Error>> {
-        let to = address.parse()?;
-        let mut state = self.state.write().unwrap();
-        let mut account = state.entry(to).or_default();
-        account.balance = account.balance + value;
-        Ok(())
+    pub fn add_balance(&self, context: u64, address: H160, value: U256) {
+        self.tx_queues.add_balance(context, address, value)
     }
 
-    pub fn sub_balance(&self, context: u64, address: &str, value: U256) -> Result<(), Box<dyn Error>> {
-        let address = address.parse()?;
-        let mut state = self.state.write().unwrap();
-        let mut account = state.get_mut(&address).unwrap();
-        if account.balance > value {
-            account.balance = account.balance - value;
-            return Ok(());
-        }
-        Err(Box::new(ErrorStr("Sub balance failed".into())))
+    pub fn sub_balance(
+        &self,
+        context: u64,
+        address: H160,
+        value: U256,
+    ) -> Result<(), Box<dyn Error>> {
+        self.tx_queues.sub_balance(context, address, value)
     }
 
     pub fn validate_raw_tx(&self, tx: &str) -> Result<SignedTx, Box<dyn Error>> {
@@ -109,7 +92,8 @@ impl EVMHandler {
     }
 
     pub fn get_context(&self) -> u64 {
-        self.tx_queues.get_context()
+        let state = self.state.read().unwrap().clone();
+        self.tx_queues.get_context(state)
     }
 
     pub fn discard_context(&self, context: u64) {
@@ -130,7 +114,7 @@ impl EVMHandler {
         let mut tx_hashes = Vec::with_capacity(self.tx_queues.len(context));
         let mut failed_tx_hashes = Vec::with_capacity(self.tx_queues.len(context));
         let vicinity = get_vicinity(None, None);
-        let state = self.state.read().unwrap().clone();
+        let state = self.tx_queues.state(context).expect("Wrong context");
         let backend = MemoryBackend::new(&vicinity, state);
         let mut executor = AinExecutor::new(backend);
 
@@ -146,6 +130,8 @@ impl EVMHandler {
             }
         }
 
+        self.tx_queues.remove(context);
+
         if update_state {
             let mut state = self.state.write().unwrap();
             *state = executor.backend().state().clone();
@@ -200,23 +186,25 @@ mod tests {
     #[test]
     fn test_finalize_block_and_update_test() {
         let handler = EVMHandler::new();
-        handler
-            .add_balance(
-                "0x6745f998a96050bb9b0449e6bd4358138a519679",
-                U256::from_str_radix("100000000000000000000", 10).unwrap(),
-            )
-            .unwrap();
         let context = handler.get_context();
+        handler.add_balance(
+            context,
+            "0x6745f998a96050bb9b0449e6bd4358138a519679"
+                .parse()
+                .unwrap(),
+            U256::from_str_radix("100000000000000000000", 10).unwrap(),
+        );
 
         let tx1: SignedTx = "f86b02830186a0830186a094a8f7c4c78c36e54c3950ad58dad24ca5e0191b2989056bc75e2d631000008025a0b0842b0c78dd7fc33584ec9a81ab5104fe70169878de188ba6c11fe7605e298aa0735dc483f625f17d68d1e1fae779b7160612628e6dde9eecf087892fe60bba4e".try_into().unwrap();
         handler.tx_queues.add_signed_tx(context, tx1.clone());
 
-        handler
-            .add_balance(
-                "0xc0cd829081485e70348975d325fe5275140277bd",
-                U256::from_str_radix("100000000000000000000", 10).unwrap(),
-            )
-            .unwrap();
+        handler.add_balance(
+            context,
+            "0xc0cd829081485e70348975d325fe5275140277bd"
+                .parse()
+                .unwrap(),
+            U256::from_str_radix("100000000000000000000", 10).unwrap(),
+        );
         let tx2: SignedTx = "f86b02830186a0830186a094a8f7c4c78c36e54c3950ad58dad24ca5e0191b2989056bc75e2d631000008025a01465e2d999c34b22bf4b8b5c9439918e46341f4f0da1b00a6b0479c541161d4aa074abe79c51bf57086e1e84b57ee483cbb2ecf30e8222bc0472436fabfc57dda8".try_into().unwrap();
         handler.tx_queues.add_signed_tx(context, tx2.clone());
 
@@ -281,13 +269,14 @@ mod tests {
     #[test]
     fn test_finalize_block_and_do_not_update_test() {
         let handler = EVMHandler::new();
-        handler
-            .add_balance(
-                "0x4a1080c5533cb89edc4b65013f08f78868e382de",
-                U256::from_str_radix("100000000000000000000", 10).unwrap(),
-            )
-            .unwrap();
         let context = handler.get_context();
+        handler.add_balance(
+            context,
+            "0x4a1080c5533cb89edc4b65013f08f78868e382de"
+                .parse()
+                .unwrap(),
+            U256::from_str_radix("100000000000000000000", 10).unwrap(),
+        );
 
         let tx1: SignedTx = "f86b02830186a0830186a094a8f7c4c78c36e54c3950ad58dad24ca5e0191b2989056bc75e2d631000008025a0b0842b0c78dd7fc33584ec9a81ab5104fe70169878de188ba6c11fe7605e298aa0735dc483f625f17d68d1e1fae779b7160612628e6dde9eecf087892fe60bba4e".try_into().unwrap();
         handler.tx_queues.add_signed_tx(context, tx1.clone());
diff --git a/src/rust/crates/ain-evm-state/src/tx_queue.rs b/src/rust/crates/ain-evm-state/src/tx_queue.rs
index 5d1c262c568..fba5af16e79 100644
--- a/src/rust/crates/ain-evm-state/src/tx_queue.rs
+++ b/src/rust/crates/ain-evm-state/src/tx_queue.rs
@@ -1,3 +1,4 @@
+use primitive_types::{H160, U256};
 use rand::Rng;
 use std::{
     collections::HashMap,
@@ -6,6 +7,8 @@ use std::{
 
 use ain_evm::transaction::SignedTx;
 
+use crate::EVMState;
+
 #[derive(Debug)]
 pub struct TransactionQueueMap {
     queues: RwLock<HashMap<u64, TransactionQueue>>,
@@ -18,19 +21,23 @@ impl TransactionQueueMap {
         }
     }
 
-    pub fn get_context(&self) -> u64 {
+    pub fn get_context(&self, state: EVMState) -> u64 {
         let mut rng = rand::thread_rng();
         loop {
             let context = rng.gen();
             let mut write_guard = self.queues.write().unwrap();
 
             if !write_guard.contains_key(&context) {
-                write_guard.insert(context, TransactionQueue::new());
+                write_guard.insert(context, TransactionQueue::new(state));
                 return context;
             }
         }
     }
 
+    pub fn remove(&self, index: u64) -> Option<TransactionQueue> {
+        self.queues.write().unwrap().remove(&index)
+    }
+
     pub fn clear(&self, index: u64) {
         if let Some(queue) = self.queues.read().unwrap().get(&index) {
             queue.clear()
@@ -56,17 +63,48 @@ impl TransactionQueueMap {
             None => 0,
         }
     }
+
+    pub fn add_balance(&self, index: u64, address: H160, value: U256) {
+        if let Some(queue) = self.queues.read().unwrap().get(&index) {
+            queue.add_balance(address, value)
+        }
+    }
+
+    pub fn sub_balance(
+        &self,
+        index: u64,
+        address: H160,
+        value: U256,
+    ) -> Result<(), Box<dyn std::error::Error>> {
+        if let Some(queue) = self.queues.read().unwrap().get(&index) {
+            queue.sub_balance(address, value)
+        } else {
+            Err(Box::new(ErrorStr(
+                "No transaction queue for this context".into(),
+            )))
+        }
+    }
+
+    pub fn state(&self, index: u64) -> Option<EVMState> {
+        if let Some(queue) = self.queues.read().unwrap().get(&index) {
+            Some(queue.state())
+        } else {
+            None
+        }
+    }
 }
 
 #[derive(Debug)]
 pub struct TransactionQueue {
     txs: Mutex<Vec<SignedTx>>,
+    state: RwLock<EVMState>,
 }
 
 impl TransactionQueue {
-    pub fn new() -> Self {
+    pub fn new(state: EVMState) -> Self {
         Self {
             txs: Mutex::new(Vec::new()),
+            state: RwLock::new(state),
         }
     }
 
@@ -89,4 +127,39 @@ impl TransactionQueue {
     pub fn len(&self) -> usize {
         self.txs.lock().unwrap().len()
     }
+
+    pub fn state(&self) -> EVMState {
+        self.state.read().unwrap().clone()
+    }
+
+    pub fn add_balance(&self, address: H160, value: U256) {
+        let mut state = self.state.write().unwrap();
+        let mut account = state.entry(address).or_default();
+        account.balance = account.balance + value;
+    }
+
+    pub fn sub_balance(
+        &self,
+        address: H160,
+        value: U256,
+    ) -> Result<(), Box<dyn std::error::Error>> {
+        let mut state = self.state.write().unwrap();
+        let mut account = state.get_mut(&address).unwrap();
+        if account.balance > value {
+            account.balance = account.balance - value;
+            return Ok(());
+        }
+        Err(Box::new(ErrorStr("Sub balance failed".into())))
+    }
 }
+
+#[derive(Debug)]
+struct ErrorStr(String);
+
+impl std::fmt::Display for ErrorStr {
+    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+        write!(f, "Error: {}", self.0)
+    }
+}
+
+impl std::error::Error for ErrorStr {}

From c42d0e3bf24107fed54a69f8463b4e71203140af Mon Sep 17 00:00:00 2001
From: jouzo <jdesclercs@gmail.com>
Date: Wed, 5 Apr 2023 14:09:07 +0100
Subject: [PATCH 7/8] Improve readability

---
 src/rust/crates/ain-evm-state/src/handler.rs  |  4 +-
 src/rust/crates/ain-evm-state/src/tx_queue.rs | 75 ++++++++++---------
 2 files changed, 41 insertions(+), 38 deletions(-)

diff --git a/src/rust/crates/ain-evm-state/src/handler.rs b/src/rust/crates/ain-evm-state/src/handler.rs
index d351ef4ebfe..2ee1b390091 100644
--- a/src/rust/crates/ain-evm-state/src/handler.rs
+++ b/src/rust/crates/ain-evm-state/src/handler.rs
@@ -57,7 +57,9 @@ impl EVMHandler {
         address: H160,
         value: U256,
     ) -> Result<(), Box<dyn Error>> {
-        self.tx_queues.sub_balance(context, address, value)
+        self.tx_queues
+            .sub_balance(context, address, value)
+            .map_err(|e| e.into())
     }
 
     pub fn validate_raw_tx(&self, tx: &str) -> Result<SignedTx, Box<dyn Error>> {
diff --git a/src/rust/crates/ain-evm-state/src/tx_queue.rs b/src/rust/crates/ain-evm-state/src/tx_queue.rs
index fba5af16e79..24a011ba545 100644
--- a/src/rust/crates/ain-evm-state/src/tx_queue.rs
+++ b/src/rust/crates/ain-evm-state/src/tx_queue.rs
@@ -34,59 +34,57 @@ impl TransactionQueueMap {
         }
     }
 
-    pub fn remove(&self, index: u64) -> Option<TransactionQueue> {
-        self.queues.write().unwrap().remove(&index)
+    pub fn remove(&self, context_id: u64) -> Option<TransactionQueue> {
+        self.queues.write().unwrap().remove(&context_id)
     }
 
-    pub fn clear(&self, index: u64) {
-        if let Some(queue) = self.queues.read().unwrap().get(&index) {
+    pub fn clear(&self, context_id: u64) {
+        if let Some(queue) = self.queues.read().unwrap().get(&context_id) {
             queue.clear()
         }
     }
 
-    pub fn add_signed_tx(&self, index: u64, signed_tx: SignedTx) {
-        if let Some(queue) = self.queues.read().unwrap().get(&index) {
+    pub fn add_signed_tx(&self, context_id: u64, signed_tx: SignedTx) {
+        if let Some(queue) = self.queues.read().unwrap().get(&context_id) {
             queue.add_signed_tx(signed_tx)
         }
     }
 
-    pub fn drain_all(&self, index: u64) -> Vec<SignedTx> {
-        match self.queues.read().unwrap().get(&index) {
+    pub fn drain_all(&self, context_id: u64) -> Vec<SignedTx> {
+        match self.queues.read().unwrap().get(&context_id) {
             Some(queue) => queue.drain_all(),
             None => Vec::new(),
         }
     }
 
-    pub fn len(&self, index: u64) -> usize {
-        match self.queues.read().unwrap().get(&index) {
+    pub fn len(&self, context_id: u64) -> usize {
+        match self.queues.read().unwrap().get(&context_id) {
             Some(queue) => queue.len(),
             None => 0,
         }
     }
 
-    pub fn add_balance(&self, index: u64, address: H160, value: U256) {
-        if let Some(queue) = self.queues.read().unwrap().get(&index) {
+    pub fn add_balance(&self, context_id: u64, address: H160, value: U256) {
+        if let Some(queue) = self.queues.read().unwrap().get(&context_id) {
             queue.add_balance(address, value)
         }
     }
 
     pub fn sub_balance(
         &self,
-        index: u64,
+        context_id: u64,
         address: H160,
         value: U256,
-    ) -> Result<(), Box<dyn std::error::Error>> {
-        if let Some(queue) = self.queues.read().unwrap().get(&index) {
+    ) -> Result<(), QueueError> {
+        if let Some(queue) = self.queues.read().unwrap().get(&context_id) {
             queue.sub_balance(address, value)
         } else {
-            Err(Box::new(ErrorStr(
-                "No transaction queue for this context".into(),
-            )))
+            Err(QueueError::NoSuchContext)
         }
     }
 
-    pub fn state(&self, index: u64) -> Option<EVMState> {
-        if let Some(queue) = self.queues.read().unwrap().get(&index) {
+    pub fn state(&self, context_id: u64) -> Option<EVMState> {
+        if let Some(queue) = self.queues.read().unwrap().get(&context_id) {
             Some(queue.state())
         } else {
             None
@@ -96,24 +94,24 @@ impl TransactionQueueMap {
 
 #[derive(Debug)]
 pub struct TransactionQueue {
-    txs: Mutex<Vec<SignedTx>>,
+    transactions: Mutex<Vec<SignedTx>>,
     state: RwLock<EVMState>,
 }
 
 impl TransactionQueue {
     pub fn new(state: EVMState) -> Self {
         Self {
-            txs: Mutex::new(Vec::new()),
+            transactions: Mutex::new(Vec::new()),
             state: RwLock::new(state),
         }
     }
 
     pub fn clear(&self) {
-        self.txs.lock().unwrap().clear()
+        self.transactions.lock().unwrap().clear()
     }
 
     pub fn drain_all(&self) -> Vec<SignedTx> {
-        self.txs
+        self.transactions
             .lock()
             .unwrap()
             .drain(..)
@@ -121,11 +119,11 @@ impl TransactionQueue {
     }
 
     pub fn add_signed_tx(&self, signed_tx: SignedTx) {
-        self.txs.lock().unwrap().push(signed_tx)
+        self.transactions.lock().unwrap().push(signed_tx)
     }
 
     pub fn len(&self) -> usize {
-        self.txs.lock().unwrap().len()
+        self.transactions.lock().unwrap().len()
     }
 
     pub fn state(&self) -> EVMState {
@@ -138,28 +136,31 @@ impl TransactionQueue {
         account.balance = account.balance + value;
     }
 
-    pub fn sub_balance(
-        &self,
-        address: H160,
-        value: U256,
-    ) -> Result<(), Box<dyn std::error::Error>> {
+    pub fn sub_balance(&self, address: H160, value: U256) -> Result<(), QueueError> {
         let mut state = self.state.write().unwrap();
         let mut account = state.get_mut(&address).unwrap();
         if account.balance > value {
             account.balance = account.balance - value;
-            return Ok(());
+            Ok(())
+        } else {
+            Err(QueueError::InsufficientBalance)
         }
-        Err(Box::new(ErrorStr("Sub balance failed".into())))
     }
 }
 
 #[derive(Debug)]
-struct ErrorStr(String);
+pub enum QueueError {
+    NoSuchContext,
+    InsufficientBalance,
+}
 
-impl std::fmt::Display for ErrorStr {
+impl std::fmt::Display for QueueError {
     fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
-        write!(f, "Error: {}", self.0)
+        match self {
+            QueueError::NoSuchContext => write!(f, "No transaction queue for this context"),
+            QueueError::InsufficientBalance => write!(f, "Insufficient balance"),
+        }
     }
 }
 
-impl std::error::Error for ErrorStr {}
+impl std::error::Error for QueueError {}

From a36861e5e0eaf7d7f04b3d746cbc5b39a5e56371 Mon Sep 17 00:00:00 2001
From: jouzo <jdesclercs@gmail.com>
Date: Wed, 5 Apr 2023 14:26:36 +0100
Subject: [PATCH 8/8] Block integration test

---
 src/rust/crates/ain-evm-state/src/handler.rs | 114 -------------------
 src/rust/crates/ain-evm-state/tests/block.rs | 110 ++++++++++++++++++
 2 files changed, 110 insertions(+), 114 deletions(-)
 create mode 100644 src/rust/crates/ain-evm-state/tests/block.rs

diff --git a/src/rust/crates/ain-evm-state/src/handler.rs b/src/rust/crates/ain-evm-state/src/handler.rs
index 2ee1b390091..97433ad1306 100644
--- a/src/rust/crates/ain-evm-state/src/handler.rs
+++ b/src/rust/crates/ain-evm-state/src/handler.rs
@@ -123,7 +123,6 @@ impl EVMHandler {
         for signed_tx in self.tx_queues.drain_all(context) {
             let tx_response = executor.exec(&signed_tx);
             println!("tx_response : {:#?}", tx_response);
-
             if tx_response.exit_reason.is_succeed() {
                 // responses.push()
                 tx_hashes.push(signed_tx.transaction);
@@ -177,116 +176,3 @@ fn get_vicinity(origin: Option<H160>, gas_price: Option<U256>) -> MemoryVicinity
         block_base_fee_per_gas: U256::MAX,
     }
 }
-
-mod tests {
-    use ain_evm::transaction::SignedTx;
-    use evm::backend::MemoryAccount;
-    use primitive_types::{H160, H256, U256};
-
-    use super::EVMHandler;
-    use rlp::Encodable;
-    #[test]
-    fn test_finalize_block_and_update_test() {
-        let handler = EVMHandler::new();
-        let context = handler.get_context();
-        handler.add_balance(
-            context,
-            "0x6745f998a96050bb9b0449e6bd4358138a519679"
-                .parse()
-                .unwrap(),
-            U256::from_str_radix("100000000000000000000", 10).unwrap(),
-        );
-
-        let tx1: SignedTx = "f86b02830186a0830186a094a8f7c4c78c36e54c3950ad58dad24ca5e0191b2989056bc75e2d631000008025a0b0842b0c78dd7fc33584ec9a81ab5104fe70169878de188ba6c11fe7605e298aa0735dc483f625f17d68d1e1fae779b7160612628e6dde9eecf087892fe60bba4e".try_into().unwrap();
-        handler.tx_queues.add_signed_tx(context, tx1.clone());
-
-        handler.add_balance(
-            context,
-            "0xc0cd829081485e70348975d325fe5275140277bd"
-                .parse()
-                .unwrap(),
-            U256::from_str_radix("100000000000000000000", 10).unwrap(),
-        );
-        let tx2: SignedTx = "f86b02830186a0830186a094a8f7c4c78c36e54c3950ad58dad24ca5e0191b2989056bc75e2d631000008025a01465e2d999c34b22bf4b8b5c9439918e46341f4f0da1b00a6b0479c541161d4aa074abe79c51bf57086e1e84b57ee483cbb2ecf30e8222bc0472436fabfc57dda8".try_into().unwrap();
-        handler.tx_queues.add_signed_tx(context, tx2.clone());
-
-        let tx3: SignedTx = "f86b02830186a0830186a094a8f7c4c78c36e54c3950ad58dad24ca5e0191b2989056bc75e2d631000008025a070b21a24cec13c0569099ee2f8221268103fd609646b73f7c9e85efeb7af5c8ea03d5de75bc12ce28a80f7c0401df6021cc82a334cb1c802c8b9d46223c5c8eb40".try_into().unwrap();
-        handler.tx_queues.add_signed_tx(context, tx3.clone());
-
-        assert_eq!(handler.tx_queues.len(context), 3);
-        assert_eq!(handler.tx_queues.len(handler.get_context()), 0);
-
-        let (block, failed_txs) = handler.finalize_block(context, true).unwrap();
-        assert_eq!(
-            block.transactions,
-            vec![tx1, tx2]
-                .into_iter()
-                .map(|t| t.transaction)
-                .collect::<Vec<_>>()
-        );
-        assert_eq!(
-            failed_txs,
-            vec![tx3]
-                .into_iter()
-                .map(|t| t.transaction)
-                .collect::<Vec<_>>()
-        );
-
-        let state = handler.state.read().unwrap();
-        assert_eq!(
-            state
-                .get(
-                    &"0xa8f7c4c78c36e54c3950ad58dad24ca5e0191b29"
-                        .parse()
-                        .unwrap()
-                )
-                .unwrap()
-                .balance,
-            U256::from_str_radix("200000000000000000000", 10).unwrap()
-        );
-        assert_eq!(
-            state
-                .get(
-                    &"0x6745f998a96050bb9b0449e6bd4358138a519679"
-                        .parse()
-                        .unwrap()
-                )
-                .unwrap()
-                .balance,
-            U256::from_str_radix("0", 10).unwrap()
-        );
-        assert_eq!(
-            state
-                .get(
-                    &"0xc0cd829081485e70348975d325fe5275140277bd"
-                        .parse()
-                        .unwrap()
-                )
-                .unwrap()
-                .balance,
-            U256::from_str_radix("0", 10).unwrap()
-        );
-    }
-
-    #[test]
-    fn test_finalize_block_and_do_not_update_test() {
-        let handler = EVMHandler::new();
-        let context = handler.get_context();
-        handler.add_balance(
-            context,
-            "0x4a1080c5533cb89edc4b65013f08f78868e382de"
-                .parse()
-                .unwrap(),
-            U256::from_str_radix("100000000000000000000", 10).unwrap(),
-        );
-
-        let tx1: SignedTx = "f86b02830186a0830186a094a8f7c4c78c36e54c3950ad58dad24ca5e0191b2989056bc75e2d631000008025a0b0842b0c78dd7fc33584ec9a81ab5104fe70169878de188ba6c11fe7605e298aa0735dc483f625f17d68d1e1fae779b7160612628e6dde9eecf087892fe60bba4e".try_into().unwrap();
-        handler.tx_queues.add_signed_tx(context, tx1.clone());
-
-        let old_state = handler.state.read().unwrap();
-        let _ = handler.finalize_block(context, false).unwrap();
-
-        let new_state = handler.state.read().unwrap();
-        assert_eq!(*new_state, *old_state);
-    }
-}
diff --git a/src/rust/crates/ain-evm-state/tests/block.rs b/src/rust/crates/ain-evm-state/tests/block.rs
new file mode 100644
index 00000000000..e3a96a9e4c7
--- /dev/null
+++ b/src/rust/crates/ain-evm-state/tests/block.rs
@@ -0,0 +1,110 @@
+use ain_evm::transaction::SignedTx;
+use primitive_types::U256;
+
+use ain_evm_state::handler::EVMHandler;
+
+#[test]
+fn test_finalize_block_and_do_not_update_state() {
+    let handler = EVMHandler::new();
+    let context = handler.get_context();
+    handler.add_balance(
+        context,
+        "0x4a1080c5533cb89edc4b65013f08f78868e382de"
+            .parse()
+            .unwrap(),
+        U256::from_str_radix("100000000000000000000", 10).unwrap(),
+    );
+
+    let tx1: SignedTx = "f86b02830186a0830186a094a8f7c4c78c36e54c3950ad58dad24ca5e0191b2989056bc75e2d631000008025a0b0842b0c78dd7fc33584ec9a81ab5104fe70169878de188ba6c11fe7605e298aa0735dc483f625f17d68d1e1fae779b7160612628e6dde9eecf087892fe60bba4e".try_into().unwrap();
+    println!("tx1 : {:#?}", tx1);
+    handler.tx_queues.add_signed_tx(context, tx1.clone());
+
+    let old_state = handler.state.read().unwrap();
+    let _ = handler.finalize_block(context, false).unwrap();
+
+    let new_state = handler.state.read().unwrap();
+    assert_eq!(*new_state, *old_state);
+}
+
+#[test]
+fn test_finalize_block_and_update_state() {
+    let handler = EVMHandler::new();
+    let context = handler.get_context();
+    handler.add_balance(
+        context,
+        "0x6745f998a96050bb9b0449e6bd4358138a519679"
+            .parse()
+            .unwrap(),
+        U256::from_str_radix("100000000000000000000", 10).unwrap(),
+    );
+
+    let tx1: SignedTx = "f86b02830186a0830186a094a8f7c4c78c36e54c3950ad58dad24ca5e0191b2989056bc75e2d631000008025a0b0842b0c78dd7fc33584ec9a81ab5104fe70169878de188ba6c11fe7605e298aa0735dc483f625f17d68d1e1fae779b7160612628e6dde9eecf087892fe60bba4e".try_into().unwrap();
+    handler.tx_queues.add_signed_tx(context, tx1.clone());
+
+    handler.add_balance(
+        context,
+        "0xc0cd829081485e70348975d325fe5275140277bd"
+            .parse()
+            .unwrap(),
+        U256::from_str_radix("100000000000000000000", 10).unwrap(),
+    );
+    let tx2: SignedTx = "f86b02830186a0830186a094a8f7c4c78c36e54c3950ad58dad24ca5e0191b2989056bc75e2d631000008025a01465e2d999c34b22bf4b8b5c9439918e46341f4f0da1b00a6b0479c541161d4aa074abe79c51bf57086e1e84b57ee483cbb2ecf30e8222bc0472436fabfc57dda8".try_into().unwrap();
+    handler.tx_queues.add_signed_tx(context, tx2.clone());
+
+    let tx3: SignedTx = "f86b02830186a0830186a094a8f7c4c78c36e54c3950ad58dad24ca5e0191b2989056bc75e2d631000008025a070b21a24cec13c0569099ee2f8221268103fd609646b73f7c9e85efeb7af5c8ea03d5de75bc12ce28a80f7c0401df6021cc82a334cb1c802c8b9d46223c5c8eb40".try_into().unwrap();
+    handler.tx_queues.add_signed_tx(context, tx3.clone());
+
+    assert_eq!(handler.tx_queues.len(context), 3);
+    assert_eq!(handler.tx_queues.len(handler.get_context()), 0);
+
+    let (block, failed_txs) = handler.finalize_block(context, true).unwrap();
+    assert_eq!(
+        block.transactions,
+        vec![tx1, tx2]
+            .into_iter()
+            .map(|t| t.transaction)
+            .collect::<Vec<_>>()
+    );
+    assert_eq!(
+        failed_txs,
+        vec![tx3]
+            .into_iter()
+            .map(|t| t.transaction)
+            .collect::<Vec<_>>()
+    );
+
+    let state = handler.state.read().unwrap();
+    assert_eq!(
+        state
+            .get(
+                &"0xa8f7c4c78c36e54c3950ad58dad24ca5e0191b29"
+                    .parse()
+                    .unwrap()
+            )
+            .unwrap()
+            .balance,
+        U256::from_str_radix("200000000000000000000", 10).unwrap()
+    );
+    assert_eq!(
+        state
+            .get(
+                &"0x6745f998a96050bb9b0449e6bd4358138a519679"
+                    .parse()
+                    .unwrap()
+            )
+            .unwrap()
+            .balance,
+        U256::from_str_radix("0", 10).unwrap()
+    );
+    assert_eq!(
+        state
+            .get(
+                &"0xc0cd829081485e70348975d325fe5275140277bd"
+                    .parse()
+                    .unwrap()
+            )
+            .unwrap()
+            .balance,
+        U256::from_str_radix("0", 10).unwrap()
+    );
+}