WIP runtime: remove Batches logic, write directly to overlay

src/consensus/validator.rs

@@ -1034,7 +1034,7 @@ impl ValidatorState {
         Ok(())
     }
 
-    pub async fn verify_transactions2(&self, txs: &[Transaction], write: bool) -> Result<()> {
+    pub async fn verify_transactions2(&self, _txs: &[Transaction], _write: bool) -> Result<()> {
         /*
         let mut sled_overlay = SledDbOverlay::new();
 

src/runtime/import/db.rs

@@ -30,10 +30,7 @@ use log::{debug, error, info};
 use wasmer::{FunctionEnvMut, WasmPtr};
 
 use crate::{
-    runtime::{
-        import,
-        vm_runtime::{ContractSection, Env, SMART_CONTRACT_ZKAS_DB_NAME},
-    },
+    runtime::vm_runtime::{ContractSection, Env, SMART_CONTRACT_ZKAS_DB_NAME},
     zk::{empty_witnesses, VerifyingKey, ZkCircuit},
     zkas::ZkBinary,
 };
@@ -120,14 +117,8 @@ pub(crate) fn db_init(ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, len: u32) -> i
 
     // TODO: Make sure we don't duplicate the DbHandle in the vec.
     //       It should behave like an ordered set.
-    // In `lookup()` we also create a `sled::Batch`. This is done for
-    // some simplicity reasons, and also for possible future changes.
-    // However, we make sure that unauthorized writes are not available
-    // from other functions that interface with the databases.
     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(import::util::Batch::default());
     (db_handles.len() - 1) as i32
 }
 
@@ -203,14 +194,8 @@ pub(crate) fn db_lookup(ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, len: u32) ->
 
     // TODO: Make sure we don't duplicate the DbHandle in the vec.
     //       It should behave like an ordered set.
-    // In `lookup()` we also create a `sled::Batch`. This is done for
-    // some simplicity reasons, and also for possible future changes.
-    // However, we make sure that unauthorized writes are not available
-    // from other functions that interface with the databases.
     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(import::util::Batch::default());
     (db_handles.len() - 1) as i32
 }
 
@@ -274,23 +259,33 @@ pub(crate) fn db_set(ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, len: u32) -> i3
     }*/
 
     let db_handles = env.db_handles.borrow();
-    let mut db_batches = env.db_batches.borrow_mut();
 
-    if db_handles.len() <= db_handle || db_batches.len() <= db_handle {
+    if db_handles.len() <= db_handle {
         error!(target: "runtime::db::db_set()", "Requested DbHandle that is out of bounds");
         return DB_SET_FAILED
     }
 
     let handle_idx = db_handle;
     let db_handle = &db_handles[handle_idx];
-    let db_batch = &mut db_batches[handle_idx];
 
     if db_handle.contract_id != env.contract_id {
         error!(target: "runtime::db::db_set()", "Unauthorized to write to DbHandle");
         return CALLER_ACCESS_DENIED
     }
 
-    db_batch.insert(key, value);
+    if env
+        .blockchain
+        .lock()
+        .unwrap()
+        .overlay
+        .lock()
+        .unwrap()
+        .insert(&db_handle.tree, &key, &value)
+        .is_err()
+    {
+        error!(target: "runtime::db::db_set()", "Couldn't insert to db_handle tree");
+        return DB_SET_FAILED
+    }
 
     DB_SUCCESS
 }
@@ -347,23 +342,25 @@ pub(crate) fn db_del(ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, len: u32) -> i3
     }*/
 
     let db_handles = env.db_handles.borrow();
-    let mut db_batches = env.db_batches.borrow_mut();
 
-    if db_handles.len() <= db_handle || db_batches.len() <= db_handle {
+    if db_handles.len() <= db_handle {
         error!(target: "runtime::db::db_del()", "Requested DbHandle that is out of bounds");
         return DB_DEL_FAILED
     }
 
     let handle_idx = db_handle;
     let db_handle = &db_handles[handle_idx];
-    let db_batch = &mut db_batches[handle_idx];
 
     if db_handle.contract_id != env.contract_id {
         error!(target: "runtime::db::db_del()", "Unauthorized to write to DbHandle");
         return CALLER_ACCESS_DENIED
     }
 
-    db_batch.remove(key);
+    if env.blockchain.lock().unwrap().overlay.lock().unwrap().remove(&db_handle.tree, &key).is_err()
+    {
+        error!(target: "runtime::db::db_del()", "Couldn't remove key from db_handle tree");
+        return DB_DEL_FAILED
+    }
 
     DB_SUCCESS
 }
@@ -565,9 +562,7 @@ pub(crate) fn zkas_db_set(ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, len: u32)
 
     // Because of `Runtime::Deploy`, we should be sure that the zkas db is index zero.
     let db_handles = env.db_handles.borrow();
-    let mut db_batches = env.db_batches.borrow_mut();
     let db_handle = &db_handles[0];
-    let db_batch = &mut db_batches[0];
     // Redundant check
     if &db_handle.contract_id != contract_id {
         error!(target: "runtime::db::zkas_db_set()", "Internal error, zkas db at index 0 incorrect");
@@ -617,7 +612,19 @@ pub(crate) fn zkas_db_set(ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, len: u32)
 
     let key = serialize(&zkbin.namespace);
     let value = serialize(&(zkas_bincode, vk_buf));
-    db_batch.insert(key, value);
+    if env
+        .blockchain
+        .lock()
+        .unwrap()
+        .overlay
+        .lock()
+        .unwrap()
+        .insert(&db_handle.tree, &key, &value)
+        .is_err()
+    {
+        error!(target: "runtime::db::zkas_db_set()", "Couldn't insert to db_handle tree");
+        return DB_SET_FAILED
+    }
 
     DB_SUCCESS
 }

src/runtime/import/merkle.rs

@@ -73,11 +73,9 @@ pub(crate) fn merkle_add(ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, len: u32) -
             let db_info = db_info as usize;
             let db_roots = db_roots as usize;
             let db_handles = env.db_handles.borrow();
-            let mut db_batches = env.db_batches.borrow_mut();
             let n_dbs = db_handles.len();
-            let n_bat = db_batches.len();
 
-            if n_dbs <= db_info || n_bat <= db_info || n_dbs <= db_roots || n_bat <= db_roots {
+            if n_dbs <= db_info || n_dbs <= db_roots {
                 error!(target: "runtime::merkle", "Requested DbHandle that is out of bounds");
                 return -2
             }
@@ -184,12 +182,17 @@ pub(crate) fn merkle_add(ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, len: u32) -
                 error!(target: "runtime::merkle", "Couldn't reserialize modified tree");
                 return -2
             }
-            let db_info_batch = &mut db_batches[info_handle_idx];
-            db_info_batch.insert(key, tree_data);
+
+            // Apply changes to overlay
+            let lock = env.blockchain.lock().unwrap();
+            let mut overlay = lock.overlay.lock().unwrap();
+            if overlay.insert(&db_info.tree, &key, &tree_data).is_err() {
+                error!(target: "runtime::merkle", "Couldn't insert to db_info tree");
+                return -2
+            }
 
             // Here we add the Merkle root to our set of roots
             // TODO: We should probably make sure that this root isn't in the set
-            let db_roots_batch = &mut db_batches[roots_handle_idx];
             for root in new_roots.iter() {
                 // FIXME: Why were we writing the set size here?
                 //let root_index: Vec<u8> = serialize(&(set_size as u32));
@@ -198,8 +201,10 @@ pub(crate) fn merkle_add(ctx: FunctionEnvMut<Env>, ptr: WasmPtr<u8>, len: u32) -
                 let root_value: Vec<u8> = serialize(root);
                 // FIXME: This assert can be used to DoS nodes from contracts
                 assert_eq!(root_value.len(), 32);
-                //db_roots_batch.insert(root_index, root_value);
-                db_roots_batch.insert(root_value, &[]);
+                if overlay.insert(&db_roots.tree, &root_value, &[]).is_err() {
+                    error!(target: "runtime::merkle", "Couldn't insert to db_roots tree");
+                    return -2
+                }
             }
 
             0

src/runtime/import/util.rs

@@ -16,10 +16,6 @@
  * 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};
 
