Ignore gas on mock-function-call (#1585)

Author: richardpringle

x/contracts/examples/counter-external/src/lib.rs

@@ -41,7 +41,7 @@ mod tests {
 
         // mock `get_value` external contract call to return `value`
         let value = 5_u64;
-        ctx.mock_function_call(external, "get_value", of, 1_000_000, 0, value);
+        ctx.mock_function_call(external, "get_value", of, 0, value);
 
         let value = get_value(&mut ctx, external, of);
         assert_eq!(value, 5);

x/contracts/wasmlanche/src/context.rs

@@ -4,7 +4,7 @@
 extern crate alloc;
 
 use crate::{
-    host::Accessor,
+    host::{Accessor, CallContractArgs},
     state::{Cache, Error, IntoPairs, Schema},
     types::{Address, ContractId},
     Gas, Id,
@@ -238,10 +238,15 @@ impl Context {
         max_units: Gas,
         value: u64,
     ) -> Result<T, ExternalCallError> {
-        let args = CallContractArgs::new(&address, function_name, args, max_units, value);
-        let args_bytes = borsh::to_vec(&args).expect("failed to serialize args");
         self.state_cache.flush();
-        let bytes = self.host_accessor.call_contract(&args_bytes);
+
+        let bytes = self.host_accessor.call_contract(&CallContractArgs {
+            address,
+            function_name,
+            args,
+            max_units,
+            value,
+        });
 
         borsh::from_slice(&bytes).expect("failed to deserialize")
     }
@@ -276,22 +281,27 @@ impl Context {
     /// Panics if serialization fails.
     pub fn mock_function_call<T, U>(
         &self,
-        target: Address,
-        function: &str,
+        address: Address,
+        function_name: &str,
         args: T,
-        // TODO: remove this
-        max_units: Gas,
-        max_value: u64,
+        value: u64,
         result: U,
     ) where
         T: BorshSerialize,
         U: BorshSerialize,
     {
         use crate::host::CALL_FUNCTION_PREFIX;
 
-        let args = borsh::to_vec(&args).expect("error serializing result");
+        let args = &borsh::to_vec(&args).expect("error serializing result");
+
+        let contract_args = CallContractArgs {
+            address,
+            function_name,
+            args,
+            value,
+            max_units: 0,
+        };
 
-        let contract_args = CallContractArgs::new(&target, function, &args, max_units, max_value);
         let contract_args = borsh::to_vec(&(CALL_FUNCTION_PREFIX, contract_args))
             .expect("error serializing result");
 
@@ -396,30 +406,3 @@ impl Default for Context {
         Self::new()
     }
 }
-
-#[derive(BorshSerialize)]
-pub struct CallContractArgs<'a> {
-    target: &'a Address,
-    function: &'a [u8],
-    args: &'a [u8],
-    max_units: Gas,
-    max_value: u64,
-}
-
-impl<'a> CallContractArgs<'a> {
-    pub fn new(
-        target: &'a Address,
-        function: &'a str,
-        args: &'a [u8],
-        max_units: Gas,
-        max_value: u64,
-    ) -> Self {
-        Self {
-            target,
-            function: function.as_bytes(),
-            args,
-            max_units,
-            max_value,
-        }
-    }
-}

x/contracts/wasmlanche/src/host.rs

@@ -1,12 +1,62 @@
 // Copyright (C) 2024, Ava Labs, Inc. All rights reserved.
 // See the file LICENSE for licensing terms.
 
+use crate::{
+    borsh::{self, BorshSerialize},
+    Address, Gas,
+};
+use cfg_if::cfg_if;
+
+cfg_if! {
+    if #[cfg(feature = "test")] {
+        pub use test_wrappers::*;
+    } else {
+        pub use external_wrappers::*;
+    }
+}
+
 pub struct StateAccessor;
 
+pub(crate) struct CallContractArgs<'a> {
+    pub(crate) address: Address,
+    pub(crate) function_name: &'a str,
+    pub(crate) args: &'a [u8],
+    pub(crate) max_units: Gas,
+    pub(crate) value: u64,
+}
+
+impl BorshSerialize for CallContractArgs<'_> {
+    fn serialize<W: borsh::io::Write>(&self, writer: &mut W) -> borsh::io::Result<()> {
+        let Self {
+            address,
+            function_name,
+            args,
+            max_units,
+            value,
+        } = self;
+
+        address.serialize(writer)?;
+        function_name.serialize(writer)?;
+        args.serialize(writer)?;
+
+        cfg_if! {
+            if #[cfg(feature = "test")] {
+                let _ = max_units;
+            } else {
+                max_units.serialize(writer)?;
+            }
+        }
+
+        value.serialize(writer)?;
+
+        Ok(())
+    }
+}
+
 #[cfg(feature = "test")]
 mod test_wrappers {
-    use crate::host::StateAccessor;
-    use crate::{Address, Gas, HostPtr};
+    use super::CallContractArgs;
+    use crate::{host::StateAccessor, Address, Gas, HostPtr};
     use core::cell::{Cell, RefCell};
 
     pub const BALANCE_PREFIX: u8 = 0;
@@ -29,7 +79,7 @@ mod test_wrappers {
     #[derive(Clone)]
     #[cfg_attr(feature = "debug", derive(Debug))]
     pub struct Accessor {
-        pub state: MockState,
+        state: MockState,
     }
 
     impl Default for Accessor {
@@ -70,12 +120,15 @@ mod test_wrappers {
             val
         }
 
-        pub fn call_contract(&self, args: &[u8]) -> HostPtr {
-            let key = [CALL_FUNCTION_PREFIX]
-                .iter()
-                .chain(args.iter())
-                .copied()
-                .collect::<Vec<u8>>();
+        pub fn call_contract(&self, args: &CallContractArgs) -> HostPtr {
+            let key = {
+                // same default as borsh::to_vec uses
+                let mut key = Vec::with_capacity(1024);
+                key.push(CALL_FUNCTION_PREFIX);
+                borsh::to_writer(&mut key, args).expect("failed to serialize call-contract args");
+                key
+            };
+
             let val = self.state.get(&key);
 
             assert!(
@@ -199,6 +252,7 @@ mod test_wrappers {
 
 #[cfg(not(feature = "test"))]
 mod external_wrappers {
+    use super::CallContractArgs;
     use crate::host::StateAccessor;
     use crate::HostPtr;
 
@@ -248,13 +302,15 @@ mod external_wrappers {
             unsafe { deploy(args.as_ptr(), args.len()) }
         }
 
-        pub fn call_contract(&self, args: &[u8]) -> HostPtr {
+        pub fn call_contract(&self, args: &CallContractArgs) -> HostPtr {
             #[link(wasm_import_module = "contract")]
             extern "C" {
                 #[link_name = "call_contract"]
                 fn call_contract(ptr: *const u8, len: usize) -> HostPtr;
             }
 
+            let args = borsh::to_vec(args).expect("failed to serialize args");
+
             unsafe { call_contract(args.as_ptr(), args.len()) }
         }
 
@@ -289,9 +345,3 @@ mod external_wrappers {
         }
     }
 }
-
-#[cfg(feature = "test")]
-pub use test_wrappers::*;
-
-#[cfg(not(feature = "test"))]
-pub use external_wrappers::*;