Allow for a HostQuery to take &mut Session

contracts/host/src/lib.rs

@@ -45,6 +45,12 @@ impl Hoster {
             false => String::from("PROOF IS INVALID"),
         }
     }
+
+    /// Call 'get_counter' function available from the host, and return the
+    /// result
+    pub fn host_get_counter(&self) -> i64 {
+        uplink::host_query("get_counter", ())
+    }
 }
 
 /// Expose `Hoster::host_hash()` to the host
@@ -60,3 +66,9 @@ unsafe fn host_verify(arg_len: u32) -> u32 {
         STATE.host_verify(proof, public_inputs)
     })
 }
+
+/// Expose `Hoster::host_counter()` to the host
+#[no_mangle]
+unsafe fn host_get_counter(arg_len: u32) -> u32 {
+    uplink::wrap_call(arg_len, |()| STATE.host_get_counter())
+}

piecrust/CHANGELOG.md

@@ -13,6 +13,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Changed
 
+- Allow for a `HostQuery` to take a `&mut Session` [#327]
 - Change to have one instance per contract function call [#325]
 - Upgrade `dusk-wasmtime` to version `17`
 
@@ -359,6 +360,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 [#234]: https://github.com/dusk-network/piecrust/pull/234
 
 <!-- ISSUES -->
+[#327]: https://github.com/dusk-network/piecrust/issues/327
 [#325]: https://github.com/dusk-network/piecrust/issues/325
 [#301]: https://github.com/dusk-network/piecrust/issues/313
 [#301]: https://github.com/dusk-network/piecrust/issues/301

piecrust/src/session.rs

@@ -552,7 +552,10 @@ impl Session {
         buf: &mut [u8],
         arg_len: u32,
     ) -> Option<u32> {
-        self.inner.host_queries.call(name, buf, arg_len)
+        let mut session = self.clone();
+        self.inner
+            .host_queries
+            .call(&mut session, name, buf, arg_len)
     }
 
     pub(crate) fn nth_from_top(&self, n: usize) -> Option<CallTreeElem> {

piecrust/src/vm.rs

@@ -236,8 +236,16 @@ impl HostQueries {
         self.map.insert(name.into(), Arc::new(query));
     }
 
-    pub fn call(&self, name: &str, buf: &mut [u8], len: u32) -> Option<u32> {
-        self.map.get(name).map(|host_query| host_query(buf, len))
+    pub fn call(
+        &self,
+        session: &mut Session,
+        name: &str,
+        buf: &mut [u8],
+        len: u32,
+    ) -> Option<u32> {
+        self.map
+            .get(name)
+            .map(|host_query| host_query(session, buf, len))
     }
 }
 
@@ -248,5 +256,18 @@ impl HostQueries {
 /// function, and should be processed first. Once this is done, the implementor
 /// should emplace the return of the query in the same buffer, and return the
 /// length written.
-pub trait HostQuery: Send + Sync + Fn(&mut [u8], u32) -> u32 {}
-impl<F> HostQuery for F where F: Send + Sync + Fn(&mut [u8], u32) -> u32 {}
+///
+/// The host query will have access to the underlying session, and can use it to
+/// perform calls to other contracts.
+///
+/// # Panics
+/// If any error occurs during the execution, the implementer is encouraged to
+/// signal this error by panicking.
+pub trait HostQuery:
+    Send + Sync + Fn(&mut Session, &mut [u8], u32) -> u32
+{
+}
+impl<F> HostQuery for F where
+    F: Send + Sync + Fn(&mut Session, &mut [u8], u32) -> u32
+{
+}

piecrust/tests/host.rs

@@ -6,7 +6,10 @@
 
 use dusk_plonk::prelude::*;
 use once_cell::sync::Lazy;
-use piecrust::{contract_bytecode, ContractData, Error, SessionData, VM};
+use piecrust::{
+    contract_bytecode, ContractData, Error, Session, SessionData, VM,
+};
+use piecrust_uplink::ContractId;
 use rand::rngs::OsRng;
 use rkyv::Deserialize;
 
@@ -30,7 +33,7 @@ fn get_prover_verifier() -> &'static (Prover, Verifier) {
     &PROVER_VERIFIER
 }
 
-fn hash(buf: &mut [u8], len: u32) -> u32 {
+fn hash(_: &mut Session, buf: &mut [u8], len: u32) -> u32 {
     let a = unsafe { rkyv::archived_root::<Vec<u8>>(&buf[..len as usize]) };
     let v: Vec<u8> = a.deserialize(&mut rkyv::Infallible).unwrap();
 
@@ -40,7 +43,7 @@ fn hash(buf: &mut [u8], len: u32) -> u32 {
     32
 }
 
-fn verify_proof(buf: &mut [u8], len: u32) -> u32 {
+fn verify_proof(_: &mut Session, buf: &mut [u8], len: u32) -> u32 {
     let a = unsafe {
         rkyv::archived_root::<(Proof, Vec<BlsScalar>)>(&buf[..len as usize])
     };
@@ -58,10 +61,24 @@ fn verify_proof(buf: &mut [u8], len: u32) -> u32 {
     valid_bytes.len() as u32
 }
 
+pub const COUNTER_ID: ContractId = ContractId::from_bytes([1; 32]);
+
+fn get_counter(session: &mut Session, buf: &mut [u8], _: u32) -> u32 {
+    let receipt = session
+        .call_raw(COUNTER_ID, "read_value", &*buf, u64::MAX)
+        .expect("calling the counter contract should succeed");
+
+    let data = receipt.data;
+    buf[..data.len()].copy_from_slice(&data);
+
+    data.len() as u32
+}
+
 fn new_ephemeral_vm() -> Result<VM, Error> {
     let mut vm = VM::ephemeral()?;
     vm.register_host_query("hash", hash);
     vm.register_host_query("verify_proof", verify_proof);
+    vm.register_host_query("get_counter", get_counter);
     Ok(vm)
 }
 
@@ -87,6 +104,35 @@ pub fn host_hash() -> Result<(), Error> {
     Ok(())
 }
 
+/// Queries a contract for the value held in the counter contract through the
+/// host, using a host query.
+#[test]
+fn host_counter() -> Result<(), Error> {
+    let vm = new_ephemeral_vm()?;
+
+    let mut session = vm.session(SessionData::builder())?;
+
+    session.deploy(
+        contract_bytecode!("counter"),
+        ContractData::builder(OWNER).contract_id(COUNTER_ID),
+        LIMIT,
+    )?;
+
+    let id = session.deploy(
+        contract_bytecode!("host"),
+        ContractData::builder(OWNER),
+        LIMIT,
+    )?;
+
+    let counter = session
+        .call::<_, i64>(id, "host_get_counter", &(), LIMIT)?
+        .data;
+
+    assert_eq!(counter, 0xfc);
+
+    Ok(())
+}
+
 /// Proves that we know a number `c` such that `a + b = c`.
 #[derive(Default)]
 struct TestCircuit {