WIP runtime: use an overlay over Blockchain

This allows us to use a sled-overlay over the original Blockchain sled db, so we can validate transactions execution without actually writting to it. Tests pass, but its far from ready. Cleaning, some todos and erroneous txs handling still missing.

Cargo.lock

@@ -1157,6 +1157,7 @@ dependencies = [
  "serde_json",
  "simplelog",
  "sled",
+ "sled-overlay",
  "smol",
  "socket2",
  "sqlx",
@@ -3836,6 +3837,15 @@ dependencies = [
  "parking_lot 0.11.2",
 ]
 
+[[package]]
+name = "sled-overlay"
+version = "0.0.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "562d0d1dc5940eae7db86d162f4682db65cdcd8b838d8e05163a70774d52ea7c"
+dependencies = [
+ "sled",
+]
+
 [[package]]
 name = "slice-group-by"
 version = "0.3.0"

Cargo.toml

@@ -129,6 +129,7 @@ sqlx = {version = "0.6.2", features = ["runtime-async-std-rustls", "sqlite"], op
 
 # Blockchain store
 sled = {version = "0.34.7", optional = true}
+sled-overlay = {version = "0.0.3", optional = true}
 
 [dev-dependencies]
 clap = {version = "4.1.4", features = ["derive"]}
@@ -156,6 +157,7 @@ blockchain = [
     "lazy_static",
     "rand",
     "sled",
+    "sled-overlay",
     "sqlx",
     "url",
 

src/blockchain/contract_store.rs

@@ -15,6 +15,7 @@ r* This program is distributed in the hope that it will be useful,
  * 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 darkfi_sdk::crypto::ContractId;
@@ -22,6 +23,7 @@ use darkfi_serial::{deserialize, serialize};
 use log::{debug, error};
 
 use crate::{
+    blockchain::SledDbOverlayPtr,
     runtime::vm_runtime::SMART_CONTRACT_ZKAS_DB_NAME,
     zk::{VerifyingKey, ZkCircuit},
     zkas::ZkBinary,
@@ -57,11 +59,23 @@ impl WasmStore {
 
         Err(Error::WasmBincodeNotFound)
     }
+}
+
+/// Overlay structure over a [`WasmStore`] instance.
+pub struct WasmStoreOverlay(SledDbOverlayPtr);
+
+impl WasmStoreOverlay {
+    pub fn new(overlay: SledDbOverlayPtr) -> Result<Self> {
+        overlay.lock().unwrap().open_tree(SLED_BINCODE_TREE)?;
+        Ok(Self(overlay))
+    }
 
     /// Inserts or replaces the bincode for a given ContractId
     pub fn insert(&self, contract_id: ContractId, bincode: &[u8]) -> Result<()> {
-        if let Err(e) = self.0.insert(serialize(&contract_id), bincode) {
-            error!(target: "blockchain::contractstore", "Failed to insert bincode to WasmStore: {}", e);
+        if let Err(e) =
+            self.0.lock().unwrap().insert(SLED_BINCODE_TREE, &serialize(&contract_id), bincode)
+        {
+            error!(target: "blockchain::contractstoreoverlay", "Failed to insert bincode to WasmStore: {}", e);
             return Err(e.into())
         }
 
@@ -89,58 +103,6 @@ impl ContractStateStore {
         Ok(Self(tree))
     }
 
-    /// Try to initialize a new contract state. Contracts can create a number
-    /// of trees, separated by `tree_name`, which they can then use from the
-    /// smart contract API. `init()` will look into the main `ContractStateStore`
-    /// tree to check if the smart contract was already deployed, and if so
-    /// it will fetch a vector of these states that were initialized. If the
-    /// state was already found, this function will return an error, because
-    /// in this case the handle should be fetched using `lookup()`.
-    /// If the tree was not initialized previously, it will be appended to
-    /// the main `ContractStateStore` tree and a `sled::Tree` handle will be
-    /// returned.
-    pub fn init(
-        &self,
-        db: &sled::Db,
-        contract_id: &ContractId,
-        tree_name: &str,
-    ) -> Result<sled::Tree> {
-        debug!(target: "blockchain::contractstore", "Initializing state tree for {}:{}", contract_id, tree_name);
-
-        let contract_id_bytes = serialize(contract_id);
-        let ptr = contract_id.hash_state_id(tree_name);
-
-        // See if there are existing state trees.
-        // If not, just start with an empty vector.
-        let mut state_pointers: Vec<[u8; 32]> = if self.0.contains_key(&contract_id_bytes)? {
-            let bytes = self.0.get(&contract_id_bytes)?.unwrap();
-            deserialize(&bytes)?
-        } else {
-            vec![]
-        };
-
-        // If the db was never initialized, it should not be in here.
-        if state_pointers.contains(&ptr) {
-            return Err(Error::ContractAlreadyInitialized)
-        }
-
-        // Now we add it so it's marked as initialized
-        state_pointers.push(ptr);
-
-        // We do this as a batch so in case of not being able to open the tree
-        // we don't write that it's initialized.
-        let mut batch = sled::Batch::default();
-        batch.insert(contract_id_bytes, serialize(&state_pointers));
-
-        // We open the tree and return its handle
-        let tree = db.open_tree(ptr)?;
-
-        // On success, apply the batch
-        self.0.apply_batch(batch)?;
-
-        Ok(tree)
-    }
-
     /// Do a lookup of an existing contract state. In order to succeed, the
     /// state must have been previously initialized with `init()`. If the
     /// state has been found, a handle to it will be returned. Otherwise, we
@@ -241,3 +203,83 @@ impl ContractStateStore {
         Ok((zkbin, vk))
     }
 }
+
+/// Overlay structure over a [`ContractStateStore`] instance.
+pub struct ContractStateStoreOverlay(SledDbOverlayPtr);
+
+impl ContractStateStoreOverlay {
+    pub fn new(overlay: SledDbOverlayPtr) -> Result<Self> {
+        overlay.lock().unwrap().open_tree(SLED_CONTRACTS_TREE)?;
+        Ok(Self(overlay))
+    }
+
+    /// Try to initialize a new contract state. Contracts can create a number
+    /// of trees, separated by `tree_name`, which they can then use from the
+    /// smart contract API. `init()` will look into the main `ContractStateStoreOverlay`
+    /// tree to check if the smart contract was already deployed, and if so
+    /// it will fetch a vector of these states that were initialized. If the
+    /// state was already found, this function will return an error, because
+    /// in this case the handle should be fetched using `lookup()`.
+    /// If the tree was not initialized previously, it will be appended to
+    /// the main `ContractStateStoreOverlay` tree and a handle to it will be
+    /// returned.
+    pub fn init(&self, contract_id: &ContractId, tree_name: &str) -> Result<[u8; 32]> {
+        debug!(target: "blockchain::contractstoreoverlay", "Initializing state overlay tree for {}:{}", contract_id, tree_name);
+
+        let contract_id_bytes = serialize(contract_id);
+        let ptr = contract_id.hash_state_id(tree_name);
+        let mut lock = self.0.lock().unwrap();
+
+        // See if there are existing state trees.
+        // If not, just start with an empty vector.
+        let mut state_pointers: Vec<[u8; 32]> =
+            if lock.contains_key(SLED_CONTRACTS_TREE, &contract_id_bytes)? {
+                let bytes = lock.get(SLED_CONTRACTS_TREE, &contract_id_bytes)?.unwrap();
+                deserialize(&bytes)?
+            } else {
+                vec![]
+            };
+
+        // If the db was never initialized, it should not be in here.
+        if state_pointers.contains(&ptr) {
+            return Err(Error::ContractAlreadyInitialized)
+        }
+
+        // Now we add it so it's marked as initialized and create its tree.
+        state_pointers.push(ptr);
+        lock.insert(SLED_CONTRACTS_TREE, &contract_id_bytes, &serialize(&state_pointers))?;
+        lock.open_tree(&ptr)?;
+
+        Ok(ptr)
+    }
+
+    /// Do a lookup of an existing contract state. In order to succeed, the
+    /// state must have been previously initialized with `init()`. If the
+    /// state has been found, a handle to it will be returned. Otherwise, we
+    /// return an error.
+    pub fn lookup(&self, contract_id: &ContractId, tree_name: &str) -> Result<[u8; 32]> {
+        debug!(target: "blockchain::contractstoreoverlay", "Looking up state tree for {}:{}", contract_id, tree_name);
+
+        let contract_id_bytes = serialize(contract_id);
+        let ptr = contract_id.hash_state_id(tree_name);
+        let mut lock = self.0.lock().unwrap();
+
+        // A guard to make sure we went through init()
+        if !lock.contains_key(SLED_CONTRACTS_TREE, &contract_id_bytes)? {
+            return Err(Error::ContractNotFound(contract_id.to_string()))
+        }
+
+        let state_pointers = lock.get(SLED_CONTRACTS_TREE, &contract_id_bytes)?.unwrap();
+        let state_pointers: Vec<[u8; 32]> = deserialize(&state_pointers)?;
+
+        // We assume the tree has been created already, so it should be listed
+        // in this array. If not, that's an error.
+        if !state_pointers.contains(&ptr) {
+            return Err(Error::ContractStateNotFound)
+        }
+
+        // We open the tree and return its handle
+        lock.open_tree(&ptr)?;
+        Ok(ptr)
+    }
+}

src/blockchain/mod.rs

@@ -16,6 +16,8 @@
  * along with this program.  If not, see <https://www.gnu.org/licenses/>.
  */
 
+use std::sync::{Arc, Mutex};
+
 use log::debug;
 
 use crate::{
@@ -34,7 +36,9 @@ pub mod tx_store;
 pub use tx_store::{PendingTxStore, TxStore};
 
 pub mod contract_store;
-pub use contract_store::{ContractStateStore, WasmStore};
+pub use contract_store::{
+    ContractStateStore, ContractStateStoreOverlay, WasmStore, WasmStoreOverlay,
+};
 
 /// Structure holding all sled trees that define the concept of Blockchain.
 #[derive(Clone)]
@@ -221,3 +225,30 @@ impl Blockchain {
         Ok(!vec.is_empty())
     }
 }
+
+/// Atomic pointer to sled db overlay.
+pub type SledDbOverlayPtr = Arc<Mutex<sled_overlay::SledDbOverlay>>;
+
+/// Atomic pointer to blockchain overlay.
+pub type BlockchainOverlayPtr = Arc<Mutex<BlockchainOverlay>>;
+
+/// Overlay structure over a [`Blockchain`] instance.
+pub struct BlockchainOverlay {
+    /// Main [`sled_overlay::SledDbOverlay`] to the sled db connection
+    pub overlay: SledDbOverlayPtr,
+    /// Contract states overlay
+    pub contracts: ContractStateStoreOverlay,
+    /// Wasm bincodes overlay
+    pub wasm_bincode: WasmStoreOverlay,
+}
+
+impl BlockchainOverlay {
+    /// Instantiate a new `BlockchainOverlay` over the given [`Blockchain`] instance.
+    pub fn new(blockchain: &Blockchain) -> Result<BlockchainOverlayPtr> {
+        let overlay = Arc::new(Mutex::new(sled_overlay::SledDbOverlay::new(&blockchain.sled_db)));
+        let contracts = ContractStateStoreOverlay::new(overlay.clone())?;
+        let wasm_bincode = WasmStoreOverlay::new(overlay.clone())?;
+
+        Ok(Arc::new(Mutex::new(Self { overlay, contracts, wasm_bincode })))
+    }
+}

src/consensus/validator.rs

@@ -43,7 +43,7 @@ use super::{
 };
 
 use crate::{
-    blockchain::Blockchain,
+    blockchain::{Blockchain, BlockchainOverlay},
     rpc::jsonrpc::JsonNotification,
     runtime::vm_runtime::Runtime,
     system::{Subscriber, SubscriberPtr},
@@ -167,12 +167,14 @@ impl ValidatorState {
         ];
 
         info!(target: "consensus::validator", "Deploying native wasm contracts");
+        let blockchain_overlay = BlockchainOverlay::new(&blockchain)?;
         for nc in native_contracts {
             info!(target: "consensus::validator", "Deploying {} with ContractID {}", nc.0, nc.1);
-            let mut runtime = Runtime::new(&nc.2[..], blockchain.clone(), nc.1)?;
+            let mut runtime = Runtime::new(&nc.2[..], blockchain_overlay.clone(), nc.1)?;
             runtime.deploy(&nc.3)?;
             info!(target: "consensus::validator", "Successfully deployed {}", nc.0);
         }
+        blockchain_overlay.lock().unwrap().overlay.lock().unwrap().apply()?;
 
         info!(target: "consensus::validator", "Finished deployment of native wasm contracts");
         // -----END NATIVE WASM CONTRACTS-----
@@ -884,6 +886,7 @@ impl ValidatorState {
     //       5. (optionally) write
     pub async fn verify_transactions(&self, txs: &[Transaction], write: bool) -> Result<()> {
         info!(target: "consensus::validator", "Verifying {} transaction(s)", txs.len());
+        let blockchain_overlay = BlockchainOverlay::new(&self.blockchain)?;
 
         for tx in txs {
             let tx_hash = blake3::hash(&serialize(tx));
@@ -916,7 +919,8 @@ impl ValidatorState {
                 tx.calls.encode(&mut payload)?; // Actual call data
 
                 // Instantiate the wasm runtime
-                let mut runtime = Runtime::new(&wasm, self.blockchain.clone(), call.contract_id)?;
+                let mut runtime =
+                    Runtime::new(&wasm, blockchain_overlay.clone(), call.contract_id)?;
 
                 info!(target: "consensus::validator", "Executing \"metadata\" call");
                 let metadata = runtime.metadata(&payload)?;
@@ -1008,7 +1012,7 @@ impl ValidatorState {
                     let wasm = self.blockchain.wasm_bincode.get(call.contract_id)?;
 
                     let mut runtime =
-                        Runtime::new(&wasm, self.blockchain.clone(), call.contract_id)?;
+                        Runtime::new(&wasm, blockchain_overlay.clone(), call.contract_id)?;
 
                     info!(target: "consensus::validator", "Executing \"apply\" call");
                     // TODO: FIXME: This should be done in an atomic tx/batch
@@ -1022,6 +1026,11 @@ impl ValidatorState {
             info!(target: "consensus::validator", "Transaction {} verified successfully", tx_hash);
         }
 
+        if write {
+            // The beauty of using Arc<Mutex<Struct{Arc<Mutex<>>}>>
+            blockchain_overlay.lock().unwrap().overlay.lock().unwrap().apply()?;
+        }
+
         Ok(())
     }
 

src/contract/money/tests/txs_verification.rs

@@ -116,10 +116,10 @@ async fn txs_verification() -> Result<()> {
 
     // Now Alice can send a little bit of funds to Bob.
     // We can duplicate this transaction to simulate double spending.
-    let DUPLICATES = 1; // Change this number to 2 to double spend
+    let duplicates = 1; // Change this number to 2 to double spend
     let mut transactions = vec![];
     let mut txs_params = vec![];
-    for i in 0..DUPLICATES {
+    for i in 0..duplicates {
         info!(target: "money", "[Alice] ======================================================");
         info!(target: "money", "[Alice] Building Money::Transfer params for payment {i} to Bob");
         info!(target: "money", "[Alice] ======================================================");
@@ -186,8 +186,8 @@ async fn txs_verification() -> Result<()> {
         txs_params.push(alice2bob_params);
     }
     alice_owncoins = vec![];
-    assert_eq!(transactions.len(), DUPLICATES);
-    assert_eq!(txs_params.len(), DUPLICATES);
+    assert_eq!(transactions.len(), duplicates);
+    assert_eq!(txs_params.len(), duplicates);
 
     // Now we can try to execute the transactions sequentialy.
     // The first transaction will get applied, while the second one(duplicate) will fail.

src/runtime/import/db.rs

@@ -30,47 +30,24 @@ use log::{debug, error, info};
 use wasmer::{FunctionEnvMut, WasmPtr};
 
 use crate::{
-    runtime::vm_runtime::{ContractSection, Env, SMART_CONTRACT_ZKAS_DB_NAME},
+    runtime::{
+        import,
+        vm_runtime::{ContractSection, Env, SMART_CONTRACT_ZKAS_DB_NAME},
+    },
     zk::{empty_witnesses, VerifyingKey, ZkCircuit},
     zkas::ZkBinary,
-    Result,
 };
 
 /// Internal wasm runtime API for sled trees
 pub struct DbHandle {
     pub contract_id: ContractId,
-    tree: sled::Tree,
+    pub tree: [u8; 32],
 }
 
 impl DbHandle {
-    pub fn new(contract_id: ContractId, tree: sled::Tree) -> Self {
+    pub fn new(contract_id: ContractId, tree: [u8; 32]) -> Self {
         Self { contract_id, tree }
     }
-
-    pub fn tree(&self) -> sled::Tree {
-        self.tree.clone()
-    }
-
-    pub fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>> {
-        if let Some(v) = self.tree.get(key)? {
-            return Ok(Some(v.to_vec()))
-        };
-
-        Ok(None)
-    }
-
-    pub fn contains_key(&self, key: &[u8]) -> Result<bool> {
-        Ok(self.tree.contains_key(key)?)
-    }
-
-    pub fn apply_batch(&self, batch: sled::Batch) -> Result<()> {
-        Ok(self.tree.apply_batch(batch)?)
-    }
-
-    pub fn flush(&self) -> Result<()> {
-        let _ = self.tree.flush()?;
-        Ok(())
-    }
 }
 
 /// Only deploy() can call this. Creates a new database instance for this contract.
@@ -84,8 +61,7 @@ pub(crate) fn db_init(ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, len: u32) -> i
     }
 
     let memory_view = env.memory_view(&ctx);
-    let db = &env.blockchain.sled_db;
-    let contracts = &env.blockchain.contracts;
+    let contracts = &env.blockchain.lock().unwrap().contracts;
     let contract_id = &env.contract_id;
 
     let Ok(mem_slice) = ptr.slice(&memory_view, len) else {
@@ -134,7 +110,7 @@ pub(crate) fn db_init(ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, len: u32) -> i
         return CALLER_ACCESS_DENIED
     }
 
-    let tree_handle = match contracts.init(db, &cid, &db_name) {
+    let tree_handle = match contracts.init(&cid, &db_name) {
         Ok(v) => v,
         Err(e) => {
             error!(target: "runtime::db::db_init()", "Failed to init db: {}", e);
@@ -151,7 +127,7 @@ pub(crate) fn db_init(ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, len: u32) -> i
     let mut db_handles = env.db_handles.borrow_mut();
     let mut db_batches = env.db_batches.borrow_mut();
     db_handles.push(DbHandle::new(cid, tree_handle));
-    db_batches.push(sled::Batch::default());
+    db_batches.push(import::util::Batch::default());
     (db_handles.len() - 1) as i32
 }
 
@@ -174,8 +150,7 @@ pub(crate) fn db_lookup(ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, len: u32) ->
     }
 
     let memory_view = env.memory_view(&ctx);
-    let db = &env.blockchain.sled_db;
-    let contracts = &env.blockchain.contracts;
+    let contracts = &env.blockchain.lock().unwrap().contracts;
 
     let Ok(mem_slice) = ptr.slice(&memory_view, len) else {
         error!(target: "runtime::db::db_lookup()", "Failed to make slice from ptr");
@@ -218,7 +193,7 @@ pub(crate) fn db_lookup(ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, len: u32) ->
         return DB_LOOKUP_FAILED
     }*/
 
-    let tree_handle = match contracts.lookup(db, &cid, &db_name) {
+    let tree_handle = match contracts.lookup(&cid, &db_name) {
         Ok(v) => v,
         Err(e) => {
             error!(target: "runtime::db::db_lookup()", "Failed to lookup db: {}", e);
@@ -235,7 +210,7 @@ pub(crate) fn db_lookup(ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, len: u32) ->
     let mut db_handles = env.db_handles.borrow_mut();
     let mut db_batches = env.db_batches.borrow_mut();
     db_handles.push(DbHandle::new(cid, tree_handle));
-    db_batches.push(sled::Batch::default());
+    db_batches.push(import::util::Batch::default());
     (db_handles.len() - 1) as i32
 }
 
@@ -455,13 +430,14 @@ pub(crate) fn db_get(ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, len: u32) -> i6
     let handle_idx = db_handle;
     let db_handle = &db_handles[handle_idx];
 
-    let ret = match db_handle.get(&key) {
-        Ok(v) => v,
-        Err(e) => {
-            error!(target: "runtime::db::db_get()", "Internal error getting from tree: {}", e);
-            return DB_GET_FAILED.into()
-        }
-    };
+    let ret =
+        match env.blockchain.lock().unwrap().overlay.lock().unwrap().get(&db_handle.tree, &key) {
+            Ok(v) => v,
+            Err(e) => {
+                error!(target: "runtime::db::db_get()", "Internal error getting from tree: {}", e);
+                return DB_GET_FAILED.into()
+            }
+        };
 
     let Some(return_data) = ret else {
         debug!(target: "runtime::db::db_get()", "returned empty vec");
@@ -470,7 +446,7 @@ pub(crate) fn db_get(ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, len: u32) -> i6
 
     // Copy Vec<u8> to the VM
     let mut objects = env.objects.borrow_mut();
-    objects.push(return_data);
+    objects.push(return_data.to_vec());
     (objects.len() - 1) as i64
 }
 
@@ -537,7 +513,8 @@ pub(crate) fn db_contains_key(ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, len: u
     let handle_idx = db_handle;
     let db_handle = &db_handles[handle_idx];
 
-    match db_handle.contains_key(&key) {
+    match env.blockchain.lock().unwrap().overlay.lock().unwrap().contains_key(&db_handle.tree, &key)
+    {
         Ok(v) => i32::from(v), // <- 0=false, 1=true
         Err(e) => {
             error!(target: "runtime::db::db_contains_key()", "sled.tree.contains_key failed: {}", e);
@@ -600,7 +577,15 @@ pub(crate) fn zkas_db_set(ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, len: u32)
     // Check if there is existing bincode and compare it. Return DB_SUCCESS if
     // they're the same. The assumption should be that VerifyingKey was generated
     // already so we can skip things after this guard.
-    match db_handle.get(&serialize(&zkbin.namespace)) {
+    match env
+        .blockchain
+        .lock()
+        .unwrap()
+        .overlay
+        .lock()
+        .unwrap()
+        .get(&db_handle.tree, &serialize(&zkbin.namespace))
+    {
         Ok(v) => {
             if let Some(bytes) = v {
                 // We allow a panic here because this db should never be corrupted in this way.

src/runtime/import/merkle.rs

@@ -114,7 +114,15 @@ pub(crate) fn merkle_add(ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, len: u32) -
             // TODO: Ensure we've read the entire buffer above.
 
             // Read the current tree
-            let ret = match db_info.get(&key) {
+            let ret = match env
+                .blockchain
+                .lock()
+                .unwrap()
+                .overlay
+                .lock()
+                .unwrap()
+                .get(&db_info.tree, &key)
+            {
                 Ok(v) => v,
                 Err(e) => {
                     error!(target: "runtime::merkle", "Internal error getting from tree: {}", e);

src/runtime/import/util.rs

@@ -16,6 +16,10 @@
  * along with this program.  If not, see <https://www.gnu.org/licenses/>.
  */
 
+// NOTE: temporary imports
+use sled::IVec;
+use std::collections::BTreeMap as Map;
+
 use log::error;
 use wasmer::{FunctionEnvMut, WasmPtr};
 
@@ -148,3 +152,40 @@ pub(crate) fn get_object_size(ctx: FunctionEnvMut<Env>, idx: u32) -> i64 {
     let obj = &objects[idx as usize];
     obj.len() as i64
 }
+
+// TODO: This is a direct copy of [`sled::Batch`](late night adventures).
+// Options:
+//  1. Upstream a get_writes() function
+//  2. Make writes public to external crates in upstream
+//  3. Drop Batches usage since we can write directly to the overlay
+//  4. Upstream batches support to sled_overlay
+#[derive(Debug, Default, Clone, PartialEq, Eq)]
+pub struct Batch {
+    pub(crate) writes: Map<IVec, Option<IVec>>,
+}
+
+impl Batch {
+    /// Set a key to a new value
+    pub fn insert<K, V>(&mut self, key: K, value: V)
+    where
+        K: Into<IVec>,
+        V: Into<IVec>,
+    {
+        self.writes.insert(key.into(), Some(value.into()));
+    }
+
+    /// Remove a key
+    pub fn remove<K>(&mut self, key: K)
+    where
+        K: Into<IVec>,
+    {
+        self.writes.insert(key.into(), None);
+    }
+
+    /// Get a value if it is present in the `Batch`.
+    /// `Some(None)` means it's present as a deletion.
+    pub fn get<K: AsRef<[u8]>>(&self, k: K) -> Option<Option<&IVec>> {
+        let inner = self.writes.get(k.as_ref())?;
+        Some(inner.as_ref())
+    }
+}

src/runtime/vm_runtime.rs

@@ -24,7 +24,6 @@ use std::{
 use darkfi_sdk::{crypto::ContractId, entrypoint};
 use darkfi_serial::serialize;
 use log::{debug, error, info};
-use sled::{transaction::ConflictableTransactionError, Transactional};
 use wasmer::{
     imports, wasmparser::Operator, AsStoreRef, CompilerConfig, Function, FunctionEnv, Instance,
     Memory, MemoryView, Module, Pages, Store, Value, WASM_PAGE_SIZE,
@@ -36,7 +35,7 @@ use wasmer_middlewares::{
 };
 
 use super::{import, import::db::DbHandle, memory::MemoryManipulation};
-use crate::{blockchain::Blockchain, Error, Result};
+use crate::{blockchain::BlockchainOverlayPtr, Error, Result};
 
 /// Name of the wasm linear memory in our guest module
 const MEMORY: &str = "memory";
@@ -75,12 +74,12 @@ impl ContractSection {
 
 /// The wasm vm runtime instantiated for every smart contract that runs.
 pub struct Env {
-    /// Blockchain access
-    pub blockchain: Blockchain,
+    /// Blockchain overlay access
+    pub blockchain: BlockchainOverlayPtr,
     /// sled tree handles used with `db_*`
     pub db_handles: RefCell<Vec<DbHandle>>,
     /// sled tree batches, indexed the same as `db_handles`.
-    pub db_batches: RefCell<Vec<sled::Batch>>,
+    pub db_batches: RefCell<Vec<import::util::Batch>>,
     /// The contract ID being executed
     pub contract_id: ContractId,
     /// The compiled wasm bincode being executed,
@@ -123,7 +122,11 @@ pub struct Runtime {
 
 impl Runtime {
     /// Create a new wasm runtime instance that contains the given wasm module.
-    pub fn new(wasm_bytes: &[u8], blockchain: Blockchain, contract_id: ContractId) -> Result<Self> {
+    pub fn new(
+        wasm_bytes: &[u8],
+        blockchain: BlockchainOverlayPtr,
+        contract_id: ContractId,
+    ) -> Result<Self> {
         info!(target: "runtime::vm_runtime", "Instantiating a new runtime");
         // TODO: Add necessary operators
         // This function will be called for each `Operator` encountered during
@@ -346,31 +349,24 @@ impl Runtime {
 
             // We always want to have the zkas db as index 0 in db handles and batches when
             // deploying.
-            let db = &env_mut.blockchain.sled_db;
-
-            let zkas_tree_handle = match env_mut.blockchain.contracts.lookup(
-                db,
-                &env_mut.contract_id,
-                SMART_CONTRACT_ZKAS_DB_NAME,
-            ) {
-                Ok(v) => v,
-                Err(_) => {
-                    // FIXME: All this is deploy code is "vulnerable" and able to init a
-                    // tree regardless of execution success. We can easily delete the db
-                    // if execution fails though, and we should charge gas for db_init.
-                    // and perhaps also for the zkas database in this specific case.
-                    env_mut.blockchain.contracts.init(
-                        db,
-                        &env_mut.contract_id,
-                        SMART_CONTRACT_ZKAS_DB_NAME,
-                    )?
-                }
-            };
+            let contracts = &env_mut.blockchain.lock().unwrap().contracts;
+
+            let zkas_tree_handle =
+                match contracts.lookup(&env_mut.contract_id, SMART_CONTRACT_ZKAS_DB_NAME) {
+                    Ok(v) => v,
+                    Err(_) => {
+                        // FIXME: All this is deploy code is "vulnerable" and able to init a
+                        // tree regardless of execution success. We can easily delete the db
+                        // if execution fails though, and we should charge gas for db_init.
+                        // and perhaps also for the zkas database in this specific case.
+                        contracts.init(&env_mut.contract_id, SMART_CONTRACT_ZKAS_DB_NAME)?
+                    }
+                };
 
             let mut db_handles = env_mut.db_handles.borrow_mut();
             let mut db_batches = env_mut.db_batches.borrow_mut();
             db_handles.push(DbHandle::new(env_mut.contract_id, zkas_tree_handle));
-            db_batches.push(sled::Batch::default());
+            db_batches.push(import::util::Batch::default());
         }
 
         debug!(target: "runtime::vm_runtime", "[wasm-runtime] payload: {:?}", payload);
@@ -381,31 +377,31 @@ impl Runtime {
 
         // Update the wasm bincode in the WasmStore
         let env_mut = self.ctx.as_mut(&mut self.store);
-        env_mut.blockchain.wasm_bincode.insert(env_mut.contract_id, &env_mut.contract_bincode)?;
+        env_mut
+            .blockchain
+            .lock()
+            .unwrap()
+            .wasm_bincode
+            .insert(env_mut.contract_id, &env_mut.contract_bincode)?;
 
         Ok(())
     }
 
-    /// Execute an atomic sled transaction to write all batches
+    /// Apply all batches to the overlay
     fn write_batches(&mut self) -> Result<()> {
-        let mut dbs = vec![];
-        let mut batches = vec![];
         let env_mut = self.ctx.as_mut(&mut self.store);
+        let batches = env_mut.db_batches.borrow();
+        let blockchain = env_mut.blockchain.lock().unwrap();
+        let mut overlay = blockchain.overlay.lock().unwrap();
         for (idx, db) in env_mut.db_handles.get_mut().iter().enumerate() {
-            let batch = env_mut.db_batches.borrow()[idx].clone();
-            dbs.push(db.tree());
-            batches.push(batch);
-        }
-
-        dbs.transaction(|dbs| {
-            for (idx, db) in dbs.iter().enumerate() {
-                db.apply_batch(&batches[idx])?;
+            let tree_handle = &db.tree;
+            for (k, v) in &batches[idx].writes {
+                match v {
+                    Some(u) => overlay.insert(tree_handle, &k, &u)?,
+                    None => overlay.remove(tree_handle, &k)?,
+                };
             }
-
-            Ok::<(), ConflictableTransactionError<sled::Error>>(())
-        })?;
-
-        env_mut.blockchain.sled_db.flush()?;
+        }
 
         Ok(())
     }