@@ -152,40 +148,3 @@ 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

@@ -76,10 +76,8 @@ impl ContractSection {
 pub struct Env {
     /// Blockchain overlay access
     pub blockchain: BlockchainOverlayPtr,
-    /// sled tree handles used with `db_*`
+    /// Overlay 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<import::util::Batch>>,
     /// The contract ID being executed
     pub contract_id: ContractId,
     /// The compiled wasm bincode being executed,
@@ -157,7 +155,6 @@ impl Runtime {
 
         // Initialize data
         let db_handles = RefCell::new(vec![]);
-        let db_batches = RefCell::new(vec![]);
         let logs = RefCell::new(vec![]);
 
         debug!(target: "runtime::vm_runtime", "Importing functions");
@@ -167,7 +164,6 @@ impl Runtime {
             Env {
                 blockchain,
                 db_handles,
-                db_batches,
                 contract_id,
                 contract_bincode: wasm_bytes.to_vec(),
                 contract_section: ContractSection::Null,
@@ -336,10 +332,10 @@ impl Runtime {
     /// The runtime will look for an `INITIALIZE` symbol in the wasm code, and execute
     /// it if found. Optionally, it is possible to pass in a payload for any kind of special
     /// instructions the developer wants to manage in the initialize function.
-    /// This process is supposed to set up the sled db trees for storing the smart contract
+    /// This process is supposed to set up the overlay trees for storing the smart contract
     /// state, and it can create, delete, modify, read, and write to databases it's allowed to.
-    /// The permissions for this are handled by the `ContractId` in the sled db API so we
-    /// assume that the contract is only able to do write operations on its own sled trees.
+    /// The permissions for this are handled by the `ContractId` in the overlay db API so we
+    /// assume that the contract is only able to do write operations on its own overlay trees.
     pub fn deploy(&mut self, payload: &[u8]) -> Result<()> {
         info!(target: "runtime::vm_runtime", "[wasm-runtime] Running deploy");
 
@@ -364,17 +360,12 @@ impl Runtime {
                 };
 
             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(import::util::Batch::default());
         }
 
         debug!(target: "runtime::vm_runtime", "[wasm-runtime] payload: {:?}", payload);
         let _ = self.call(ContractSection::Deploy, payload)?;
 
-        // If the above didn't fail, we write the batches.
-        self.write_batches()?;
-
         // Update the wasm bincode in the WasmStore
         let env_mut = self.ctx.as_mut(&mut self.store);
         env_mut
@@ -387,25 +378,6 @@ impl Runtime {
         Ok(())
     }
 
-    /// Apply all batches to the overlay
-    fn write_batches(&mut self) -> Result<()> {
-        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 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(())
-    }
-
     /// This funcion runs when someone wants to execute a smart contract.
     /// The runtime will look for an `ENTRYPOINT` symbol in the wasm code, and
     /// execute it if found. A payload is also passed as an instruction that can
@@ -416,7 +388,7 @@ impl Runtime {
     }
 
     /// This function runs after successful execution of `exec` and tries to
-    /// apply the state change to the sled databases.
+    /// apply the state change to the overlay databases.
     /// The runtime will lok for an `UPDATE` symbol in the wasm code, and execute
     /// it if found. The function does not take an arbitrary payload, but just takes
     /// a state update from `env` and passes it into the wasm runtime.
@@ -424,9 +396,6 @@ impl Runtime {
         debug!(target: "runtime::vm_runtime", "apply: {:?}", update);
         let _ = self.call(ContractSection::Update, update)?;
 
-        // If the above didn't fail, we write the batches.
-        self.write_batches()?;
-
         Ok(())
     }