runtime/import: Implement tx-local db functions

src/runtime/import/db/db_contains_key.rs

@@ -21,12 +21,10 @@ use tracing::error;
 use wasmer::{FunctionEnvMut, WasmPtr};
 
 use crate::runtime::{
-    import::acl::acl_allow,
+    import::{acl::acl_allow, util::wasm_mem_read},
     vm_runtime::{ContractSection, Env},
 };
 
-use super::util::wasm_mem_read;
-
 /// Check if an on-chain database contains a given key.
 ///
 /// Returns `1` if the key is found.
@@ -37,7 +35,42 @@ use super::util::wasm_mem_read;
 /// * `ContractSection::Deploy`
 /// * `ContractSection::Metadata`
 /// * `ContractSection::Exec`
-pub(crate) fn db_contains_key(mut ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, ptr_len: u32) -> i64 {
+pub(crate) fn db_contains_key(ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, ptr_len: u32) -> i64 {
+    db_contains_key_internal(ctx, ptr, ptr_len, false)
+}
+
+/// Check if a tx-local database contains a given key.
+///
+/// Returns `1` if the key is found.
+/// Returns `0` if the key is not found and there are no errors.
+/// Otherwise, returns an error code.
+///
+/// ## Permissions
+/// * `ContractSection::Deploy`
+/// * `ContractSection::Metadata`
+/// * `ContractSection::Exec`
+pub(crate) fn db_contains_key_local(
+    ctx: FunctionEnvMut<Env>,
+    ptr: WasmPtr<u8>,
+    ptr_len: u32,
+) -> i64 {
+    db_contains_key_internal(ctx, ptr, ptr_len, true)
+}
+
+/// Internal `db_contains_key` function which branches to either on-chain or
+/// tx-local.
+///
+/// ## Permissions
+/// * `ContractSection::Deploy`
+/// * `ContractSection::Metadata`
+/// * `ContractSection::Exec`
+fn db_contains_key_internal(
+    mut ctx: FunctionEnvMut<Env>,
+    ptr: WasmPtr<u8>,
+    ptr_len: u32,
+    local: bool,
+) -> i64 {
+    let lt = if local { "db_contains_key_local" } else { "db_contains_key" };
     let (env, mut store) = ctx.data_and_store_mut();
     let cid = env.contract_id;
 
@@ -46,22 +79,23 @@ pub(crate) fn db_contains_key(mut ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, pt
         acl_allow(env, &[ContractSection::Deploy, ContractSection::Metadata, ContractSection::Exec])
     {
         error!(
-            target: "runtime::db::db_contains_key",
-            "[WASM] [{cid}] db_contains_key(): Called in unauthorized section: {e}",
+            target: "runtime::db::{lt}",
+            "[WASM] [{cid}] {lt}(): Called in unauthorized section: {e}",
         );
         return darkfi_sdk::error::CALLER_ACCESS_DENIED
     }
 
-    // Subtract used gas. Reading is free.
+    // Subtract used gas.
+    // Reading is free.
     env.subtract_gas(&mut store, 1);
 
-    // Get the wasm memory reader
+    // Get the WASM memory reader
     let mut buf_reader = match wasm_mem_read(env, &store, ptr, ptr_len) {
         Ok(v) => v,
         Err(e) => {
             error!(
-                target: "runtime::db::db_contains_key",
-                "[WASM] [{cid}] db_contains_key(): Failed to read wasm memory: {e}",
+                target: "runtime::db::{lt}",
+                "[WASM] [{cid}] {lt}(): Failed to read wasm memory: {e}",
             );
             return darkfi_sdk::error::DB_CONTAINS_KEY_FAILED
         }
@@ -72,8 +106,8 @@ pub(crate) fn db_contains_key(mut ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, pt
         Ok(v) => v,
         Err(e) => {
             error!(
-                target: "runtime::db::db_contains_key",
-                "[WASM] [{cid}] db_contains_key(): Failed to decode DbHandle: {e}",
+                target: "runtime::db::{lt}",
+                "[WASM] [{cid}] {lt}(): Failed to decode DbHandle: {e}",
             );
             return darkfi_sdk::error::DB_CONTAINS_KEY_FAILED
         }
@@ -86,30 +120,30 @@ pub(crate) fn db_contains_key(mut ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, pt
         Ok(v) => v,
         Err(e) => {
             error!(
-                target: "runtime::db::db_contains_key",
-                "[WASM] [{cid}] db_contains_key(): Failed to decode key vec: {e}",
+                target: "runtime::db::{lt}",
+                "[WASM] [{cid}] {lt}(): Failed to decode key vec: {e}",
             );
             return darkfi_sdk::error::DB_CONTAINS_KEY_FAILED
         }
     };
 
-    // Make sure there are no trailing bytes in the buffer.
-    // This means we've used all data that was supplied.
+    // Make sure we've read the entire buffer
     if buf_reader.position() != ptr_len as u64 {
         error!(
-            target: "runtime::db::db_contains_key",
-            "[WASM] [{cid}] db_contains_key(): Trailing bytes in argument stream",
+            target: "runtime::db::{lt}",
+            "[WASM] [{cid}] {lt}(): Trailing bytes in argument stream",
         );
         return darkfi_sdk::error::DB_CONTAINS_KEY_FAILED
     }
 
-    let db_handles = env.db_handles.borrow();
+    // Fetch requested db handles
+    let db_handles = if local { env.local_db_handles.borrow() } else { env.db_handles.borrow() };
 
     // Ensure DbHandle index is within bounds
     if db_handles.len() <= db_handle_index {
         error!(
-            target: "runtime::db::db_contains_key",
-            "[WASM] [{cid}] db_contains_key(): Requested DbHandle that is out of bounds",
+            target: "runtime::db::{lt}",
+            "[WASM] [{cid}] {lt}(): Requested DbHandle out of bounds",
         );
         return darkfi_sdk::error::DB_CONTAINS_KEY_FAILED
     }
@@ -117,14 +151,39 @@ pub(crate) fn db_contains_key(mut ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, pt
     // Retrieve DbHandle using the index
     let db_handle = &db_handles[db_handle_index];
 
-    // Lookup key parameter in the database
+    // Lookup key parameter in the appropriate db
+    if local {
+        let db = env.tx_local.lock();
+        let Some(db_cid) = db.get(&db_handle.contract_id) else {
+            error!(
+                target: "runtime::db::{lt}",
+                "[WASM] [{cid}] {lt}(): Could not find db for {}",
+                db_handle.contract_id,
+            );
+            return darkfi_sdk::error::DB_CONTAINS_KEY_FAILED
+        };
+
+        let Some(tree) = db_cid.get(&db_handle.tree) else {
+            error!(
+                target: "runtime::db::{lt}",
+                "[WASM] [{cid}] {lt}(): Could not find db tree for {}",
+                db_handle.contract_id,
+            );
+            return darkfi_sdk::error::DB_CONTAINS_KEY_FAILED
+        };
+
+        // 0=false, 1=true. Convert bool to i64.
+        return i64::from(tree.contains_key(&key))
+    }
+
+    // On-chain db
     match env.blockchain.lock().unwrap().overlay.lock().unwrap().contains_key(&db_handle.tree, &key)
     {
         Ok(v) => i64::from(v), // <- 0=false, 1=true. Convert bool to i64.
         Err(e) => {
             error!(
-                target: "runtime::db::db_contains_key",
-                "[WASM] [{cid}] db_contains_key(): sled.tree.contains_key failed: {e}",
+                target: "runtime::db::{lt}",
+                "[WASM] [{cid}] {lt}(): sled.tree.contains_key failed: {e}",
             );
             darkfi_sdk::error::DB_CONTAINS_KEY_FAILED
         }

src/runtime/import/db/db_del.rs

@@ -22,42 +22,67 @@ use tracing::error;
 use wasmer::{FunctionEnvMut, WasmPtr};
 
 use crate::runtime::{
-    import::acl::acl_allow,
+    import::{acl::acl_allow, util::wasm_mem_read},
     vm_runtime::{ContractSection, Env},
 };
 
-use super::util::wasm_mem_read;
+/// Remove a key from the on-chain database.
+///
+/// Returns `SUCCESS` on success, otherwise returns an error value.
+///
+/// ## Permissions
+/// * `ContractSection::Deploy`
+/// * `ContractSection::Update`
+pub(crate) fn db_del(ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, ptr_len: u32) -> i64 {
+    db_del_internal(ctx, ptr, ptr_len, false)
+}
 
-/// Remove a key from an on-chain database.
+/// Remove a key from the tx-local database.
 ///
 /// Returns `SUCCESS` on success, otherwise returns an error value.
 ///
 /// ## Permissions
 /// * `ContractSection::Deploy`
 /// * `ContractSection::Update`
-pub(crate) fn db_del(mut ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, ptr_len: u32) -> i64 {
+pub(crate) fn db_del_local(ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, ptr_len: u32) -> i64 {
+    db_del_internal(ctx, ptr, ptr_len, true)
+}
+
+/// Internal `db_del` function which branches to either on-chain or tx-local.
+///
+/// ## Permissions
+/// * `ContractSection::Deploy`
+/// * `ContractSection::Update`
+fn db_del_internal(
+    mut ctx: FunctionEnvMut<Env>,
+    ptr: WasmPtr<u8>,
+    ptr_len: u32,
+    local: bool,
+) -> i64 {
+    let lt = if local { "db_del_local" } else { "db_del" };
     let (env, mut store) = ctx.data_and_store_mut();
     let cid = env.contract_id;
 
     // Enforce function ACL
     if let Err(e) = acl_allow(env, &[ContractSection::Deploy, ContractSection::Update]) {
         error!(
-            target: "runtime::db::db_del",
-            "[WASM] [{cid}] db_del(): Called in unauthorized section: {e}",
+            target: "runtime::db::{lt}",
+            "[WASM] [{cid}] {lt}(): Called in unauthorized section: {e}",
         );
         return darkfi_sdk::error::CALLER_ACCESS_DENIED
     }
 
-    // Subtract used gas. We make deletion free.
+    // Subtract used gas.
+    // We make deletion free.
     env.subtract_gas(&mut store, 1);
 
-    // Get the wasm memory reader
+    // Get the WASM memory reader
     let mut buf_reader = match wasm_mem_read(env, &store, ptr, ptr_len) {
         Ok(v) => v,
         Err(e) => {
             error!(
-                target: "runtime::db::db_del",
-                "[WASM] [{cid}] db_del(): Failed to read WASM memory: {e}",
+                target: "runtime::db::{lt}",
+                "[WASM] [{cid}] {lt}(): Failed to read WASM memory: {e}",
             );
             return darkfi_sdk::error::DB_DEL_FAILED
         }
@@ -68,8 +93,8 @@ pub(crate) fn db_del(mut ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, ptr_len: u3
         Ok(v) => v,
         Err(e) => {
             error!(
-                target: "runtime::db::db_del",
-                "[WASM] [{cid}] db_del(): Failed to decode DbHandle: {e}"
+                target: "runtime::db::{lt}",
+                "[WASM] [{cid}] {lt}(): Failed to decode DbHandle: {e}",
             );
             return darkfi_sdk::error::DB_DEL_FAILED
         }
@@ -81,8 +106,8 @@ pub(crate) fn db_del(mut ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, ptr_len: u3
     // We should disallow writing with this.
     if env.contract_section == ContractSection::Deploy && db_handle_index == 0 {
         error!(
-            target: "runtime::db::db_del",
-            "[WASM] [{cid}] db_del(): Tried to write to zkas db",
+            target: "runtime::db::{lt}",
+            "[WASM] [{cid}] {lt}(): Tried to write to zkas db",
         );
         return darkfi_sdk::error::CALLER_ACCESS_DENIED
     }
@@ -92,8 +117,8 @@ pub(crate) fn db_del(mut ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, ptr_len: u3
         Ok(v) => v,
         Err(e) => {
             error!(
-                target: "runtime::db::db_del",
-                "[WASM] [{cid}] db_del(): Failed to decode key vec: {e}",
+                target: "runtime::db::{lt}",
+                "[WASM] [{cid}] {lt}(): Failed to decode key Vec: {e}",
             );
             return darkfi_sdk::error::DB_DEL_FAILED
         }
@@ -102,41 +127,64 @@ pub(crate) fn db_del(mut ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, ptr_len: u3
     // Make sure we've read the entire buffer
     if buf_reader.position() != ptr_len as u64 {
         error!(
-            target: "runtime::db::db_del",
-            "[WASM] [{cid}] db_del(): Trailing bytes in argument stream",
+            target: "runtime::db::{lt}",
+            "[WASM] [{cid}] {lt}(): Trailing bytes in argument stream",
         );
         return darkfi_sdk::error::DB_DEL_FAILED
     }
 
-    let db_handles = env.db_handles.borrow();
+    // Fetch requested db handles
+    let db_handles = if local { env.local_db_handles.borrow() } else { env.db_handles.borrow() };
 
     // Check DbHandle index is within bounds
     if db_handles.len() <= db_handle_index {
         error!(
-            target: "runtime::db::db_del",
-            "[WASM] [{cid}] db_del(): Requested DbHandle that is out of bounds",
+            target: "runtime::db::{lt}",
+            "[WASM] [{cid}] {lt}(): Requested DbHandle out of bounds",
         );
         return darkfi_sdk::error::DB_DEL_FAILED
     }
 
-    // Retrive DbHandle using the index
+    // Retrieve DbHandle using the index
     let db_handle = &db_handles[db_handle_index];
 
-    // Validate that the DbHandle matches the contract ID
+    // Validate that the DbHandle matches the contract ID.
+    // We're not letting foreign contracts write to others' dbs.
     if db_handle.contract_id != cid {
         error!(
-            target: "runtime::db::db_del",
-            "[WASM] [{cid}] db_del(): Unauthorized to write to DbHandle",
+            target: "runtime::db::{lt}",
+            "[WASM] [{cid}] {lt}(): Unauthorized write to DbHandle",
         );
         return darkfi_sdk::error::CALLER_ACCESS_DENIED
     }
 
-    // Remove key-value pair from the database corresponding to this contract
-    if env.blockchain.lock().unwrap().overlay.lock().unwrap().remove(&db_handle.tree, &key).is_err()
+    // Delete from appropriate db
+    if local {
+        // Safe to unwrap here.
+        let mut db = env.tx_local.lock();
+        let db_cid = db.get_mut(&cid).unwrap();
+        let Some(tree) = db_cid.get_mut(&db_handle.tree) else {
+            error!(
+                target: "runtime::db::{lt}",
+                "[WASM] [{cid}] {lt}(): Could not remove key from tx-local tree",
+            );
+            return darkfi_sdk::error::DB_DEL_FAILED
+        };
+
+        tree.remove(&key);
+    } else if env
+        .blockchain
+        .lock()
+        .unwrap()
+        .overlay
+        .lock()
+        .unwrap()
+        .remove(&db_handle.tree, &key)
+        .is_err()
     {
         error!(
-            target: "runtime::db::db_del",
-            "[WASM] [{cid}] db_del(): Couldn't remove key from db_handle tree",
+            target: "runtime::db::{lt}",
+            "[WASM] [{cid}] {lt}(): Could not remove key from on-chain tree",
         );
         return darkfi_sdk::error::DB_DEL_FAILED
     }

src/runtime/import/db/db_get.rs

@@ -21,12 +21,10 @@ use tracing::{debug, error};
 use wasmer::{FunctionEnvMut, WasmPtr};
 
 use crate::runtime::{
-    import::acl::acl_allow,
+    import::{acl::acl_allow, util::wasm_mem_read},
     vm_runtime::{ContractSection, Env},
 };
 
-use super::util::wasm_mem_read;
-
 /// Reads a value by key from the on-chain key-value store.
 ///
 /// On success, returns the length of the `objects` Vector in the environment.
@@ -36,7 +34,36 @@ use super::util::wasm_mem_read;
 /// * `ContractSection::Deploy`
 /// * `ContractSection::Metadata`
 /// * `ContractSection::Exec`
-pub(crate) fn db_get(mut ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, ptr_len: u32) -> i64 {
+pub(crate) fn db_get(ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, ptr_len: u32) -> i64 {
+    db_get_internal(ctx, ptr, ptr_len, false)
+}
+
+/// Reads a value by key from the tx-local key-value store.
+///
+/// On success, returns the length of the `objects` Vector in the environment.
+/// Otherwise, returns an error code.
+///
+/// ## Permissions
+/// * `ContractSection::Deploy`
+/// * `ContractSection::Metadata`
+/// * `ContractSection::Exec`
+pub(crate) fn db_get_local(ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, ptr_len: u32) -> i64 {
+    db_get_internal(ctx, ptr, ptr_len, true)
+}
+
+/// Internal `db_get` function which branches to either on-chain or tx-local.
+///
+/// ## Permissions
+/// * `ContractSection::Deploy`
+/// * `ContractSection::Metadata`
+/// * `ContractSection::Exec`
+pub(crate) fn db_get_internal(
+    mut ctx: FunctionEnvMut<Env>,
+    ptr: WasmPtr<u8>,
+    ptr_len: u32,
+    local: bool,
+) -> i64 {
+    let lt = if local { "db_get_local" } else { "db_get" };
     let (env, mut store) = ctx.data_and_store_mut();
     let cid = env.contract_id;
 
@@ -45,22 +72,22 @@ pub(crate) fn db_get(mut ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, ptr_len: u3
         acl_allow(env, &[ContractSection::Deploy, ContractSection::Metadata, ContractSection::Exec])
     {
         error!(
-            target: "runtime::db::db_get",
-            "[WASM] [{cid}] db_get(): Called in unauthorized section: {e}",
+            target: "runtime::db::{lt}",
+            "[WASM] [{cid}] {lt}(): Called in unauthorized section: {e}",
         );
         return darkfi_sdk::error::CALLER_ACCESS_DENIED
     }
 
-    // Subtract used gas. Reading is free.
+    // Subtract used gas.
     env.subtract_gas(&mut store, 1);
 
-    // Get the wasm memory reader
+    // Get the WASM memory reader
     let mut buf_reader = match wasm_mem_read(env, &store, ptr, ptr_len) {
         Ok(v) => v,
         Err(e) => {
             error!(
-                target: "runtime::db::db_get",
-                "[WASM] [{cid}] db_get(): Failed to read wasm memory: {e}",
+                target: "runtime::db::{lt}",
+                "[WASM] [{cid}] {lt}(): Failed to read wasm memory: {e}",
             );
             return darkfi_sdk::error::DB_GET_FAILED
         }
@@ -71,8 +98,8 @@ pub(crate) fn db_get(mut ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, ptr_len: u3
         Ok(v) => v,
         Err(e) => {
             error!(
-                target: "runtime::db::db_get",
-                "[WASM] [{cid}] db_get(): Failed to decode DbHandle: {e}",
+                target: "runtime::db::{lt}",
+                "[WASM] [{cid}] {lt}(): Failed to decode DbHandle: {e}",
             );
             return darkfi_sdk::error::DB_GET_FAILED
         }
@@ -95,19 +122,20 @@ pub(crate) fn db_get(mut ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, ptr_len: u3
     // Make sure there are no trailing bytes in the buffer.
     if buf_reader.position() != ptr_len as u64 {
         error!(
-            target: "runtime::db::db_get",
-            "[WASM] [{cid}] db_get(): Trailing bytes in argument stream",
+            target: "runtime::db::{lt}",
+            "[WASM] [{cid}] {lt}(): Trailing bytes in argument stream",
         );
         return darkfi_sdk::error::DB_GET_FAILED
     }
 
-    let db_handles = env.db_handles.borrow();
+    // Fetch requested db handles
+    let db_handles = if local { env.local_db_handles.borrow() } else { env.db_handles.borrow() };
 
     // Ensure that the index is within bounds
     if db_handles.len() <= db_handle_index {
         error!(
-            target: "runtime::db::db_get",
-            "[WASM] [{cid}] db_get(): Requested DbHandle that is out of bounds",
+            target: "runtime::db::{lt}",
+            "[WASM] [{cid}] {lt}(): Requested DbHandle that is out of bounds",
         );
         return darkfi_sdk::error::DB_GET_FAILED
     }
@@ -116,24 +144,47 @@ pub(crate) fn db_get(mut ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, ptr_len: u3
     let db_handle = &db_handles[db_handle_index];
 
     // Retrieve data using the `key`
-    let ret =
+    let ret: Option<Vec<u8>> = if local {
+        // tx-local db
+        let db = env.tx_local.lock();
+        let Some(db_cid) = db.get(&db_handle.contract_id) else {
+            error!(
+                target: "runtime::db::{lt}",
+                "[WASM] [{cid}] {lt}(): Could not find db for {}",
+                db_handle.contract_id,
+            );
+            return darkfi_sdk::error::DB_GET_FAILED
+        };
+
+        let Some(tree) = db_cid.get(&db_handle.tree) else {
+            error!(
+                target: "runtime::db::{lt}",
+                "[WASM] [{cid}] {lt}(): Could not find db tree for {}",
+                db_handle.contract_id,
+            );
+            return darkfi_sdk::error::DB_GET_FAILED
+        };
+
+        tree.get(&key).cloned()
+    } else {
         match env.blockchain.lock().unwrap().overlay.lock().unwrap().get(&db_handle.tree, &key) {
-            Ok(v) => v,
+            Ok(v) => v.map(|iv| iv.to_vec()),
             Err(e) => {
                 error!(
-                    target: "runtime::db::db_get",
-                    "[WASM] [{cid}] db_get(): Internal error getting from tree: {e}",
+                    target: "runtime::db::{lt}",
+                    "[WASM] [{cid}] {lt}(): Internal error getting from tree: {e}",
                 );
                 return darkfi_sdk::error::DB_GET_FAILED
             }
-        };
+        }
+    };
     drop(db_handles);
 
     // Return special error if the data is empty
     let Some(return_data) = ret else {
         debug!(
-            target: "runtime::db::db_get",
-            "[WASM] [{cid}] db_get(): Return data is empty",
+            target: "runtime::db::{lt}",
+            "[WASM] [{cid}] {lt}(): Return data is empty",
         );
         return darkfi_sdk::error::DB_GET_EMPTY
     };

src/runtime/import/db/db_init.rs

@@ -24,11 +24,11 @@ use tracing::error;
 use wasmer::{FunctionEnvMut, WasmPtr};
 
 use crate::runtime::{
-    import::acl::acl_allow,
+    import::{acl::acl_allow, util::wasm_mem_read},
     vm_runtime::{ContractSection, Env},
 };
 
-use super::{util::wasm_mem_read, DbHandle};
+use super::DbHandle;
 
 /// Create a new on-chain database instance for the calling contract.
 /// When created, push it to the list of db_handles.

src/runtime/import/db/db_lookup.rs

@@ -24,11 +24,11 @@ use tracing::error;
 use wasmer::{FunctionEnvMut, WasmPtr};
 
 use crate::runtime::{
-    import::acl::acl_allow,
+    import::{acl::acl_allow, util::wasm_mem_read},
     vm_runtime::{ContractSection, Env},
 };
 
-use super::{util::wasm_mem_read, DbHandle};
+use super::DbHandle;
 
 /// Lookup an on-chain database handle from its name.
 /// If it exists, push it to the list of db_handles.
@@ -43,7 +43,45 @@ use super::{util::wasm_mem_read, DbHandle};
 /// * `ContractSection::Metadata`
 /// * `ContractSection::Exec`
 /// * `ContractSection::Update`
-pub(crate) fn db_lookup(mut ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, ptr_len: u32) -> i64 {
+pub(crate) fn db_lookup(ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, ptr_len: u32) -> i64 {
+    db_lookup_internal(ctx, ptr, ptr_len, false)
+}
+
+/// Lookup a tx-local database handle from its name.
+/// Unlike the on-chain version, this will also initialize the database
+/// in-memory if it does not exist and the caller is allowed to write.
+/// Then it will push it to the list of transaction-local db_handles.
+///
+/// Returns the index of the DbHandle in the local_db_handles Vector on success.
+/// Otherwise, returns an error value.
+///
+/// This function can be called from any [`ContractSection`].
+///
+/// ## Permissions
+/// * `ContractSection::Deploy`
+/// * `ContractSection::Metadata`
+/// * `ContractSection::Exec`
+/// * `ContractSection::Update`
+pub(crate) fn db_lookup_local(ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, ptr_len: u32) -> i64 {
+    db_lookup_internal(ctx, ptr, ptr_len, true)
+}
+
+/// Internal `db_lookup` function which branches to either on-chain or
+/// tx-local.
+///
+///
+/// ## Permissions
+/// * `ContractSection::Deploy`
+/// * `ContractSection::Metadata`
+/// * `ContractSection::Exec`
+/// * `ContractSection::Update`
+fn db_lookup_internal(
+    mut ctx: FunctionEnvMut<Env>,
+    ptr: WasmPtr<u8>,
+    ptr_len: u32,
+    local: bool,
+) -> i64 {
+    let lt = if local { "db_lookup_local" } else { "db_lookup" };
     let (env, mut store) = ctx.data_and_store_mut();
     let cid = env.contract_id;
 
@@ -58,13 +96,14 @@ pub(crate) fn db_lookup(mut ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, ptr_len:
         ],
     ) {
         error!(
-            target: "runtime::db::db_lookup",
-            "[WASM] [{cid}] db_lookup() called in unauthorized section: {e}",
+            target: "runtime::db::{lt}",
+            "[WASM] [{cid}] {lt}() called in unauthorized section: {e}",
         );
         return darkfi_sdk::error::CALLER_ACCESS_DENIED
     }
 
-    // Subtract used gas. Opening an existing db should be free (i.e. 1 gas unit).
+    // Subtract used gas.
+    // Opening an existing db should be free (i.e. 1 gas unit).
     env.subtract_gas(&mut store, 1);
 
     // Get the wasm memory reader
@@ -72,23 +111,20 @@ pub(crate) fn db_lookup(mut ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, ptr_len:
         Ok(v) => v,
         Err(e) => {
             error!(
-                target: "runtime::db::db_lookup",
-                "[WASM] [{cid}] db_lookup(): Failed to read WASM memory: {e}",
+                target: "runtime::db::{lt}",
+                "[WASM] [{cid}] {lt}(): Failed to read WASM memory: {e}",
             );
             return darkfi_sdk::error::DB_LOOKUP_FAILED
         }
     };
 
-    // This takes lock of the blockchain overlay reference in the wasm env
-    let contracts = &env.blockchain.lock().unwrap().contracts;
-
     // Decode ContractId from memory
-    let cid: ContractId = match Decodable::decode(&mut buf_reader) {
+    let read_cid: ContractId = match Decodable::decode(&mut buf_reader) {
         Ok(v) => v,
         Err(e) => {
             error!(
-                target: "runtime::db::db_lookup",
-                "[WASM] [{cid}] db_lookup(): Failed to decode ContractId: {e}",
+                target: "runtime::db::{lt}",
+                "[WASM] [{cid}] {lt}(): Failed to decode ContractId: {e}",
             );
             return darkfi_sdk::error::DB_LOOKUP_FAILED
         }
@@ -99,8 +135,8 @@ pub(crate) fn db_lookup(mut ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, ptr_len:
         Ok(v) => v,
         Err(e) => {
             error!(
-                target: "runtime::db::db_lookup",
-                "[WASM] [{cid}] db_lookup(): Failed to decode db_name: {e}",
+                target: "runtime::db::{lt}",
+                "[WASM] [{cid}] {lt}(): Failed to decode db_name: {e}",
             );
             return darkfi_sdk::error::DB_LOOKUP_FAILED
         }
@@ -109,8 +145,8 @@ pub(crate) fn db_lookup(mut ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, ptr_len:
     // Make sure we've read the entire buffer
     if buf_reader.position() != ptr_len as u64 {
         error!(
-            target: "runtime::db::db_lookup",
-            "[WASM] [{cid}] db_lookup(), Trailing bytes in argument stream",
+            target: "runtime::db::{lt}",
+            "[WASM] [{cid}] {lt}(), Trailing bytes in argument stream",
         );
         return darkfi_sdk::error::DB_LOOKUP_FAILED
     }
@@ -118,21 +154,53 @@ pub(crate) fn db_lookup(mut ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, ptr_len:
     // We won't allow reading from the special zkas db or monotree db
     if [SMART_CONTRACT_ZKAS_DB_NAME, SMART_CONTRACT_MONOTREE_DB_NAME].contains(&db_name.as_str()) {
         error!(
-            target: "runtime::db::db_lookup",
-            "[WASM] [{cid}] db_lookup(): Attempted to lookup special db ({db_name})"
+            target: "runtime::db::{lt}",
+            "[WASM] [{cid}] {lt}(): Attempted to lookup special db ({db_name})"
         );
         return darkfi_sdk::error::CALLER_ACCESS_DENIED
     }
 
-    // Lookup contract state
-    let tree_handle = match contracts.lookup(&cid, &db_name) {
-        Ok(v) => v,
-        Err(_) => return darkfi_sdk::error::DB_LOOKUP_FAILED,
+    // Fetch the appropriate db
+    let tree_handle = if local {
+        let tree_handle = read_cid.hash_state_id(&db_name);
+
+        // Acquire the tx-local state
+        let mut db = env.tx_local.lock();
+
+        // If the caller is allowed to write, initialize the tx-local db
+        if read_cid == cid {
+            // Should be safe to unwrap here.
+            let db_cid = db.get_mut(&cid).unwrap();
+            db_cid.entry(tree_handle).or_default();
+        }
+
+        let Some(db_cid) = db.get(&read_cid) else {
+            // DB non-existent
+            return darkfi_sdk::error::DB_LOOKUP_FAILED
+        };
+
+        // Now check if the contract's db contains the db_name tree
+        if !db_cid.contains_key(&tree_handle) {
+            return darkfi_sdk::error::DB_LOOKUP_FAILED
+        }
+
+        // If it does, we can return the handle
+        tree_handle
+    } else {
+        // This takes lock of the blockchain overlay reference in the wasm env
+        let contracts = &env.blockchain.lock().unwrap().contracts;
+
+        // Lookup contract state
+        match contracts.lookup(&read_cid, &db_name) {
+            Ok(v) => v,
+            Err(_) => return darkfi_sdk::error::DB_LOOKUP_FAILED,
+        }
     };
 
     // Create the DbHandle
-    let db_handle = DbHandle::new(cid, tree_handle);
-    let mut db_handles = env.db_handles.borrow_mut();
+    let db_handle = DbHandle::new(read_cid, tree_handle);
+    let mut db_handles =
+        if local { env.local_db_handles.borrow_mut() } else { env.db_handles.borrow_mut() };
 
     // Make sure we don't duplicate the DbHandle in the vec
     if let Some(index) = db_handles.iter().position(|x| x == &db_handle) {
@@ -147,8 +215,8 @@ pub(crate) fn db_lookup(mut ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, ptr_len:
         }
         Err(_) => {
             error!(
-                target: "runtime::db::db_lookup",
-                "[WASM] [{cid}] db_lookup(): Too many open DbHandles",
+                target: "runtime::db::{lt}",
+                "[WASM] [{cid}] {lt}(): Too many open DbHandles",
             );
             darkfi_sdk::error::DB_LOOKUP_FAILED
         }

src/runtime/import/db/db_set.rs

@@ -22,12 +22,10 @@ use tracing::error;
 use wasmer::{FunctionEnvMut, WasmPtr};
 
 use crate::runtime::{
-    import::acl::acl_allow,
+    import::{acl::acl_allow, util::wasm_mem_read},
     vm_runtime::{ContractSection, Env},
 };
 
-use super::util::wasm_mem_read;
-
 /// Set a value in the on-chain database for the given DbHandle.
 ///
 /// * `ptr` must contain the DbHandle index and the key-value pair.
@@ -38,15 +36,44 @@ use super::util::wasm_mem_read;
 /// ## Permissions
 /// * `ContractSection::Deploy`
 /// * `ContractSection::Update`
-pub(crate) fn db_set(mut ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, ptr_len: u32) -> i64 {
+pub(crate) fn db_set(ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, ptr_len: u32) -> i64 {
+    db_set_internal(ctx, ptr, ptr_len, false)
+}
+
+/// Set a value in the tx-local database for the given DbHandle.
+///
+/// * `ptr` must contain the DbHandle index and the key-value pair.
+/// * The DbHandle must match the ContractId.
+///
+/// Returns `SUCCESS` on success, otherwise returns an error value.
+///
+/// ## Permissions
+/// * `ContractSection::Deploy`
+/// * `ContractSection::Update`
+pub(crate) fn db_set_local(ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, ptr_len: u32) -> i64 {
+    db_set_internal(ctx, ptr, ptr_len, true)
+}
+
+/// Internal `db_set` function which branches to either on-chain or tx-local.
+///
+/// ## Permissions
+/// * `ContractSection::Deploy`
+/// * `ContractSection::Update`
+fn db_set_internal(
+    mut ctx: FunctionEnvMut<Env>,
+    ptr: WasmPtr<u8>,
+    ptr_len: u32,
+    local: bool,
+) -> i64 {
+    let lt = if local { "db_set_local" } else { "db_set" };
     let (env, mut store) = ctx.data_and_store_mut();
     let cid = env.contract_id;
 
     // Enforce function ACL
     if let Err(e) = acl_allow(env, &[ContractSection::Deploy, ContractSection::Update]) {
         error!(
-            target: "runtime::db::db_set",
-            "[WASM] [{cid}] db_set(): Called in unauthorized section: {e}",
+            target: "runtime::db::{lt}",
+            "[WASM] [{cid}] {lt}(): Called in unauthorized section: {e}",
         );
         return darkfi_sdk::error::CALLER_ACCESS_DENIED
     }
@@ -61,8 +88,8 @@ pub(crate) fn db_set(mut ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, ptr_len: u3
         Ok(v) => v,
         Err(e) => {
             error!(
-                target: "runtime::db::db_set",
-                "[WASM] [{cid}] db_set(): Failed to read wasm memory: {e}",
+                target: "runtime::db::{lt}",
+                "[WASM] [{cid}] {lt}(): Failed to read wasm memory: {e}",
             );
             return darkfi_sdk::error::DB_SET_FAILED
         }
@@ -73,8 +100,8 @@ pub(crate) fn db_set(mut ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, ptr_len: u3
         Ok(v) => v,
         Err(e) => {
             error!(
-                target: "runtime::db::db_set",
-                "[WASM] [{cid}] db_set(): Failed to decode DbHandle: {e}",
+                target: "runtime::db::{lt}",
+                "[WASM] [{cid}] {lt}(): Failed to decode DbHandle: {e}",
             );
             return darkfi_sdk::error::DB_SET_FAILED
         }
@@ -86,8 +113,8 @@ pub(crate) fn db_set(mut ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, ptr_len: u3
     // We should disallow writing with this.
     if env.contract_section == ContractSection::Deploy && db_handle_index == 0 {
         error!(
-            target: "runtime::db::db_set",
-            "[WASM] [{cid}] db_set(): Tried to write to zkas db",
+            target: "runtime::db::{lt}",
+            "[WASM] [{cid}] {lt}(): Tried to write to zkas db",
         );
         return darkfi_sdk::error::CALLER_ACCESS_DENIED
     }
@@ -97,8 +124,8 @@ pub(crate) fn db_set(mut ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, ptr_len: u3
         Ok(v) => v,
         Err(e) => {
             error!(
-                target: "runtime::db::db_set",
-                "[WASM] [{cid}] db_set(): Failed to decode key vec: {e}",
+                target: "runtime::db::{lt}",
+                "[WASM] [{cid}] {lt}(): Failed to decode key vec: {e}",
             );
             return darkfi_sdk::error::DB_SET_FAILED
         }
@@ -108,8 +135,8 @@ pub(crate) fn db_set(mut ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, ptr_len: u3
         Ok(v) => v,
         Err(e) => {
             error!(
-                target: "runtime::db::db_set",
-                "[WASM] [{cid}] db_set(): Failed to decode value vec: {e}",
+                target: "runtime::db::{lt}",
+                "[WASM] [{cid}] {lt}(): Failed to decode value vec: {e}",
             );
             return darkfi_sdk::error::DB_SET_FAILED
         }
@@ -118,19 +145,20 @@ pub(crate) fn db_set(mut ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, ptr_len: u3
     // Make sure we've read the entire buffer
     if buf_reader.position() != ptr_len as u64 {
         error!(
-            target: "runtime::db::db_set",
-            "[WASM] [{cid}] db_set(): Trailing bytes in argument stream",
+            target: "runtime::db::{lt}",
+            "[WASM] [{cid}] {lt}(): Trailing bytes in argument stream",
         );
         return darkfi_sdk::error::DB_SET_FAILED
     }
 
-    let db_handles = env.db_handles.borrow();
+    // Fetch requested db handles
+    let db_handles = if local { env.local_db_handles.borrow() } else { env.db_handles.borrow() };
 
     // Check DbHandle index is within bounds
     if db_handles.len() <= db_handle_index {
         error!(
-            target: "runtime::db::db_set",
-            "[WASM] [{cid}] db_set(): Requested DbHandle that is out of bounds",
+            target: "runtime::db::{lt}",
+            "[WASM] [{cid}] {lt}(): Requested DbHandle that is out of bounds",
         );
         return darkfi_sdk::error::DB_SET_FAILED
     }
@@ -141,14 +169,27 @@ pub(crate) fn db_set(mut ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, ptr_len: u3
     // Validate that the DbHandle matches the contract ID
     if db_handle.contract_id != env.contract_id {
         error!(
-            target: "runtime::db::db_set",
-            "[WASM] [{cid}] db_set(): Unauthorized to write to DbHandle",
+            target: "runtime::db::{lt}",
+            "[WASM] [{cid}] {lt}(): Unauthorized to write to DbHandle",
         );
         return darkfi_sdk::error::CALLER_ACCESS_DENIED
     }
 
     // Insert key-value pair into the database corresponding to this contract
-    if env
+    if local {
+        // Safe to unwrap here.
+        let mut db = env.tx_local.lock();
+        let db_cid = db.get_mut(&cid).unwrap();
+        let Some(tree) = db_cid.get_mut(&db_handle.tree) else {
+            error!(
+                target: "runtime::db::{lt}",
+                "[WASM] [{cid}] {lt}(): Could not insert to tx-local tree",
+            );
+            return darkfi_sdk::error::DB_SET_FAILED
+        };
+
+        tree.insert(key, value);
+    } else if env
         .blockchain
         .lock()
         .unwrap()
@@ -159,8 +200,8 @@ pub(crate) fn db_set(mut ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, ptr_len: u3
         .is_err()
     {
         error!(
-            target: "runtime::db::db_set",
-            "[WASM] [{cid}] db_set(): Couldn't insert to db_handle tree",
+            target: "runtime::db::{lt}",
+            "[WASM] [{cid}] {lt}(): Couldn't insert to on-chain tree",
         );
         return darkfi_sdk::error::DB_SET_FAILED
     }

src/runtime/import/db/mod.rs

@@ -35,21 +35,19 @@ pub(crate) mod db_init;
 pub(crate) use db_init::db_init;
 
 pub(crate) mod db_lookup;
-pub(crate) use db_lookup::db_lookup;
+pub(crate) use db_lookup::{db_lookup, db_lookup_local};
 
 pub(crate) mod db_set;
-pub(crate) use db_set::db_set;
+pub(crate) use db_set::{db_set, db_set_local};
 
 pub(crate) mod db_del;
-pub(crate) use db_del::db_del;
+pub(crate) use db_del::{db_del, db_del_local};
 
 pub(crate) mod db_get;
-pub(crate) use db_get::db_get;
+pub(crate) use db_get::{db_get, db_get_local};
 
 pub(crate) mod db_contains_key;
-pub(crate) use db_contains_key::db_contains_key;
+pub(crate) use db_contains_key::{db_contains_key, db_contains_key_local};
 
 pub(crate) mod zkas_db_set;
 pub(crate) use zkas_db_set::zkas_db_set;
-
-mod util;

src/runtime/import/db/util.rs

@@ -1,43 +0,0 @@
-/* This file is part of DarkFi (https://dark.fi)
- *
- * Copyright (C) 2020-2026 Dyne.org foundation
- *
- * This program is free software: you can redistribute it and/or modify
- * it under the terms of the GNU Affero General Public License as
- * published by the Free Software Foundation, either version 3 of the
- * License, or (at your option) any later version.
- *
- * This program is distributed in the hope that it will be useful,
- * but WITHOUT ANY WARRANTY; without even the implied warranty of
- * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
- * GNU Affero General Public License for more details.
- *
- * You should have received a copy of the GNU Affero General Public License
- * along with this program.  If not, see <https://www.gnu.org/licenses/>.
- */
-
-use std::io::Cursor;
-
-use wasmer::{StoreMut, WasmPtr};
-
-use crate::{runtime::vm_runtime::Env, Result};
-
-/// Create a mem slice of the WASM VM memory given a pointer and its length,
-/// and return a `Cursor` from which callers are able to read as a stream.
-pub fn wasm_mem_read(
-    env: &Env,
-    store: &StoreMut<'_>,
-    ptr: WasmPtr<u8>,
-    ptr_len: u32,
-) -> Result<Cursor<Vec<u8>>> {
-    let memory_view = env.memory_view(&store);
-    let mem_slice = ptr.slice(&memory_view, ptr_len)?;
-
-    // Allocate a buffer and copy all the data from the pointer
-    // into the buffer
-    let mut buf = vec![0u8; ptr_len as usize];
-    mem_slice.read_slice(&mut buf)?;
-
-    // Once the data is copied, we'll return a Cursor over it
-    Ok(Cursor::new(buf))
-}

src/runtime/import/db/zkas_db_set.rs

@@ -23,15 +23,13 @@ use wasmer::{FunctionEnvMut, WasmPtr};
 
 use crate::{
     runtime::{
-        import::acl::acl_allow,
+        import::{acl::acl_allow, util::wasm_mem_read},
         vm_runtime::{ContractSection, Env},
     },
     zk::{empty_witnesses, VerifyingKey, ZkCircuit},
     zkas::ZkBinary,
 };
 
-use super::util::wasm_mem_read;
-
 /// Given a zkas circuit, create a VerifyingKey and insert them both into
 /// the on-chain db.
 ///

src/runtime/import/merkle.rs

@@ -27,205 +27,247 @@ use darkfi_serial::{serialize, Decodable, Encodable, WriteExt};
 use tracing::{debug, error};
 use wasmer::{FunctionEnvMut, WasmPtr};
 
-use super::acl::acl_allow;
-use crate::runtime::vm_runtime::{ContractSection, Env};
-
-/// Adds data to merkle tree. The tree, database connection, and new data to add is
-/// read from `ptr` at offset specified by `len`.
-/// Returns `0` on success; otherwise, returns an error-code corresponding to a
-/// [`darkfi_sdk::error::ContractError`] (defined in the SDK).
-/// See also the method `merkle_add` in `sdk/src/merkle.rs`.
+use crate::runtime::{
+    import::{acl::acl_allow, util::wasm_mem_read},
+    vm_runtime::{ContractSection, Env},
+};
+
+/// Add data to an on-chain Merkle tree.
+///
+/// Expects:
+/// * `db_info`: Handle where the Merkle tree is stored
+/// * `db_roots`: Handle where all new Merkle roots are stored
+/// * `root_key`: Serialized key pointing to latest root in `db_info`
+/// * `tree_key`: Serialized key pointing to the Merkle tree in `db_info`
+/// * `coins`: Items we want to add to the Merkle tree
+///
+/// ## Permissions
+/// * `ContractSection::Update`
+pub(crate) fn merkle_add(ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, ptr_len: u32) -> i64 {
+    merkle_add_internal(ctx, ptr, ptr_len, false)
+}
+
+/// Add data to a tx-local Merkle tree.
 ///
-/// Permissions: update
-pub(crate) fn merkle_add(mut ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, len: u32) -> i64 {
+/// Expects:
+/// * `db_info`: Handle where the Merkle tree is stored
+/// * `db_roots`: Handle where all new Merkle roots are stored
+/// * `root_key`: Serialized key pointing to latest root in `db_info`
+/// * `tree_key`: Serialized key pointing to the Merkle tree in `db_info`
+/// * `coins`: Items we want to add to the Merkle tree
+///
+/// ## Permissions
+/// * `ContractSection::Update`
+pub(crate) fn merkle_add_local(ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, ptr_len: u32) -> i64 {
+    merkle_add_internal(ctx, ptr, ptr_len, true)
+}
+
+/// Internal function for `merkle_add` which branches to either on-chain
+/// or transaction-local.
+pub(crate) fn merkle_add_internal(
+    mut ctx: FunctionEnvMut<Env>,
+    ptr: WasmPtr<u8>,
+    ptr_len: u32,
+    local: bool,
+) -> i64 {
+    let lt = if local { "merkle_add_local" } else { "merkle_add" };
     let (env, mut store) = ctx.data_and_store_mut();
     let cid = env.contract_id;
 
     // Enforce function ACL
     if let Err(e) = acl_allow(env, &[ContractSection::Update]) {
         error!(
-            target: "runtime::merkle::merkle_add",
-            "[WASM] [{cid}] merkle_add(): Called in unauthorized section: {e}"
+            target: "runtime::merkle::{lt}",
+            "[WASM] [{cid}] {lt}(): Called in unauthorized section: {e}",
         );
         return darkfi_sdk::error::CALLER_ACCESS_DENIED
     }
 
-    // Subtract used gas.
-    // This makes calling the function which returns early have some (small) cost.
-    env.subtract_gas(&mut store, 1);
-
-    // Subtract written bytes as gas
-    env.subtract_gas(&mut store, 33 /* value_data.len() as u64 */);
+    // Subtract used gas. 1 for opcode, 33 for value_data.len().
+    env.subtract_gas(&mut store, 34);
 
-    let memory_view = env.memory_view(&store);
-    let Ok(mem_slice) = ptr.slice(&memory_view, len) else {
-        error!(
-            target: "runtime::merkle::merkle_add",
-            "[WASM] [{cid}] merkle_add(): Failed to make slice from ptr"
-        );
-        return darkfi_sdk::error::INTERNAL_ERROR
+    // Get the wasm memory reader
+    let mut buf_reader = match wasm_mem_read(env, &store, ptr, ptr_len) {
+        Ok(v) => v,
+        Err(e) => {
+            error!(
+                target: "runtime::merkle::{lt}",
+                "[WASM] [{cid}] {lt}(): Failed to read wasm memory: {e}",
+            );
+            return darkfi_sdk::error::INTERNAL_ERROR
+        }
     };
 
-    let mut buf = vec![0_u8; len as usize];
-    if let Err(e) = mem_slice.read_slice(&mut buf) {
-        error!(
-            target: "runtime::merkle::merkle_add",
-            "[WASM] [{cid}] merkle_add(): Failed to read from memory slice: {e}"
-        );
-        return darkfi_sdk::error::INTERNAL_ERROR
-    };
+    // The buffer should deserialize intto:
+    // - db_info (DbHandle)
+    // - db_roots (DbHandle)
+    // - root_key (Vec<u8>)
+    // - tree_key (Vec<u8>)
+    // - coins (Vec<MerkleNode>)
 
-    // The buffer should deserialize into:
-    // - db_info
-    // - db_roots
-    // - root_key (as Vec<u8>) (key being the name of the sled key in info_db where the latest root is)
-    // - tree_key (as Vec<u8>) (key being the name of the sled key in info_db where the Merkle tree is)
-    // - coins (as Vec<MerkleNode>) (the coins being added into the Merkle tree)
-    let mut buf_reader = Cursor::new(buf);
-    // FIXME: There's a type DbHandle=u32, but this should maybe be renamed
     let db_info_index: u32 = match Decodable::decode(&mut buf_reader) {
         Ok(v) => v,
         Err(e) => {
             error!(
-                target: "runtime::merkle::merkle_add",
-                "[WASM] [{cid}] merkle_add(): Failed to decode db_info DbHandle: {e}"
+                target: "runtime::merkle::{lt}",
+                "[WASM] [{cid}] {lt}(): Failed to decode db_info DbHandle: {e}",
             );
             return darkfi_sdk::error::INTERNAL_ERROR
         }
     };
-    let db_info_index = db_info_index as usize;
 
     let db_roots_index: u32 = match Decodable::decode(&mut buf_reader) {
         Ok(v) => v,
         Err(e) => {
             error!(
-                target: "runtime::merkle::merkle_add",
-                "[WASM] [{cid}] merkle_add(): Failed to decode db_roots DbHandle: {e}"
+                target: "runtime::merkle::{lt}",
+                "[WASM] [{cid}] {lt}(): Failed to decode db_roots DbHandle: {e}",
             );
             return darkfi_sdk::error::INTERNAL_ERROR
         }
     };
-    let db_roots_index = db_roots_index as usize;
 
-    let db_handles = env.db_handles.borrow();
+    // Fetch the required db handles
+    let db_info_index = db_info_index as usize;
+    let db_roots_index = db_roots_index as usize;
+    let db_handles = if local { env.local_db_handles.borrow() } else { env.db_handles.borrow() };
     let n_dbs = db_handles.len();
 
     if n_dbs <= db_info_index || n_dbs <= db_roots_index {
         error!(
-            target: "runtime::merkle::merkle_add",
-            "[WASM] [{cid}] merkle_add(): Requested DbHandle that is out of bounds"
+            target: "runtime::merkle::{lt}",
+            "[WASM] [{cid}] {lt}(): Requested DbHandle that is out of bounds",
         );
         return darkfi_sdk::error::INTERNAL_ERROR
     }
+
     let db_info = &db_handles[db_info_index];
     let db_roots = &db_handles[db_roots_index];
 
     // Make sure that the contract owns the dbs it wants to write to
-    if db_info.contract_id != env.contract_id || db_roots.contract_id != env.contract_id {
+    if db_info.contract_id != cid || db_roots.contract_id != cid {
         error!(
-            target: "runtime::merkle::merkle_add",
-            "[WASM] [{cid}] merkle_add(): Unauthorized to write to DbHandle"
+            target: "runtime::merkle::{lt}",
+            "[WASM] [{cid}] {lt}(): Unauthorized write to DbHandle",
         );
         return darkfi_sdk::error::CALLER_ACCESS_DENIED
     }
 
-    // This `key` represents the sled key in info where the latest root is
+    // This key represents the key in db_info where the latest root is
     let root_key: Vec<u8> = match Decodable::decode(&mut buf_reader) {
         Ok(v) => v,
         Err(e) => {
             error!(
-                target: "runtime::merkle::merkle_add",
-                "[WASM] [{cid}] merkle_add(): Failed to decode key vec: {e}"
+                target: "runtime::merkle::{lt}",
+                "[WASM] [{cid}] {lt}(): Failed to decode root_key Vec: {e}",
             );
             return darkfi_sdk::error::INTERNAL_ERROR
         }
     };
 
-    // This `key` represents the sled key in info where the Merkle tree is
+    // This key represents the key in db_info where the Merkle tree is
     let tree_key: Vec<u8> = match Decodable::decode(&mut buf_reader) {
         Ok(v) => v,
         Err(e) => {
             error!(
-                target: "runtime::merkle::merkle_add",
-                "[WASM] [{cid}] merkle_add(): Failed to decode key vec: {e}"
+                target: "runtime::merkle::{lt}",
+                "[WASM] [{cid}] {lt}(): Failed to decode tree_key Vec: {e}",
             );
             return darkfi_sdk::error::INTERNAL_ERROR
         }
     };
 
-    // This `coin` represents the leaf we're adding to the Merkle tree
+    // Coins represent the leaf(s) we're adding to the Merkle tree
     let coins: Vec<MerkleNode> = match Decodable::decode(&mut buf_reader) {
         Ok(v) => v,
         Err(e) => {
             error!(
-                target: "runtime::merkle::merkle_add",
-                "[WASM] [{cid}] merkle_add(): Failed to decode MerkleNode: {e}"
+                target: "runtime::merkle::{lt}",
+                "[WASM] [{cid}] {lt}(): Failed to decode Vec<MerkleNode>: {e}",
             );
             return darkfi_sdk::error::INTERNAL_ERROR
         }
     };
 
     // Make sure we've read the entire buffer
-    if buf_reader.position() != (len as u64) {
+    if buf_reader.position() != ptr_len as u64 {
         error!(
-            target: "runtime::merkle::merkle_add",
-            "[WASM] [{cid}] merkle_add(): Mismatch between given length, and cursor length"
+            target: "runtime::merkle::{lt}",
+            "[WASM] [{cid}] {lt}(): Trailing bytes in argument stream",
         );
         return darkfi_sdk::error::INTERNAL_ERROR
     }
 
-    // Locking should happen for the entire duration of this fn. This is unsafe otherwise.
-    let lock = env.blockchain.lock().unwrap();
-    let mut overlay = lock.overlay.lock().unwrap();
-    // Read the current tree
-    let ret = match overlay.get(&db_info.tree, &tree_key) {
-        Ok(v) => v,
-        Err(e) => {
+    // Even with tx-local, we will lock the blockchain db to make sure
+    // it does not change for any reason during this execution.
+    let blockchain = env.blockchain.lock().unwrap();
+    let mut overlay = blockchain.overlay.lock().unwrap();
+    let mut tx_local_db = env.tx_local.lock();
+
+    // Read the current Merkle tree.
+    let tree_bytes = if local {
+        let Some(db_cid) = tx_local_db.get(&db_info.contract_id) else {
+            error!(
+                target: "runtime::db::{lt}",
+                "[WASM] [{cid}] {lt}(): Could not find db for {}",
+                db_info.contract_id,
+            );
+            return darkfi_sdk::error::INTERNAL_ERROR
+        };
+
+        let Some(tree) = db_cid.get(&db_info.tree) else {
             error!(
-                target: "runtime::merkle::merkle_add",
-                "[WASM] [{cid}] merkle_add(): Internal error getting from tree: {e}"
+                target: "runtime::db::{lt}",
+                "[WASM] [{cid}] {lt}(): Could not find db tree for {}",
+                db_info.contract_id,
             );
             return darkfi_sdk::error::INTERNAL_ERROR
+        };
+
+        tree.get(&tree_key).cloned()
+    } else {
+        match overlay.get(&db_info.tree, &tree_key) {
+            Ok(v) => v.map(|iv| iv.to_vec()),
+            Err(e) => {
+                error!(
+                    target: "runtime::merkle::{lt}",
+                    "[WASM] [{cid}] {lt}(): Error getting from sled tree: {e}",
+                );
+                return darkfi_sdk::error::INTERNAL_ERROR
+            }
         }
     };
 
-    let Some(return_data) = ret else {
+    let Some(tree_bytes) = tree_bytes else {
         error!(
-            target: "runtime::merkle::merkle_add",
-            "[WASM] [{cid}] merkle_add(): Return data is empty"
+            target: "runtime::merkle::{lt}",
+            "[WASM] [{cid}] {lt}(): Merkle tree k/v is empty",
         );
         return darkfi_sdk::error::INTERNAL_ERROR
     };
 
-    debug!(
-        target: "runtime::merkle::merkle_add",
-        "Serialized tree: {} bytes",
-        return_data.len()
-    );
-    debug!(
-        target: "runtime::merkle::merkle_add",
-        "                 {}",
-        return_data.hex()
-    );
+    // Deserialize the tree
+    debug!(target: "runtime::merkle::{lt}", "Serialized tree: {} bytes", tree_bytes.len());
+    debug!(target: "runtime::merkle::{lt}", "{}", tree_bytes.hex());
 
-    let mut decoder = Cursor::new(&return_data);
+    let mut decoder = Cursor::new(&tree_bytes);
     let set_size: u32 = match Decodable::decode(&mut decoder) {
         Ok(v) => v,
         Err(e) => {
             error!(
-                target: "runtime::merkle::merkle_add",
-                "[WASM] [{cid}] merkle_add(): Unable to read set size: {e}"
+                target: "runtime::merkle::{lt}",
+                "[WASM] [{cid}] {lt}(): Unable to decode set size: {e}",
             );
             return darkfi_sdk::error::INTERNAL_ERROR
         }
     };
 
-    let mut tree: MerkleTree = match Decodable::decode(&mut decoder) {
+    let mut merkle_tree: MerkleTree = match Decodable::decode(&mut decoder) {
         Ok(v) => v,
         Err(e) => {
             error!(
-                target: "runtime::merkle::merkle_add",
-                "[WASM] [{cid}] merkle_add(): Unable to deserialize Merkle tree: {e}"
+                target: "runtime::merkle::{lt}",
+                "[WASM] [{cid}] {lt}(): Unable to deserialize Merkle tree: {e}",
             );
             return darkfi_sdk::error::INTERNAL_ERROR
         }
@@ -234,44 +276,52 @@ pub(crate) fn merkle_add(mut ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, len: u3
     // Here we add the new coins into the tree.
     let coins_len = coins.len();
     for coin in coins {
-        tree.append(coin);
+        merkle_tree.append(coin);
     }
 
     // And we serialize the tree back to bytes
-    let mut tree_data = Vec::new();
-    if tree_data.write_u32(set_size + coins_len as u32).is_err() ||
-        tree.encode(&mut tree_data).is_err()
+    let mut merkle_tree_data = vec![];
+    if merkle_tree_data.write_u32(set_size + coins_len as u32).is_err() ||
+        merkle_tree.encode(&mut merkle_tree_data).is_err()
     {
         error!(
-            target: "runtime::merkle::merkle_add",
-            "[WASM] [{cid}] merkle_add(): Couldn't reserialize modified tree"
+            target: "runtime::merkle::{lt}",
+            "[WASM] [{cid}] {lt}(): Could not serialize modified Merkle tree",
         );
         return darkfi_sdk::error::INTERNAL_ERROR
     }
 
-    // Apply changes to overlay
-    if overlay.insert(&db_info.tree, &tree_key, &tree_data).is_err() {
+    // Apply changes
+    if local {
+        // We unwrap here because we already know the databases exist
+        // from when we fetched the tree.
+        let db_cid = tx_local_db.get_mut(&db_info.contract_id).unwrap();
+        let tree = db_cid.get_mut(&db_info.tree).unwrap();
+        tree.insert(tree_key, merkle_tree_data);
+    } else if let Err(e) = overlay.insert(&db_info.tree, &tree_key, &merkle_tree_data) {
         error!(
-            target: "runtime::merkle::merkle_add",
-            "[WASM] [{cid}] merkle_add(): Couldn't insert to db_info tree"
+            target: "runtime::merkle::{lt}",
+            "[WASM] [{cid}] {lt}(): Could not insert tree to db_info: {e}",
         );
         return darkfi_sdk::error::INTERNAL_ERROR
     }
 
     // Here we add the Merkle root to our set of roots
-    // Since each update to the tree is atomic, we only need to add the last root.
-    let Some(latest_root) = tree.root(0) else {
+    // Since each update to the tree is atomic, we only need to add the last
+    // known root.
+    let Some(latest_root) = merkle_tree.root(0) else {
         error!(
-            target: "runtime::merkle::merkle_add",
-            "[WASM] [{cid}] merkle_add(): Unable to read the root of tree"
+            target: "runtime::merkle::{lt}",
+            "[WASM] [{cid}] {lt}(): Unable to read Merkle tree root",
         );
         return darkfi_sdk::error::INTERNAL_ERROR
     };
 
     debug!(
-        target: "runtime::merkle::merkle_add",
-        "[WASM] [{cid}] merkle_add(): Appending Merkle root to db: {latest_root:?}"
+        target: "runtime::merkle::{lt}",
+        "[WASM] [{cid}] {lt}(): Appending Merkle root to db: {latest_root:?}",
     );
+
     let latest_root_data = serialize(&latest_root);
     assert_eq!(latest_root_data.len(), 32);
 
@@ -280,35 +330,38 @@ pub(crate) fn merkle_add(mut ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, len: u3
     env.call_idx.encode(&mut value_data).expect("Unable to serialize call_idx");
     assert_eq!(value_data.len(), 32 + 1);
 
-    if overlay.insert(&db_roots.tree, &latest_root_data, &value_data).is_err() {
-        error!(
-            target: "runtime::merkle::merkle_add",
-            "[WASM] [{cid}] merkle_add(): Couldn't insert to db_roots tree"
-        );
-        return darkfi_sdk::error::INTERNAL_ERROR
-    }
+    if local {
+        // We unwrap here because we already know the databases exist
+        // from when we fetched the tree.
+        let db_cid = tx_local_db.get_mut(&db_info.contract_id).unwrap();
 
-    // Write a pointer to the latest known root
-    debug!(
-        target: "runtime::merkle::merkle_add",
-        "[WASM] [{cid}] merkle_add(): Replacing latest Merkle root pointer"
-    );
+        let info_tree = db_cid.get_mut(&db_info.tree).unwrap();
+        info_tree.insert(root_key, latest_root_data.clone());
 
-    if overlay.insert(&db_info.tree, &root_key, &latest_root_data).is_err() {
-        error!(
-            target: "runtime::merkle::merkle_add",
-            "[WASM] [{cid}] merkle_add(): Couldn't insert latest root to db_info tree"
-        );
-        return darkfi_sdk::error::INTERNAL_ERROR
+        let roots_tree = db_cid.get_mut(&db_roots.tree).unwrap();
+        roots_tree.insert(latest_root_data, value_data);
+    } else {
+        if let Err(e) = overlay.insert(&db_roots.tree, &latest_root_data, &value_data) {
+            error!(
+                target: "runtime::merkle::{lt}",
+                "[WASM] [{cid}] {lt}(): Could not insert to db_roots tree: {e}",
+            );
+            return darkfi_sdk::error::INTERNAL_ERROR
+        }
+
+        if let Err(e) = overlay.insert(&db_info.tree, &root_key, &latest_root_data) {
+            error!(
+                target: "runtime::merkle::{lt}",
+                "[WASM] [{cid}] {lt}(): Could not insert latest root to db_info: {e}",
+            );
+            return darkfi_sdk::error::INTERNAL_ERROR
+        }
     }
 
     // Subtract used gas.
-    // Here we count:
-    // * The size of the Merkle tree we deserialized from the db.
-    // * The size of the Merkle tree we serialized into the db.
-    // * The size of the new Merkle roots we wrote into the db.
+    drop(tx_local_db);
     drop(overlay);
-    drop(lock);
+    drop(blockchain);
     drop(db_handles);
     let spent_gas = coins_len * 32;
     env.subtract_gas(&mut store, spent_gas as u64);

src/runtime/import/util.rs

@@ -21,10 +21,13 @@ use std::io::Cursor;
 use darkfi_sdk::wasm;
 use darkfi_serial::Decodable;
 use tracing::{debug, error};
-use wasmer::{FunctionEnvMut, WasmPtr};
+use wasmer::{FunctionEnvMut, StoreMut, WasmPtr};
 
 use super::acl::acl_allow;
-use crate::runtime::vm_runtime::{ContractSection, Env};
+use crate::{
+    runtime::vm_runtime::{ContractSection, Env},
+    Result,
+};
 
 /// Host function for logging strings.
 pub(crate) fn drk_log(mut ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, len: u32) {
@@ -50,6 +53,26 @@ pub(crate) fn drk_log(mut ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, len: u32)
     }
 }
 
+/// Create a mem slice of the WASM VM memory given a pointer and its length,
+/// and return a `Cursor` from which callers are able to read as a stream.
+pub(crate) fn wasm_mem_read(
+    env: &Env,
+    store: &StoreMut<'_>,
+    ptr: WasmPtr<u8>,
+    ptr_len: u32,
+) -> Result<Cursor<Vec<u8>>> {
+    let memory_view = env.memory_view(&store);
+    let mem_slice = ptr.slice(&memory_view, ptr_len)?;
+
+    // Allocate a buffer and copy all the data from the pointer
+    // into the buffer
+    let mut buf = vec![0u8; ptr_len as usize];
+    mem_slice.read_slice(&mut buf)?;
+
+    // Once the data is copied, we'll return a Cursor over it
+    Ok(Cursor::new(buf))
+}
+
 /// Writes data to the `contract_return_data` field of [`Env`].
 /// The data will be read from `ptr` at a memory offset specified by `len`.
 ///

src/runtime/vm_runtime.rs

@@ -270,30 +270,60 @@ impl Runtime {
                     import::db::db_lookup,
                 ),
 
+                "db_lookup_local_" => Function::new_typed_with_env(
+                    &mut store,
+                    &ctx,
+                    import::db::db_lookup_local,
+                ),
+
                 "db_get_" => Function::new_typed_with_env(
                     &mut store,
                     &ctx,
                     import::db::db_get,
                 ),
 
+                "db_get_local_" => Function::new_typed_with_env(
+                    &mut store,
+                    &ctx,
+                    import::db::db_get_local,
+                ),
+
                 "db_contains_key_" => Function::new_typed_with_env(
                     &mut store,
                     &ctx,
                     import::db::db_contains_key,
                 ),
 
+                "db_contains_key_local_" => Function::new_typed_with_env(
+                    &mut store,
+                    &ctx,
+                    import::db::db_contains_key_local,
+                ),
+
                 "db_set_" => Function::new_typed_with_env(
                     &mut store,
                     &ctx,
                     import::db::db_set,
                 ),
 
+                "db_set_local_" => Function::new_typed_with_env(
+                    &mut store,
+                    &ctx,
+                    import::db::db_set_local,
+                ),
+
                 "db_del_" => Function::new_typed_with_env(
                     &mut store,
                     &ctx,
                     import::db::db_del,
                 ),
 
+                "db_del_local_" => Function::new_typed_with_env(
+                    &mut store,
+                    &ctx,
+                    import::db::db_del_local,
+                ),
+
                 "zkas_db_set_" => Function::new_typed_with_env(
                     &mut store,
                     &ctx,
@@ -318,6 +348,12 @@ impl Runtime {
                     import::merkle::merkle_add,
                 ),
 
+                "merkle_add_local_" => Function::new_typed_with_env(
+                    &mut store,
+                    &ctx,
+                    import::merkle::merkle_add_local,
+                ),
+
                 "sparse_merkle_insert_batch_" => Function::new_typed_with_env(
                     &mut store,
                     &ctx,