From 9fb2febfb373fe68237e66f5cf32d783dfa07bf2 Mon Sep 17 00:00:00 2001
From: parazyd <parazyd@dyne.org>
Date: Thu, 18 Jan 2024 12:50:30 +0100
Subject: [PATCH] runtime/import/db: Implement host gas cost for zkas_db_set()

This is unfinished, as we should properly account for the ZkBinary
rather than just charging a fixed cost for every literal/witness/opcode.
---
 src/runtime/import/db.rs   | 26 ++++++++++++++++++++------
 src/runtime/import/util.rs |  1 -
 2 files changed, 20 insertions(+), 7 deletions(-)

diff --git a/src/runtime/import/db.rs b/src/runtime/import/db.rs
index 7cf6e43e4..d419c71a0 100644
--- a/src/runtime/import/db.rs
+++ b/src/runtime/import/db.rs
@@ -786,9 +786,9 @@ pub(crate) fn db_contains_key(ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, ptr_le
 /// Given a zkas circuit, create a VerifyingKey and insert them both into the db.
 /// This function can called only from the Deploy [`ContractSection`].
 /// Returns `0` on success, otherwise returns a (negative) error code.
-pub(crate) fn zkas_db_set(ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, ptr_len: u32) -> i64 {
-    let env = ctx.data();
-    let cid = &env.contract_id;
+pub(crate) fn zkas_db_set(mut ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, ptr_len: u32) -> i64 {
+    let (env, mut store) = ctx.data_and_store_mut();
+    let cid = env.contract_id;
 
     if let Err(e) = acl_allow(env, &[ContractSection::Deploy]) {
         error!(
@@ -798,8 +798,10 @@ pub(crate) fn zkas_db_set(ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, ptr_len: u
         return darkfi_sdk::error::CALLER_ACCESS_DENIED
     }
 
-    let memory_view = env.memory_view(&ctx);
-    let contract_id = &env.contract_id;
+    // Subtract used gas. Here we count the length read from the memory slice.
+    env.subtract_gas(&mut store, ptr_len as u64);
+
+    let memory_view = env.memory_view(&store);
 
     // Ensure that the memory is readable
     let Ok(mem_slice) = ptr.slice(&memory_view, ptr_len) else {
@@ -843,11 +845,18 @@ pub(crate) fn zkas_db_set(ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, ptr_len: u
         }
     };
 
+    // Subtract used gas. We count 100 gas per opcode, witness, and literal.
+    // This is likely bad.
+    // TODO: This should be better-priced.
+    let gas_cost =
+        (zkbin.literals.len() + zkbin.witnesses.len() + zkbin.opcodes.len()) as u64 * 100;
+    env.subtract_gas(&mut store, gas_cost);
+
     // Because of `Runtime::Deploy`, we should be sure that the zkas db is index zero.
     let db_handles = env.db_handles.borrow();
     let db_handle = &db_handles[0];
     // Redundant check
-    if &db_handle.contract_id != contract_id {
+    if db_handle.contract_id != cid {
         error!(
             target: "runtime::db::zkas_db_set",
             "[WASM] [{}] zkas_db_set(): Internal error, zkas db at index 0 incorrect", cid,
@@ -898,6 +907,7 @@ pub(crate) fn zkas_db_set(ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, ptr_len: u
         "[WASM] [{}] zkas_db_set(): Creating VerifyingKey for {} zkas circuit",
         cid, zkbin.namespace,
     );
+
     let witnesses = match empty_witnesses(&zkbin) {
         Ok(w) => w,
         Err(e) => {
@@ -940,6 +950,10 @@ pub(crate) fn zkas_db_set(ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, ptr_len: u
         );
         return darkfi_sdk::error::DB_SET_FAILED
     }
+    drop(db_handles);
+
+    // Subtract used gas. Here we count the bytes written into the db.
+    env.subtract_gas(&mut store, (key.len() + value.len()) as u64);
 
     darkfi_sdk::entrypoint::SUCCESS
 }
diff --git a/src/runtime/import/util.rs b/src/runtime/import/util.rs
index a525234ae..c53b34989 100644
--- a/src/runtime/import/util.rs
+++ b/src/runtime/import/util.rs
@@ -280,7 +280,6 @@ pub(crate) fn get_slot(mut ctx: FunctionEnvMut<Env>, slot: u64) -> i64 {
         return darkfi_sdk::error::CALLER_ACCESS_DENIED
     }
 
-    // TODO: Gas cost
     let ret = match env.blockchain.lock().unwrap().slots.get_by_id(slot) {
         Ok(v) => v,
         Err(e) => {