sdk/monotree: Add a MonotreeStorageAdapter trait for sled-overlay support

Cargo.lock

@@ -1922,6 +1922,7 @@ dependencies = [
  "rand 0.8.5",
  "rand_core 0.6.4",
  "sha2",
+ "sled-overlay",
  "subtle",
  "thiserror 2.0.12",
 ]

src/blockchain/contract_store.rs

@@ -23,7 +23,7 @@ use darkfi_sdk::{
         ContractId, NATIVE_CONTRACT_IDS_BYTES, NATIVE_CONTRACT_ZKAS_DB_NAMES,
         SMART_CONTRACT_ZKAS_DB_NAME,
     },
-    monotree::Monotree,
+    monotree::{self, Monotree},
 };
 use darkfi_serial::{deserialize, serialize};
 use sled_overlay::{serial::parse_record, sled, SledDbOverlay};
@@ -260,10 +260,11 @@ impl ContractStore {
     /// checksums, along with the wasm bincodes checksum.
     ///
     /// Note: native contracts zkas tree and wasm bincodes are excluded.
-    pub fn get_state_monotree(&self, db: &sled::Db) -> Result<Monotree> {
+    pub fn get_state_monotree(&self, db: &sled::Db) -> Result<Monotree<monotree::MemoryDb>> {
         // Initialize the monotree
         let mut root = None;
-        let mut tree = Monotree::new();
+        let monotree_db = monotree::MemoryDb::new();
+        let mut tree = Monotree::new(monotree_db);
 
         // Iterate over current contracts states records
         // TODO: parallelize this with a threadpool
@@ -455,7 +456,7 @@ impl ContractStoreOverlay {
     /// Be carefull as this will open all states trees in the overlay.
     ///
     /// Note: native contracts zkas tree and wasm bincodes are excluded.
-    pub fn get_state_monotree(&self) -> Result<Monotree> {
+    pub fn get_state_monotree(&self) -> Result<Monotree<monotree::MemoryDb>> {
         let mut lock = self.0.lock().unwrap();
 
         // Grab all states pointers
@@ -473,7 +474,8 @@ impl ContractStoreOverlay {
 
         // Initialize the monotree
         let mut root = None;
-        let mut tree = Monotree::new();
+        let monotree_db = monotree::MemoryDb::new();
+        let mut tree = Monotree::new(monotree_db);
 
         // Iterate over contract states pointers
         // TODO: parallelize this with a threadpool
@@ -520,7 +522,7 @@ impl ContractStoreOverlay {
     /// Monotree(SMT).
     ///
     /// Note: native contracts zkas tree and wasm bincodes are excluded.
-    pub fn update_state_monotree(&self, tree: &mut Monotree) -> Result<()> {
+    pub fn update_state_monotree(&self, tree: &mut Monotree<monotree::MemoryDb>) -> Result<()> {
         let lock = self.0.lock().unwrap();
 
         // Iterate over overlay's caches

src/blockchain/mod.rs

@@ -18,7 +18,10 @@
 
 use std::sync::{Arc, Mutex};
 
-use darkfi_sdk::{monotree::Monotree, tx::TransactionHash};
+use darkfi_sdk::{
+    monotree::{self, Monotree},
+    tx::TransactionHash,
+};
 use sled_overlay::{sled, sled::Transactional};
 use tracing::debug;
 
@@ -403,7 +406,7 @@ impl Blockchain {
     /// checksums, along with the wasm bincodes checksum.
     ///
     /// Note: native contracts zkas tree and wasm bincodes are excluded.
-    pub fn get_state_monotree(&self) -> Result<Monotree> {
+    pub fn get_state_monotree(&self) -> Result<Monotree<monotree::MemoryDb>> {
         self.contracts.get_state_monotree(&self.sled_db)
     }
 }
@@ -608,7 +611,7 @@ impl BlockchainOverlay {
     /// during checksum computing.
     ///
     /// Note: native contracts zkas tree and wasm bincodes are excluded.
-    pub fn get_state_monotree(&self) -> Result<Monotree> {
+    pub fn get_state_monotree(&self) -> Result<Monotree<monotree::MemoryDb>> {
         self.full_clone()?.lock().unwrap().contracts.get_state_monotree()
     }
 }

src/sdk/Cargo.toml

@@ -45,6 +45,9 @@ lazy_static = "1.5.0"
 subtle = "2.6.1"
 hashbrown = "0.15.4"
 
+# Storage
+sled-overlay = "0.1.9"
+
 [dev-dependencies]
 halo2_proofs = {version = "0.3.1", features = ["dev-graph", "sanity-checks"]}
 halo2_gadgets = {version = "0.3.1", features = ["test-dependencies"]}

src/sdk/src/monotree/mod.rs

@@ -41,7 +41,7 @@ pub mod bits;
 pub mod node;
 
 pub mod tree;
-pub use tree::Monotree;
+pub use tree::{MemoryDb, Monotree, SledDb};
 
 pub mod utils;
 

src/sdk/src/monotree/tests.rs

@@ -18,7 +18,7 @@
  */
 
 use super::{
-    tree::verify_proof,
+    tree::{verify_proof, MemoryDb},
     utils::{random_hashes, shuffle},
     Hash, Monotree,
 };
@@ -29,7 +29,8 @@ fn monotree_test_insert_then_verify_values() {
     let values = random_hashes(100);
 
     let mut root = None;
-    let mut tree = Monotree::new();
+    let db = MemoryDb::new();
+    let mut tree = Monotree::new(db);
 
     for (i, (key, value)) in keys.iter().zip(values.iter()).enumerate() {
         root = tree.insert(root.as_ref(), key, value).unwrap();
@@ -49,7 +50,8 @@ fn monotree_test_insert_keys_then_gen_and_verify_proof() {
     let values = random_hashes(100);
 
     let mut root = None;
-    let mut tree = Monotree::new();
+    let db = MemoryDb::new();
+    let mut tree = Monotree::new(db);
 
     for (i, (key, value)) in keys.iter().zip(values.iter()).enumerate() {
         root = tree.insert(root.as_ref(), key, value).unwrap();
@@ -70,7 +72,8 @@ fn monotree_test_insert_keys_then_delete_keys_in_order() {
     let values = random_hashes(100);
 
     let mut root = None;
-    let mut tree = Monotree::new();
+    let db = MemoryDb::new();
+    let mut tree = Monotree::new(db);
 
     // pre-insertion for removal test
     root = tree.inserts(root.as_ref(), &keys, &values).unwrap();
@@ -102,7 +105,8 @@ fn monotree_test_insert_then_delete_keys_reverse() {
     let values = random_hashes(100);
 
     let mut root = None;
-    let mut tree = Monotree::new();
+    let db = MemoryDb::new();
+    let mut tree = Monotree::new(db);
 
     // pre-insertion for removal test
     root = tree.inserts(root.as_ref(), &keys, &values).unwrap();
@@ -134,7 +138,8 @@ fn monotree_test_insert_then_delete_keys_random() {
     let values = random_hashes(100);
 
     let mut root = None;
-    let mut tree = Monotree::new();
+    let db = MemoryDb::new();
+    let mut tree = Monotree::new(db);
 
     // pre-insertion for removal test
     root = tree.inserts(root.as_ref(), &keys, &values).unwrap();
@@ -171,10 +176,12 @@ fn monotree_test_deterministic_ordering() {
     let values = random_hashes(100);
 
     let mut root1 = None;
-    let mut tree1 = Monotree::new();
+    let db = MemoryDb::new();
+    let mut tree1 = Monotree::new(db);
 
     let mut root2 = None;
-    let mut tree2 = Monotree::new();
+    let db = MemoryDb::new();
+    let mut tree2 = Monotree::new(db);
 
     // Insert in normal order
     root1 = tree1.inserts(root1.as_ref(), &keys, &values).unwrap();

src/sdk/src/monotree/tree.rs

@@ -18,6 +18,7 @@
  */
 
 use hashbrown::{HashMap, HashSet};
+use sled_overlay::SledTreeOverlay;
 
 use super::{
     bits::Bits,
@@ -25,7 +26,7 @@ use super::{
     utils::{get_sorted_indices, slice_to_hash},
     Hash, Proof, HASH_LEN, ROOT_KEY,
 };
-use crate::GenericResult;
+use crate::{ContractError, GenericResult};
 
 #[derive(Clone, Debug)]
 pub(crate) struct MemCache {
@@ -57,11 +58,26 @@ impl MemCache {
         self.set.remove(key);
     }
 
-    pub(crate) fn delete(&mut self, key: &[u8]) {
+    pub(crate) fn del(&mut self, key: &[u8]) {
         self.set.insert(slice_to_hash(key));
     }
 }
 
+/// Trait for implementing Monotree's storage system
+pub trait MonotreeStorageAdapter {
+    /// Insert a Key/Value pair into the Monotree
+    fn put(&mut self, key: &Hash, value: Vec<u8>) -> GenericResult<()>;
+    /// Query the Monotree for a Key
+    fn get(&self, key: &[u8]) -> GenericResult<Option<Vec<u8>>>;
+    /// Delete an entry in the Monotree
+    fn del(&mut self, key: &Hash) -> GenericResult<()>;
+    /// Initialize a batch
+    fn init_batch(&mut self) -> GenericResult<()>;
+    /// Finalize and write a batch
+    fn finish_batch(&mut self) -> GenericResult<()>;
+}
+
+/// In-memory storage for Monotree
 #[derive(Clone, Debug)]
 pub struct MemoryDb {
     db: HashMap<Hash, Vec<u8>>,
@@ -69,11 +85,23 @@ pub struct MemoryDb {
     batch_on: bool,
 }
 
-#[allow(dead_code)]
+#[allow(clippy::new_without_default)]
 impl MemoryDb {
-    fn new() -> Self {
+    pub fn new() -> Self {
         Self { db: HashMap::new(), batch: MemCache::new(), batch_on: false }
     }
+}
+
+impl MonotreeStorageAdapter for MemoryDb {
+    fn put(&mut self, key: &Hash, value: Vec<u8>) -> GenericResult<()> {
+        if self.batch_on {
+            self.batch.put(key, value);
+        } else {
+            self.db.insert(slice_to_hash(key), value);
+        }
+
+        Ok(())
+    }
 
     fn get(&self, key: &[u8]) -> GenericResult<Option<Vec<u8>>> {
         if self.batch_on && self.batch.contains(key) {
@@ -86,21 +114,13 @@ impl MemoryDb {
         }
     }
 
-    fn put(&mut self, key: &[u8], value: Vec<u8>) -> GenericResult<()> {
-        if self.batch_on {
-            self.batch.put(key, value);
-        } else {
-            self.db.insert(slice_to_hash(key), value);
-        }
-        Ok(())
-    }
-
-    fn delete(&mut self, key: &[u8]) -> GenericResult<()> {
+    fn del(&mut self, key: &Hash) -> GenericResult<()> {
         if self.batch_on {
-            self.batch.delete(key);
+            self.batch.del(key);
         } else {
             self.db.remove(key);
         }
+
         Ok(())
     }
 
@@ -109,6 +129,7 @@ impl MemoryDb {
             self.batch.clear();
             self.batch_on = true;
         }
+
         Ok(())
     }
 
@@ -122,25 +143,98 @@ impl MemoryDb {
             }
             self.batch_on = false;
         }
+
         Ok(())
     }
 }
 
-/// A structure for `monotree`
+/// sled-overlay-based storage for Monotree
 #[derive(Clone, Debug)]
-pub struct Monotree {
-    db: MemoryDb,
+pub struct SledDb {
+    db: SledTreeOverlay,
+    batch: MemCache,
+    batch_on: bool,
 }
 
-impl Default for Monotree {
-    fn default() -> Self {
-        Self::new()
+impl SledDb {
+    pub fn new(db: SledTreeOverlay) -> Self {
+        Self { db, batch: MemCache::new(), batch_on: false }
     }
 }
 
-impl Monotree {
-    pub fn new() -> Self {
-        Self { db: MemoryDb::new() }
+impl MonotreeStorageAdapter for SledDb {
+    fn put(&mut self, key: &Hash, value: Vec<u8>) -> GenericResult<()> {
+        if self.batch_on {
+            self.batch.put(key, value);
+        } else if let Err(e) = self.db.insert(&slice_to_hash(key), &value) {
+            return Err(ContractError::IoError(e.to_string()))
+        }
+
+        Ok(())
+    }
+
+    fn get(&self, key: &[u8]) -> GenericResult<Option<Vec<u8>>> {
+        if self.batch_on && self.batch.contains(key) {
+            return Ok(self.batch.get(key));
+        }
+
+        match self.db.get(key) {
+            Ok(Some(v)) => Ok(Some(v.to_vec())),
+            Ok(None) => Ok(None),
+            Err(e) => Err(ContractError::IoError(e.to_string())),
+        }
+    }
+
+    fn del(&mut self, key: &Hash) -> GenericResult<()> {
+        if self.batch_on {
+            self.batch.del(key);
+        } else if let Err(e) = self.db.remove(key) {
+            return Err(ContractError::IoError(e.to_string()));
+        }
+
+        Ok(())
+    }
+
+    fn init_batch(&mut self) -> GenericResult<()> {
+        if !self.batch_on {
+            self.batch.clear();
+            self.batch_on = true;
+        }
+
+        Ok(())
+    }
+
+    fn finish_batch(&mut self) -> GenericResult<()> {
+        if self.batch_on {
+            for (key, value) in self.batch.map.drain() {
+                if let Err(e) = self.db.insert(&key, &value) {
+                    return Err(ContractError::IoError(e.to_string()))
+                }
+            }
+            for key in self.batch.set.drain() {
+                if let Err(e) = self.db.remove(&key) {
+                    return Err(ContractError::IoError(e.to_string()))
+                }
+            }
+            self.batch_on = false;
+        }
+
+        Ok(())
+    }
+}
+
+/// A structure for `monotree`
+///
+/// To use this, first create a `MonotreeStorageAdapter` implementor,
+/// and then just manually create this struct with `::new()`.
+#[derive(Clone, Debug)]
+pub struct Monotree<D: MonotreeStorageAdapter> {
+    db: D,
+}
+
+impl<D: MonotreeStorageAdapter> Monotree<D> {
+    pub fn new(db: D) -> Self {
+        Self { db }
     }
 
     fn hash_digest(bytes: &[u8]) -> Hash {
@@ -403,6 +497,8 @@ impl Monotree {
 }
 
 /// Verify a MerkleProof with the given root and leaf.
+///
+/// NOTE: We use `Monotree::<MemoryDb>` to `hash_digest()` but it doesn't matter.
 pub fn verify_proof(root: Option<&Hash>, leaf: &Hash, proof: Option<&Proof>) -> bool {
     match proof {
         None => false,
@@ -412,10 +508,10 @@ pub fn verify_proof(root: Option<&Hash>, leaf: &Hash, proof: Option<&Proof>) ->
                 if *right {
                     let l = cut.len();
                     let o = [&cut[..l - 1], &hash[..], &cut[l - 1..]].concat();
-                    hash = Monotree::hash_digest(&o);
+                    hash = Monotree::<MemoryDb>::hash_digest(&o);
                 } else {
                     let o = [&hash[..], &cut[..]].concat();
-                    hash = Monotree::hash_digest(&o);
+                    hash = Monotree::<MemoryDb>::hash_digest(&o);
                 }
             });
             root.expect("verify_proof(): root") == &hash

src/validator/consensus.rs

@@ -18,7 +18,11 @@
 
 use std::collections::{HashMap, HashSet};
 
-use darkfi_sdk::{crypto::MerkleTree, monotree::Monotree, tx::TransactionHash};
+use darkfi_sdk::{
+    crypto::MerkleTree,
+    monotree::{self, Monotree},
+    tx::TransactionHash,
+};
 use darkfi_serial::{async_trait, SerialDecodable, SerialEncodable};
 use num_bigint::BigUint;
 use sled_overlay::database::SledDbOverlayStateDiff;
@@ -726,7 +730,7 @@ pub struct Fork {
     /// Current PoW module state
     pub module: PoWModule,
     /// Current contracts states checksums Monotree(SMT)
-    pub state_monotree: Monotree,
+    pub state_monotree: Monotree<monotree::MemoryDb>,
     /// Fork proposal hashes sequence
     pub proposals: Vec<HeaderHash>,
     /// Fork proposal overlay diffs sequence

src/validator/verification.rs

@@ -26,7 +26,7 @@ use darkfi_sdk::{
     },
     dark_tree::dark_forest_leaf_vec_integrity_check,
     deploy::DeployParamsV1,
-    monotree::Monotree,
+    monotree::{self, Monotree},
     pasta::pallas,
 };
 use darkfi_serial::{deserialize_async, serialize_async, AsyncDecodable, AsyncEncodable};
@@ -214,7 +214,7 @@ pub fn validate_blockchain(
 pub async fn verify_block(
     overlay: &BlockchainOverlayPtr,
     module: &PoWModule,
-    state_monotree: &mut Monotree,
+    state_monotree: &mut Monotree<monotree::MemoryDb>,
     block: &BlockInfo,
     previous: &BlockInfo,
     verify_fees: bool,
@@ -297,7 +297,7 @@ pub async fn verify_block(
 /// Verify given checkpoint [`BlockInfo`], and apply it to the provided overlay.
 pub async fn verify_checkpoint_block(
     overlay: &BlockchainOverlayPtr,
-    state_monotree: &mut Monotree,
+    state_monotree: &mut Monotree<monotree::MemoryDb>,
     block: &BlockInfo,
     header: &HeaderHash,
     block_target: u32,