Fix implementation for balance_of (#79)

The balance_of syscall is now available in pallet-revive.
- Fix balance_of implementation to use correct runtime api
- Add build_address_argument_store helper to be used for address arguments

crates/integration/contracts/Value.sol

@@ -5,6 +5,15 @@ pragma solidity ^0.8;
 {
     "differential": true,
     "actions": [
+        {
+            "Upload": {
+                "code": {
+                    "Solidity": {
+                        "contract": "ValueTester"
+                    }
+                }
+            }
+        },
         {
             "Instantiate": {
                 "value": 1024,
@@ -23,31 +32,35 @@ pragma solidity ^0.8;
                 "value": 123,
                 "data": "3fa4f245"
             }
-        },
-        {
-            "Call": {
-                "dest": {
-                    "Instantiated": 0
-                },
-                "data": "52da5fa0"
-            }
         }
     ]
 }
 */
 
-contract Value {
+contract ValueTester {
     constructor() payable {}
 
-    function value() public payable returns (uint ret) {
-        ret = msg.value;
-    }
-
     function balance_self() public view returns (uint ret) {
         ret = address(this).balance;
     }
+}
 
-    function balance_of(address _address) public view returns (uint ret) {
+contract Value {
-        ret = _address.balance;
+    constructor() payable {
+        ValueTester tester = new ValueTester{value: msg.value}();
+
+        // own account
+        assert(address(this).balance == 0);
+
+        // tester account
+        assert(address(tester).balance == msg.value);
+        assert(tester.balance_self() == msg.value);
+
+        // non-existant account
+        assert(address(0xdeadbeef).balance == 0);
+    }
+
+    function value() public payable returns (uint ret) {
+        ret = msg.value;
     }
 }

crates/integration/src/tests.rs

@@ -263,32 +263,6 @@ fn create2_failure() {
     assert_eq!(output.flags, ReturnFlags::Revert);
 }
 
-
-#[test]
-fn balance() {
-    let (_, output) = assert_success(&Contract::value_balance_of(Default::default()), false);
-
-    let expected = U256::ZERO;
-    let received = U256::from_be_slice(&output.data);
-    assert_eq!(expected, received);
-
-    let expected = U256::from(54589);
-    let (mut state, address) = State::new_deployed(Contract::value_balance_of(Default::default()));
-    state.accounts_mut().get_mut(&address).unwrap().value = expected;
-
-    let contract = Contract::value_balance_of(address);
-    let (_, output) = state
-        .transaction()
-        .with_default_account(&contract.pvm_runtime)
-        .calldata(contract.calldata)
-        .call();
-
-    assert_eq!(ReturnFlags::Success, output.flags);
-
-    let received = U256::from_be_slice(&output.data);
-    assert_eq!(expected, received)
-}
-
 #[test]
 fn ext_code_size() {
     let contract = Contract::ext_code_size(Transaction::default_address());

crates/llvm-context/src/polkavm/const/runtime_api.rs

@@ -17,6 +17,8 @@ pub mod imports {
 
     pub static BALANCE: &str = "balance";
 
+    pub static BALANCE_OF: &str = "balance_of";
+
     pub static BLOCK_NUMBER: &str = "block_number";
 
     pub static CHAIN_ID: &str = "chain_id";
@@ -55,9 +57,10 @@ pub mod imports {
 
     /// All imported runtime API symbols.
     /// Useful for configuring common attributes and linkage.
-    pub static IMPORTS: [&str; 20] = [
+    pub static IMPORTS: [&str; 21] = [
         ADDRESS,
         BALANCE,
+        BALANCE_OF,
         BLOCK_NUMBER,
         CALL,
         CALLER,

crates/llvm-context/src/polkavm/context/mod.rs

@@ -607,6 +607,23 @@ where
         Pointer::new(r#type, AddressSpace::Stack, pointer)
     }
 
+    /// Truncate `address` to the ethereum address length and store it as bytes on the stack.
+    /// The stack allocation will be at the function entry. Returns the stack pointer.
+    /// This helper should be used when passing address arguments to the runtime, ensuring correct size and endianness.
+    pub fn build_address_argument_store(
+        &self,
+        address: inkwell::values::IntValue<'ctx>,
+    ) -> anyhow::Result<Pointer<'ctx>> {
+        let address_type = self.integer_type(revive_common::BIT_LENGTH_ETH_ADDRESS);
+        let address_pointer = self.build_alloca_at_entry(address_type, "address_pointer");
+        let address_truncated =
+            self.builder()
+                .build_int_truncate(address, address_type, "address_truncated")?;
+        let address_swapped = self.build_byte_swap(address_truncated.into())?;
+        self.build_store(address_pointer, address_swapped)?;
+        Ok(address_pointer)
+    }
+
     /// Load the address at given pointer and zero extend it to the VM word size.
     pub fn build_load_address(
         &self,

crates/llvm-context/src/polkavm/evm/call.rs

@@ -27,14 +27,7 @@ pub fn call<'ctx, D>(
 where
     D: Dependency + Clone,
 {
-    let address_type = context.integer_type(revive_common::BIT_LENGTH_ETH_ADDRESS);
+    let address_pointer = context.build_address_argument_store(address)?;
-    let address_pointer = context.build_alloca_at_entry(address_type, "address_pointer");
-    let address_truncated =
-        context
-            .builder()
-            .build_int_truncate(address, address_type, "address_truncated")?;
-    let address_swapped = context.build_byte_swap(address_truncated.into())?;
-    context.build_store(address_pointer, address_swapped)?;
 
     let value = value.unwrap_or_else(|| context.word_const(0));
     let value_pointer = context.build_alloca_at_entry(context.word_type(), "value_pointer");

crates/llvm-context/src/polkavm/evm/ether_gas.rs

@@ -40,22 +40,19 @@ pub fn balance<'ctx, D>(
 where
     D: Dependency + Clone,
 {
-    let balance_pointer = context.build_alloca(context.word_type(), "balance_pointer");
+    let address_pointer = context.build_address_argument_store(address)?;
-    let address_pointer = context.build_alloca(context.word_type(), "address_pointer");
-    context.build_store(address_pointer, address)?;
 
+    let balance_pointer = context.build_alloca(context.word_type(), "balance_pointer");
     let balance = context.builder().build_ptr_to_int(
         balance_pointer.value,
         context.xlen_type(),
         "balance",
     )?;
-    let _address = context.builder().build_ptr_to_int(
-        address_pointer.value,
-        context.xlen_type(),
-        "address",
-    )?;
 
-    context.build_runtime_call(runtime_api::imports::BALANCE, &[balance.into()]);
+    context.build_runtime_call(
+        runtime_api::imports::BALANCE_OF,
+        &[address_pointer.to_int(context).into(), balance.into()],
+    );
 
     context.build_load(balance_pointer, "balance")
 }

crates/runtime-api/src/polkavm_imports.c

@@ -105,6 +105,8 @@ POLKAVM_IMPORT(void, gas_left, uint32_t, uint32_t)
 
 POLKAVM_IMPORT(void, balance, uint32_t)
 
+POLKAVM_IMPORT(void, balance_of, uint32_t, uint32_t)
+
 POLKAVM_IMPORT(void, chain_id, uint32_t)
 
 POLKAVM_IMPORT(void, now, uint32_t)