wasm: make functions use the same generic underlying call(), and able to return data through env with set_return_data()

bin/dao/daod/src/util.rs

@@ -145,6 +145,7 @@ pub struct Transaction {
     // TODO: this is wrong. It should be Vec<Vec<Signature>>
     // each Vec<Signature> correspond to ONE function call
     pub signatures: Vec<Signature>,
+    // pub proofs: Vec<Proof>,
 }
 // ANCHOR_END: transaction
 

example/smart-contract/src/lib.rs

@@ -4,8 +4,8 @@ use darkfi_sdk::{
     define_contract,
     error::ContractResult,
     msg,
-    state::set_update,
     tx::FuncCall,
+    util::set_return_data,
 };
 use darkfi_serial::{deserialize, serialize, SerialDecodable, SerialEncodable};
 
@@ -34,6 +34,13 @@ pub struct FooCallData {
     pub b: u64,
 }
 
+impl FooCallData {
+    //fn zk_public_values(&self) -> Vec<(String, Vec<DrkCircuitField>)>;
+
+    fn get_metadata(&self) {
+    }
+}
+
 #[derive(SerialEncodable, SerialDecodable)]
 pub struct BarArgs {
     pub x: u32,
@@ -45,7 +52,13 @@ pub struct FooUpdate {
     pub age: u32,
 }
 
-define_contract!(init: init_contract, exec: process_instruction, apply: process_update);
+define_contract!(
+    init: init_contract,
+    exec: process_instruction,
+    apply: process_update,
+
+    metadata: get_metadata
+);
 
 fn init_contract(cid: ContractId, _ix: &[u8]) -> ContractResult {
     msg!("wakeup wagies!");
@@ -62,6 +75,25 @@ fn init_contract(cid: ContractId, _ix: &[u8]) -> ContractResult {
     Ok(())
 }
 
+fn get_metadata(cid: ContractId, ix: &[u8]) -> ContractResult {
+    match Function::from(ix[0]) {
+        Function::Foo => {
+            let tx_data = &ix[1..];
+            // ...
+            let (func_call_index, func_calls): (u32, Vec<FuncCall>) = deserialize(tx_data)?;
+            let call_data: FooCallData =
+                deserialize(&func_calls[func_call_index as usize].call_data)?;
+
+            // Convert call_data to halo2 public inputs
+            // Pass this to the env
+        }
+        Function::Bar => {
+            // ...
+        }
+    }
+    Ok(())
+}
+
 // This is the main entrypoint function where the payload is fed.
 // Through here, you can branch out into different functions inside
 // this library.
@@ -79,7 +111,7 @@ fn process_instruction(cid: ContractId, ix: &[u8]) -> ContractResult {
 
             let mut update_data = vec![Function::Foo as u8];
             update_data.extend_from_slice(&serialize(&update));
-            set_update(&update_data)?;
+            set_return_data(&update_data)?;
             msg!("update is set!");
 
             // Example: try to get a value from the db

example/smart-contract/tests/runtime.rs

@@ -77,12 +77,12 @@ fn run_contract() -> Result<()> {
     // ============================================================
     // Serialize the payload into the runtime format and execute it
     // ============================================================
-    runtime.exec(&payload)?;
+    let update = runtime.exec(&payload)?;
 
     // =====================================================
     // If exec was successful, try to apply the state change
     // =====================================================
-    runtime.apply()?;
+    runtime.apply(&update)?;
 
     Ok(())
 }

src/runtime/import/chain_state.rs

@@ -21,33 +21,6 @@ use wasmer::{FunctionEnvMut, WasmPtr};
 
 use crate::runtime::vm_runtime::{ContractSection, Env};
 
-pub(crate) fn set_update(mut ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, len: u32) -> i32 {
-    let env = ctx.data();
-    match env.contract_section {
-        ContractSection::Exec => {
-            let memory_view = env.memory_view(&ctx);
-
-            let Ok(slice) = ptr.slice(&memory_view, len) else {
-                return -2
-            };
-
-            let Ok(update_data) = slice.read_to_vec() else {
-                return -2;
-            };
-
-            // This function should only ever be called once on the runtime.
-            if !env.contract_update.take().is_none() {
-                return -3
-            }
-            let func_id = update_data[0];
-            let update_data = &update_data[1..];
-            env.contract_update.set(Some((func_id, update_data.to_vec())));
-            0
-        }
-        _ => -1,
-    }
-}
-
 /// Try to read a `Nullifier` from the given pointer and check if it's
 /// an existing nullifier in the blockchain state machine.
 pub(crate) fn nullifier_exists(mut ctx: FunctionEnvMut<Env>, ptr: u32, len: u32) -> i32 {

src/runtime/import/util.rs

@@ -19,7 +19,7 @@
 use log::error;
 use wasmer::{FunctionEnvMut, WasmPtr};
 
-use crate::runtime::vm_runtime::Env;
+use crate::runtime::vm_runtime::{ContractSection, Env};
 
 /// Host function for logging strings.
 /// This is injected into the runtime with wasmer's `imports!` macro.
@@ -38,3 +38,29 @@ pub(crate) fn drk_log(mut ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, len: u32)
         }
     }
 }
+
+pub(crate) fn set_return_data(mut ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, len: u32) -> i32 {
+    let env = ctx.data();
+    match env.contract_section {
+        ContractSection::Exec => {
+            let memory_view = env.memory_view(&ctx);
+
+            let Ok(slice) = ptr.slice(&memory_view, len) else {
+                return -2
+            };
+
+            let Ok(update_data) = slice.read_to_vec() else {
+                return -2;
+            };
+
+            // This function should only ever be called once on the runtime.
+            if !env.contract_return_data.take().is_none() {
+                return -3
+            }
+            env.contract_return_data.set(Some(update_data));
+            0
+        }
+        _ => -1,
+    }
+}
+

src/runtime/vm_runtime.rs

@@ -38,20 +38,34 @@ use crate::{blockchain::Blockchain, Error, Result};
 
 /// Name of the wasm linear memory in our guest module
 const MEMORY: &str = "memory";
+
 /// Hardcoded setup function of a contract
-pub const INITIALIZE: &str = "__initialize";
+const INITIALIZE: &str = "__initialize";
 /// Hardcoded entrypoint function of a contract
-pub const ENTRYPOINT: &str = "__entrypoint";
+const ENTRYPOINT: &str = "__entrypoint";
 /// Hardcoded apply function of a contract
-pub const UPDATE: &str = "__update";
+const UPDATE: &str = "__update";
+
 /// Gas limit for a contract
 const GAS_LIMIT: u64 = 200000;
 
+#[derive(Clone, Copy)]
 pub enum ContractSection {
-    Null,
     Deploy,
     Exec,
     Update,
+    Null,
+}
+
+impl ContractSection {
+    fn name(&self) -> &str {
+        match self {
+            Self::Deploy => "__initialize",
+            Self::Exec => "__entrypoint",
+            Self::Update => "__update",
+            Self::Null => unreachable!(),
+        }
+    }
 }
 
 /// The wasm vm runtime instantiated for every smart contract that runs.
@@ -65,7 +79,7 @@ pub struct Env {
     /// The contract section being executed
     pub contract_section: ContractSection,
     /// State update produced by a smart contract function call
-    pub contract_update: Cell<Option<(u8, Vec<u8>)>>,
+    pub contract_return_data: Cell<Option<Vec<u8>>>,
     /// Logs produced by the contract
     pub logs: RefCell<Vec<String>>,
     /// Direct memory access to the VM
@@ -139,7 +153,7 @@ impl Runtime {
                 db_handles,
                 contract_id,
                 contract_section: ContractSection::Null,
-                contract_update: Cell::new(None),
+                contract_return_data: Cell::new(None),
                 logs,
                 memory: None,
             },
@@ -165,10 +179,10 @@ impl Runtime {
                     import::chain_state::is_valid_merkle,
                 ),
 
-                "set_update_" => Function::new_typed_with_env(
+                "set_return_data_" => Function::new_typed_with_env(
                     &mut store,
                     &ctx,
-                    import::chain_state::set_update,
+                    import::util::set_return_data,
                 ),
 
                 "db_init_" => Function::new_typed_with_env(
@@ -218,6 +232,63 @@ impl Runtime {
         Ok(Self { instance, store, ctx })
     }
 
+    fn call(&mut self, section: ContractSection, payload: &[u8]) -> Result<Vec<u8>> {
+        debug!(target: "runtime", "Calling {} method", section.name());
+
+        let mut env_mut = self.ctx.as_mut(&mut self.store);
+        env_mut.contract_section = section;
+        assert!(env_mut.contract_return_data.take().is_none());
+        env_mut.contract_return_data.set(None);
+
+        // Serialize the payload for the format the wasm runtime is expecting.
+        let payload = Self::serialize_payload(&env_mut.contract_id, payload);
+
+        // Allocate enough memory for the payload and copy it into the memory.
+        let pages_required = payload.len() / WASM_PAGE_SIZE + 1;
+        self.set_memory_page_size(pages_required as u32)?;
+        self.copy_to_memory(&payload)?;
+
+        debug!(target: "runtime", "Getting initialize function");
+        let entrypoint = self.instance.exports.get_function(section.name())?;
+
+        debug!(target: "runtime", "Executing wasm");
+        let ret = match entrypoint.call(&mut self.store, &[Value::I32(0 as i32)]) {
+            Ok(retvals) => {
+                self.print_logs();
+                debug!(target: "runtime", "{}", self.gas_info());
+                retvals
+            }
+            Err(e) => {
+                self.print_logs();
+                debug!(target: "runtime", "{}", self.gas_info());
+                // WasmerRuntimeError panics are handled here. Return from run() immediately.
+                return Err(e.into())
+            }
+        };
+
+        debug!(target: "runtime", "wasm executed successfully");
+        debug!(target: "runtime", "Contract returned: {:?}", ret[0]);
+
+        let mut env_mut = self.ctx.as_mut(&mut self.store);
+        env_mut.contract_section = ContractSection::Null;
+        let retdata = match env_mut.contract_return_data.take() {
+            Some(retdata) => retdata,
+            None => Vec::new()
+        };
+
+        let retval = match ret[0] {
+            Value::I64(v) => v as u64,
+            _ => unreachable!(),
+        };
+
+        match retval {
+            entrypoint::SUCCESS => Ok(retdata),
+            // FIXME: we should be able to see the error returned from the contract
+            // We can put sdk::Error inside of this.
+            _ => Err(Error::ContractInitError(retval)),
+        }
+    }
+
     /// This function runs when a smart contract is initially deployed, or re-deployed.
     /// The runtime will look for an [`INITIALIZE`] symbol in the wasm code, and execute
     /// it if found. Optionally, it is possible to pass in a payload for any kind of special
@@ -227,8 +298,13 @@ impl Runtime {
     /// The permissions for this are handled by the `ContractId` in the sled db API so we
     /// assume that the contract is only able to do write operations on its own sled trees.
     pub fn deploy(&mut self, payload: &[u8]) -> Result<()> {
+        let _ = self.call(ContractSection::Deploy, payload)?;
+        Ok(())
+        /*
         let mut env_mut = self.ctx.as_mut(&mut self.store);
         env_mut.contract_section = ContractSection::Deploy;
+        assert!(env_mut.contract_return_data.take().is_none());
+        env_mut.contract_return_data.set(None);
 
         // Serialize the payload for the format the wasm runtime is expecting.
         let payload = Self::serialize_payload(&env_mut.contract_id, payload);
@@ -259,6 +335,10 @@ impl Runtime {
         debug!(target: "wasm_runtime::deploy", "wasm executed successfully");
         debug!(target: "wasm_runtime::deploy", "Contract returned: {:?}", ret[0]);
 
+        let mut env_mut = self.ctx.as_mut(&mut self.store);
+        env_mut.contract_section = ContractSection::Null;
+        let retdata = env_mut.contract_return_data.take();
+
         let retval = match ret[0] {
             Value::I64(v) => v as u64,
             _ => unreachable!(),
@@ -270,15 +350,20 @@ impl Runtime {
             // We can put sdk::Error inside of this.
             _ => Err(Error::ContractInitError(retval)),
         }
+        */
     }
 
     /// This funcion runs when someone wants to execute a smart contract.
     /// The runtime will look for an [`ENTRYPOINT`] symbol in the wasm code, and
     /// execute it if found. A payload is also passed as an instruction that can
     /// be used inside the vm by the runtime.
-    pub fn exec(&mut self, payload: &[u8]) -> Result<()> {
+    pub fn exec(&mut self, payload: &[u8]) -> Result<Vec<u8>> {
+        self.call(ContractSection::Exec, payload)
+        /*
         let mut env_mut = self.ctx.as_mut(&mut self.store);
         env_mut.contract_section = ContractSection::Exec;
+        assert!(env_mut.contract_return_data.take().is_none());
+        env_mut.contract_return_data.set(None);
 
         // Serialize the payload for the format the wasm runtime is expecting.
         let payload = Self::serialize_payload(&env_mut.contract_id, payload);
@@ -319,6 +404,7 @@ impl Runtime {
             entrypoint::SUCCESS => Ok(()),
             _ => Err(Error::ContractExecError(retval)),
         }
+        */
     }
 
     /// This function runs after successful execution of [`exec`] and tries to
@@ -326,9 +412,14 @@ impl Runtime {
     /// The runtime will lok for an [`UPDATE`] symbol in the wasm code, and execute
     /// it if found. The function does not take an arbitrary payload, but just takes
     /// a state update from `env` and passes it into the wasm runtime.
-    pub fn apply(&mut self) -> Result<()> {
+    pub fn apply(&mut self, update: &[u8]) -> Result<()> {
+        let _ = self.call(ContractSection::Update, update)?;
+        Ok(())
+        /*
         let mut env_mut = self.ctx.as_mut(&mut self.store);
         env_mut.contract_section = ContractSection::Update;
+        assert!(env_mut.contract_return_data.take().is_none());
+        env_mut.contract_return_data.set(None);
 
         // Take the update data from env, and serialize it for the format the wasm
         // runtime is expecting.
@@ -374,12 +465,13 @@ impl Runtime {
             entrypoint::SUCCESS => Ok(()),
             _ => Err(Error::ContractInitError(retval)),
         }
+        */
     }
 
     fn print_logs(&self) {
         let logs = self.ctx.as_ref(&self.store).logs.borrow();
         for msg in logs.iter() {
-            debug!(target: "wasm_runtime::print_logs", "Contract log: {}", msg);
+            debug!(target: "runtime", "Contract log: {}", msg);
         }
     }
 

src/sdk/src/entrypoint.rs

@@ -25,7 +25,12 @@ pub const SUCCESS: u64 = 0;
 
 #[macro_export]
 macro_rules! define_contract {
-    (init: $init_func:ident, exec: $exec_func:ident, apply: $apply_func:ident) => {
+    (
+        init: $init_func:ident,
+        exec: $exec_func:ident,
+        apply: $apply_func:ident,
+        metadata: $metadata_func:ident
+    ) => {
         /// # Safety
         #[no_mangle]
         pub unsafe extern "C" fn __initialize(input: *mut u8) -> u64 {
@@ -54,6 +59,15 @@ macro_rules! define_contract {
                 Err(e) => e.into(),
             }
         }
+        #[no_mangle]
+        pub unsafe extern "C" fn __get_metadata(input: *mut u8) -> u64 {
+            let (contract_id, instruction_data) = $crate::entrypoint::deserialize(input);
+
+            match $metadata_func(contract_id, &instruction_data) {
+                Ok(()) => $crate::entrypoint::SUCCESS,
+                Err(e) => e.into(),
+            }
+        }
     };
 }
 

src/sdk/src/lib.rs

@@ -39,3 +39,7 @@ pub mod state;
 
 /// Transaction structure
 pub mod tx;
+
+/// Utilities
+pub mod util;
+

src/sdk/src/state.rs

@@ -21,21 +21,6 @@ use super::{
     error::ContractError,
 };
 
-pub fn set_update(update_data: &[u8]) -> Result<(), ContractError> {
-    #[cfg(target_arch = "wasm32")]
-    unsafe {
-        return match set_update_(update_data.as_ptr(), update_data.len() as u32) {
-            0 => Ok(()),
-            -1 => Err(ContractError::SetUpdateError),
-            -2 => Err(ContractError::UpdateAlreadySet),
-            _ => unreachable!(),
-        }
-    }
-
-    #[cfg(not(target_arch = "wasm32"))]
-    unimplemented!();
-}
-
 pub fn nullifier_exists(nullifier: &Nullifier) -> Result<bool, ContractError> {
     #[cfg(target_arch = "wasm32")]
     unsafe {
@@ -74,8 +59,6 @@ pub fn is_valid_merkle(merkle_root: &MerkleNode) -> Result<bool, ContractError>
 
 #[cfg(target_arch = "wasm32")]
 extern "C" {
-    fn get_update_() -> i32;
-    fn set_update_(ptr: *const u8, len: u32) -> i32;
     fn nullifier_exists_(ptr: *const u8, len: u32) -> i32;
     fn is_valid_merkle_(ptr: *const u8, len: u32) -> i32;
 }

src/sdk/src/util.rs

@@ -0,0 +1,22 @@
+use super::error::ContractError;
+
+pub fn set_return_data(data: &[u8]) -> Result<(), ContractError> {
+    #[cfg(target_arch = "wasm32")]
+    unsafe {
+        return match set_return_data_(data.as_ptr(), data.len() as u32) {
+            0 => Ok(()),
+            -1 => Err(ContractError::SetUpdateError),
+            -2 => Err(ContractError::UpdateAlreadySet),
+            _ => unreachable!(),
+        }
+    }
+
+    #[cfg(not(target_arch = "wasm32"))]
+    unimplemented!();
+}
+
+#[cfg(target_arch = "wasm32")]
+extern "C" {
+    fn set_return_data_(ptr: *const u8, len: u32) -> i32;
+}
+