From 8ed51919dcea3dbd06d384ff1a592bdd3f8c26e4 Mon Sep 17 00:00:00 2001 From: Arkadiusz Konior Date: Fri, 7 Jun 2024 12:24:14 +0200 Subject: [PATCH] Add recursive bool flag to get_balance --- vlayer/ethereum/circuits/lib/src/token.nr | 19 +++++++------------ .../circuits/lib/src/token_int_test.nr | 2 +- .../circuits/is_dao_worthy/src/main.nr | 2 +- .../is_dao_worthy_recursive/src/main.nr | 2 +- 4 files changed, 10 insertions(+), 15 deletions(-) diff --git a/vlayer/ethereum/circuits/lib/src/token.nr b/vlayer/ethereum/circuits/lib/src/token.nr index 35440b853..6c24d3e4c 100644 --- a/vlayer/ethereum/circuits/lib/src/token.nr +++ b/vlayer/ethereum/circuits/lib/src/token.nr @@ -22,22 +22,17 @@ impl ERC20Token { } trait ERC20 { - fn get_balance(self, wallet_address: Address, block_number: u64) -> U128; - fn get_balance_recursive(self, wallet_address: Address, block_number: u64) -> U128; + fn get_balance(self, wallet_address: Address, block_number: u64, recursive: bool) -> U128; } impl ERC20 for ERC20Token { - fn get_balance(self, wallet_address: Address, block_number: u64) -> U128 { + fn get_balance(self, wallet_address: Address, block_number: u64, recursive: bool) -> U128 { let storage_key = self.calculate_balance_storage_key(wallet_address); - let account = get_account_with_storage(self.chain_id, block_number, self.address, storage_key); - let balance = account.values[TOKEN_BALANCE_INDEX]; - - U128::from_integer(bytes32_to_field(balance)) - } - - fn get_balance_recursive(self, wallet_address: Address, block_number: u64) -> U128 { - let storage_key = self.calculate_balance_storage_key(wallet_address); - let account = get_account_with_storage_recursive(self.chain_id, block_number, self.address, storage_key); + let account = if recursive { + get_account_with_storage_recursive(self.chain_id, block_number, self.address, storage_key) + } else { + get_account_with_storage(self.chain_id, block_number, self.address, storage_key) + }; let balance = account.values[TOKEN_BALANCE_INDEX]; U128::from_integer(bytes32_to_field(balance)) diff --git a/vlayer/ethereum/circuits/lib/src/token_int_test.nr b/vlayer/ethereum/circuits/lib/src/token_int_test.nr index 79c8e9edf..544a72574 100644 --- a/vlayer/ethereum/circuits/lib/src/token_int_test.nr +++ b/vlayer/ethereum/circuits/lib/src/token_int_test.nr @@ -38,7 +38,7 @@ mod test_ERC20Token { 0x55, 0xFE, 0x00, 0x2a, 0xef, 0xF0, 0x2F, 0x77, 0x36, 0x4d, 0xe3, 0x39, 0xa1, 0x29, 0x29, 0x23, 0xA1, 0x58, 0x44, 0xB8 ]; let block_number = 19000000; - let balance = usdc_token.get_balance(circle_address, block_number); + let balance = usdc_token.get_balance(circle_address, block_number, false); let expected_usdc_balnce_of_circle = U128::from_integer(125761774888720); assert_eq(expected_usdc_balnce_of_circle, balance); } diff --git a/vlayer/examples/circuits/is_dao_worthy/src/main.nr b/vlayer/examples/circuits/is_dao_worthy/src/main.nr index 250bd8967..07d2235ee 100644 --- a/vlayer/examples/circuits/is_dao_worthy/src/main.nr +++ b/vlayer/examples/circuits/is_dao_worthy/src/main.nr @@ -6,7 +6,7 @@ mod main_test; global MIN_BALANCE = U128::from_integer(100_000_000_000); // $100k fn main(wallet_address: Address, block_number: pub u64) { - let wallet_balance = USDC.get_balance(wallet_address, block_number); + let wallet_balance = USDC.get_balance(wallet_address, block_number, false); let is_balance_greater_or_equal = MIN_BALANCE <= wallet_balance; assert(is_balance_greater_or_equal, "Insufficient USDC balance"); diff --git a/vlayer/examples/circuits/is_dao_worthy_recursive/src/main.nr b/vlayer/examples/circuits/is_dao_worthy_recursive/src/main.nr index c11a10e89..574389b1b 100644 --- a/vlayer/examples/circuits/is_dao_worthy_recursive/src/main.nr +++ b/vlayer/examples/circuits/is_dao_worthy_recursive/src/main.nr @@ -6,7 +6,7 @@ mod main_test; global MIN_BALANCE = U128::from_integer(100_000_000_000); // $100k fn main(wallet_address: Address, block_number: pub u64) { - let wallet_balance = USDC.get_balance_recursive(wallet_address, block_number); + let wallet_balance = USDC.get_balance(wallet_address, block_number, true); let is_balance_greater_or_equal = MIN_BALANCE <= wallet_balance; assert(is_balance_greater_or_equal, "Insufficient USDC balance");