blockchain/contract_store: unified wasm and states trees into a single store structure

src/blockchain/contract_store.rs

@@ -36,103 +36,47 @@ const SLED_BINCODE_TREE: &[u8] = b"_wasm_bincode";
 /// The hardcoded db name for the zkas circuits database tree
 pub const SMART_CONTRACT_ZKAS_DB_NAME: &str = "_zkas";
 
-/// The `WasmStore` is a `sled` tree that stores the wasm bincode for deployed
-/// contracts.
+/// The `ContractStore` is a structure representing all `sled` trees related
+/// to storing the blockchain's contracts information.
 #[derive(Clone)]
-pub struct WasmStore(sled::Tree);
-
-impl WasmStore {
-    /// Opens or creates a `WasmStore`. This tree holds the wasm bincode.
+pub struct ContractStore {
+    /// The `sled` tree storing the wasm bincode for deployed contracts.
     /// The layout looks like this:
     /// ```plaintext
     ///  tree: "_wasm_bincode"
     ///   key: ContractId
     /// value: Vec<u8>
-    pub fn new(db: &sled::Db) -> Result<Self> {
-        let tree = db.open_tree(SLED_BINCODE_TREE)?;
-        Ok(Self(tree))
-    }
-
-    /// Fetches the bincode for a given ContractId
-    /// Returns an error if the bincode is not found.
-    pub fn get(&self, contract_id: ContractId) -> Result<Vec<u8>> {
-        if let Some(bincode) = self.0.get(serialize(&contract_id))? {
-            return Ok(bincode.to_vec())
-        }
-
-        Err(Error::WasmBincodeNotFound)
-    }
-
-    /// Retrieve all wasm bincodes from the `WasmStore` in the form of a tuple
-    /// (`contract_id`, `bincode`).
-    /// Be careful as this will try to load everything in memory.
-    pub fn get_all(&self) -> Result<Vec<(ContractId, Vec<u8>)>> {
-        let mut bincodes = vec![];
-
-        for bincode in self.0.iter() {
-            let bincode = bincode.unwrap();
-            let contract_id = deserialize(&bincode.0)?;
-            bincodes.push((contract_id, bincode.1.to_vec()));
-        }
-
-        Ok(bincodes)
-    }
+    pub wasm: sled::Tree,
+    /// The `sled` tree storing the pointers to contracts' databases.
+    /// See the rustdoc for the impl functions for more info.
+    /// The layout looks like this:
+    /// ```plaintext
+    ///  tree: "_contracts"
+    ///   key: ContractId
+    /// value: Vec<blake3(ContractId || tree_name)>
+    /// ```
+    /// These values get mutated with `init()` and `remove()`.
+    pub state: sled::Tree,
 }
 
-/// 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.clone()))
+impl ContractStore {
+    /// Opens a new or existing `ContractStore` on the given sled database.
+    pub fn new(db: &sled::Db) -> Result<Self> {
+        let wasm = db.open_tree(SLED_BINCODE_TREE)?;
+        let state = db.open_tree(SLED_CONTRACTS_TREE)?;
+        Ok(Self { wasm, state })
     }
 
-    /// Fetches the bincode for a given ContractId
+    /// Fetches the bincode for a given ContractId from the store's wasm tree.
     /// Returns an error if the bincode is not found.
     pub fn get(&self, contract_id: ContractId) -> Result<Vec<u8>> {
-        if let Some(bincode) =
-            self.0.lock().unwrap().get(SLED_BINCODE_TREE, &serialize(&contract_id))?
-        {
+        if let Some(bincode) = self.wasm.get(serialize(&contract_id))? {
             return Ok(bincode.to_vec())
         }
 
         Err(Error::WasmBincodeNotFound)
     }
 
-    /// 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.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())
-        }
-
-        Ok(())
-    }
-}
-
-/// The `ContractStateStore` is a `sled` tree that stores pointers to contracts'
-/// databases. See the rustdoc for the impl functions for more info.
-#[derive(Clone)]
-pub struct ContractStateStore(sled::Tree);
-
-impl ContractStateStore {
-    /// Opens or creates a `ContractStateStore`. This main tree holds the links
-    /// of contracts' states.
-    /// The layout looks like this:
-    /// ```plaintext
-    ///  tree: "_contracts"
-    ///   key: ContractId
-    /// value: Vec<blake3(ContractId || tree_name)>
-    /// ```
-    /// These values get mutated with `init()` and `remove()`.
-    pub fn new(db: &sled::Db) -> Result<Self> {
-        let tree = db.open_tree(SLED_CONTRACTS_TREE)?;
-        Ok(Self(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
@@ -149,11 +93,11 @@ impl ContractStateStore {
         let ptr = contract_id.hash_state_id(tree_name);
 
         // A guard to make sure we went through init()
-        if !self.0.contains_key(&contract_id_bytes)? {
+        if !self.state.contains_key(&contract_id_bytes)? {
             return Err(Error::ContractNotFound(contract_id.to_string()))
         }
 
-        let state_pointers = self.0.get(&contract_id_bytes)?.unwrap();
+        let state_pointers = self.state.get(&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
@@ -181,11 +125,11 @@ impl ContractStateStore {
         let ptr = contract_id.hash_state_id(tree_name);
 
         // A guard to make sure we went through init()
-        if !self.0.contains_key(&contract_id_bytes)? {
+        if !self.state.contains_key(&contract_id_bytes)? {
             return Err(Error::ContractNotFound(contract_id.to_string()))
         }
 
-        let state_pointers = self.0.get(&contract_id_bytes)?.unwrap();
+        let state_pointers = self.state.get(&contract_id_bytes)?.unwrap();
         let mut state_pointers: Vec<[u8; 32]> = deserialize(&state_pointers)?;
 
         // We assume the tree has been created already, so it should be listed
@@ -196,7 +140,7 @@ impl ContractStateStore {
 
         // Remove the deleted tree from the state pointer set.
         state_pointers.retain(|x| *x != ptr);
-        self.0.insert(contract_id_bytes, serialize(&state_pointers))?;
+        self.state.insert(contract_id_bytes, serialize(&state_pointers))?;
 
         // Drop the deleted tree from the database
         db.drop_tree(ptr)?;
@@ -237,13 +181,28 @@ impl ContractStateStore {
         Ok((zkbin, vk))
     }
 
-    /// Retrieve all contract states from the `ContractStateStore` in the form of a tuple
-    /// (`contract_id`, `state_hashes`).
+    /// Retrieve all wasm bincodes from the store's wasm tree in the form
+    /// of a tuple (`contract_id`, `bincode`).
+    /// Be careful as this will try to load everything in memory.
+    pub fn get_all_wasm(&self) -> Result<Vec<(ContractId, Vec<u8>)>> {
+        let mut bincodes = vec![];
+
+        for bincode in self.wasm.iter() {
+            let bincode = bincode.unwrap();
+            let contract_id = deserialize(&bincode.0)?;
+            bincodes.push((contract_id, bincode.1.to_vec()));
+        }
+
+        Ok(bincodes)
+    }
+
+    /// Retrieve all contract states from the store's state tree in the
+    /// form of a tuple (`contract_id`, `state_hashes`).
     /// Be careful as this will try to load everything in memory.
-    pub fn get_all(&self) -> Result<Vec<(ContractId, Vec<blake3::Hash>)>> {
+    pub fn get_all_states(&self) -> Result<Vec<(ContractId, Vec<blake3::Hash>)>> {
         let mut contracts = vec![];
 
-        for contract in self.0.iter() {
+        for contract in self.state.iter() {
             contracts.push(parse_record(contract.unwrap())?);
         }
 
@@ -251,15 +210,41 @@ impl ContractStateStore {
     }
 }
 
-/// Overlay structure over a [`ContractStateStore`] instance.
-pub struct ContractStateStoreOverlay(SledDbOverlayPtr);
+/// Overlay structure over a [`ContractStore`] instance.
+pub struct ContractStoreOverlay(SledDbOverlayPtr);
 
-impl ContractStateStoreOverlay {
+impl ContractStoreOverlay {
     pub fn new(overlay: &SledDbOverlayPtr) -> Result<Self> {
+        overlay.lock().unwrap().open_tree(SLED_BINCODE_TREE)?;
         overlay.lock().unwrap().open_tree(SLED_CONTRACTS_TREE)?;
         Ok(Self(overlay.clone()))
     }
 
+    /// Fetches the bincode for a given ContractId from the overlay's wasm tree.
+    /// Returns an error if the bincode is not found.
+    pub fn get(&self, contract_id: ContractId) -> Result<Vec<u8>> {
+        if let Some(bincode) =
+            self.0.lock().unwrap().get(SLED_BINCODE_TREE, &serialize(&contract_id))?
+        {
+            return Ok(bincode.to_vec())
+        }
+
+        Err(Error::WasmBincodeNotFound)
+    }
+
+    /// Inserts or replaces the bincode for a given ContractId into the overlay's
+    /// wasm tree.
+    pub fn insert(&self, contract_id: ContractId, bincode: &[u8]) -> Result<()> {
+        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 Wasm tree: {}", e);
+            return Err(e.into())
+        }
+
+        Ok(())
+    }
+
     /// 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`

src/blockchain/mod.rs

@@ -42,9 +42,7 @@ pub use tx_store::{TxStore, TxStoreOverlay};
 
 /// Contracts and Wasm storage implementations
 pub mod contract_store;
-pub use contract_store::{
-    ContractStateStore, ContractStateStoreOverlay, WasmStore, WasmStoreOverlay,
-};
+pub use contract_store::{ContractStore, ContractStoreOverlay};
 
 /// Structure holding all sled trees that define the concept of Blockchain.
 #[derive(Clone)]
@@ -61,10 +59,8 @@ pub struct Blockchain {
     pub difficulties: BlockDifficultyStore,
     /// Transactions related sled trees
     pub transactions: TxStore,
-    /// Contract states
-    pub contracts: ContractStateStore,
-    /// Wasm bincodes
-    pub wasm_bincode: WasmStore,
+    /// Contracts related sled trees
+    pub contracts: ContractStore,
 }
 
 impl Blockchain {
@@ -75,8 +71,7 @@ impl Blockchain {
         let order = BlockOrderStore::new(db)?;
         let difficulties = BlockDifficultyStore::new(db)?;
         let transactions = TxStore::new(db)?;
-        let contracts = ContractStateStore::new(db)?;
-        let wasm_bincode = WasmStore::new(db)?;
+        let contracts = ContractStore::new(db)?;
 
         Ok(Self {
             sled_db: db.clone(),
@@ -86,7 +81,6 @@ impl Blockchain {
             difficulties,
             transactions,
             contracts,
-            wasm_bincode,
         })
     }
 
@@ -357,10 +351,8 @@ pub struct BlockchainOverlay {
     pub difficulties: BlockDifficultyStoreOverlay,
     /// Transactions overlay
     pub transactions: TxStoreOverlay,
-    /// Contract states overlay
-    pub contracts: ContractStateStoreOverlay,
-    /// Wasm bincodes overlay
-    pub wasm_bincode: WasmStoreOverlay,
+    /// Contract overlay
+    pub contracts: ContractStoreOverlay,
 }
 
 impl BlockchainOverlay {
@@ -372,8 +364,7 @@ impl BlockchainOverlay {
         let order = BlockOrderStoreOverlay::new(&overlay)?;
         let difficulties = BlockDifficultyStoreOverlay::new(&overlay)?;
         let transactions = TxStoreOverlay::new(&overlay)?;
-        let contracts = ContractStateStoreOverlay::new(&overlay)?;
-        let wasm_bincode = WasmStoreOverlay::new(&overlay)?;
+        let contracts = ContractStoreOverlay::new(&overlay)?;
 
         Ok(Arc::new(Mutex::new(Self {
             overlay,
@@ -383,7 +374,6 @@ impl BlockchainOverlay {
             difficulties,
             transactions,
             contracts,
-            wasm_bincode,
         })))
     }
 
@@ -517,8 +507,7 @@ impl BlockchainOverlay {
         let order = BlockOrderStoreOverlay::new(&overlay)?;
         let difficulties = BlockDifficultyStoreOverlay::new(&overlay)?;
         let transactions = TxStoreOverlay::new(&overlay)?;
-        let contracts = ContractStateStoreOverlay::new(&overlay)?;
-        let wasm_bincode = WasmStoreOverlay::new(&overlay)?;
+        let contracts = ContractStoreOverlay::new(&overlay)?;
 
         Ok(Arc::new(Mutex::new(Self {
             overlay,
@@ -528,7 +517,6 @@ impl BlockchainOverlay {
             difficulties,
             transactions,
             contracts,
-            wasm_bincode,
         })))
     }
 }

src/contract/test-harness/src/lib.rs

@@ -288,7 +288,7 @@ fn benchmark_wasm_calls(
 
     for (idx, call) in tx.calls.iter().enumerate() {
         let overlay = BlockchainOverlay::new(&validator.blockchain).expect("blockchain overlay");
-        let wasm = overlay.lock().unwrap().wasm_bincode.get(call.data.contract_id).unwrap();
+        let wasm = overlay.lock().unwrap().contracts.get(call.data.contract_id).unwrap();
         let mut runtime = Runtime::new(
             &wasm,
             overlay.clone(),

src/runtime/vm_runtime.rs

@@ -452,13 +452,13 @@ impl Runtime {
         //debug!(target: "runtime::vm_runtime", "[WASM] payload: {:?}", payload);
         let _ = self.call(ContractSection::Deploy, payload)?;
 
-        // Update the wasm bincode in the WasmStore if the deploy exec passed successfully.
+        // Update the wasm bincode in the ContractStore wasm tree if the deploy exec passed successfully.
         let env_mut = self.ctx.as_mut(&mut self.store);
         env_mut
             .blockchain
             .lock()
             .unwrap()
-            .wasm_bincode
+            .contracts
             .insert(env_mut.contract_id, &env_mut.contract_bincode)?;
 
         info!(target: "runtime::vm_runtime", "[WASM] Successfully deployed ContractID: {}", cid);

src/validator/verification.rs

@@ -280,7 +280,7 @@ pub async fn verify_producer_transaction(
     tx.calls.encode_async(&mut payload).await?; // Actual call data
 
     debug!(target: "validator::verification::verify_producer_transaction", "Instantiating WASM runtime");
-    let wasm = overlay.lock().unwrap().wasm_bincode.get(call.data.contract_id)?;
+    let wasm = overlay.lock().unwrap().contracts.get(call.data.contract_id)?;
 
     let mut runtime = Runtime::new(
         &wasm,
@@ -447,7 +447,7 @@ pub async fn verify_transaction(
         tx.calls.encode_async(&mut payload).await?; // Actual call data
 
         debug!(target: "validator::verification::verify_transaction", "Instantiating WASM runtime");
-        let wasm = overlay.lock().unwrap().wasm_bincode.get(call.data.contract_id)?;
+        let wasm = overlay.lock().unwrap().contracts.get(call.data.contract_id)?;
 
         let mut runtime = Runtime::new(
             &wasm,