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/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_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 address_pointer = context.build_address_argument_store(address)?;
 
     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_alloca(context.word_type(), "address_pointer");
-    context.build_store(address_pointer, address)?;
+    let address_pointer = context.build_address_argument_store(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")
 }