Add recursive bool flag to get_balance

vlayer-xyz · Jun 7, 2024 · 8ed5191 · 8ed5191
1 parent 7fae049
commit 8ed5191
Show file tree

Hide file tree

Showing 4 changed files with 10 additions and 15 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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");