sdk/crypto/smt: Implement leaf removal support

src/contract/deployooor/src/entrypoint.rs

@@ -19,9 +19,11 @@
 use darkfi_sdk::{
     crypto::ContractId,
     dark_tree::DarkLeaf,
-    db::{db_init, db_lookup, db_set},
     error::ContractResult,
-    util::set_return_data,
+    wasm::{
+        db::{db_init, db_lookup, db_set},
+        util::set_return_data,
+    },
     ContractCall,
 };
 use darkfi_serial::{deserialize, serialize};

src/contract/deployooor/src/entrypoint/deploy_v1.rs

@@ -19,11 +19,11 @@
 use darkfi_sdk::{
     crypto::{ContractId, PublicKey},
     dark_tree::DarkLeaf,
-    db::{db_get, db_lookup, db_set},
     deploy::DeployParamsV1,
     error::{ContractError, ContractResult},
     msg,
     pasta::pallas,
+    wasm::db::{db_get, db_lookup, db_set},
     ContractCall,
 };
 use darkfi_serial::{deserialize, serialize, Encodable, WriteExt};

src/contract/deployooor/src/entrypoint/lock_v1.rs

@@ -19,10 +19,10 @@
 use darkfi_sdk::{
     crypto::{ContractId, PublicKey},
     dark_tree::DarkLeaf,
-    db::{db_contains_key, db_get, db_lookup, db_set},
     error::{ContractError, ContractResult},
     msg,
     pasta::pallas,
+    wasm::db::{db_contains_key, db_get, db_lookup, db_set},
     ContractCall,
 };
 use darkfi_serial::{deserialize, serialize, Encodable, WriteExt};

src/runtime/import/smt.rs

@@ -23,6 +23,7 @@ use darkfi_sdk::{
         pasta_prelude::*,
         smt::{PoseidonFp, SparseMerkleTree, StorageAdapter, EMPTY_NODES_FP, SMT_FP_DEPTH},
     },
+    error::{ContractError, ContractResult},
     wasm,
 };
 use darkfi_serial::{serialize, Decodable, Encodable};
@@ -42,35 +43,50 @@ pub struct SledStorage<'a> {
 impl<'a> StorageAdapter for SledStorage<'a> {
     type Value = pallas::Base;
 
-    fn put(&mut self, key: BigUint, value: pallas::Base) -> bool {
-        if self.overlay.insert(self.tree_key, &key.to_bytes_le(), &value.to_repr()).is_err() {
+    fn put(&mut self, key: BigUint, value: pallas::Base) -> ContractResult {
+        if let Err(e) = self.overlay.insert(self.tree_key, &key.to_bytes_le(), &value.to_repr()) {
             error!(
                 target: "runtime::smt::SledStorage::put",
-                "[WASM] sparse_merkle_insert_batch(): inserting key {:?}, value {:?} into DB tree: {:?}",
-                key, value, self.tree_key
+                "[WASM] sparse_merkle_insert_batch(): inserting key {:?}, value {:?} into DB tree: {:?}: {}",
+                key, value, self.tree_key, e,
             );
-            return false
+            return Err(ContractError::SmtPutFailed)
         }
-        true
+
+        Ok(())
     }
+
     fn get(&self, key: &BigUint) -> Option<pallas::Base> {
-        let Ok(value) = self.overlay.get(self.tree_key, &key.to_bytes_le()) else {
-            error!(
-                target: "runtime::smt::SledStorage::get",
-                "[WASM] sparse_merkle_insert_batch(): fetching key {:?} from DB tree: {:?}",
-                key, self.tree_key
-            );
-            return None
+        let value = match self.overlay.get(self.tree_key, &key.to_bytes_le()) {
+            Ok(v) => v,
+            Err(e) => {
+                error!(
+                    target: "runtime::smt::SledStorage::get",
+                    "[WASM] SledStorage::get(): Fetching key {:?} from DB tree: {:?}: {}",
+                    key, self.tree_key, e,
+                );
+                return None
+            }
         };
+
         let value = value?;
         let mut repr = [0; 32];
         repr.copy_from_slice(&value);
-        let value = pallas::Base::from_repr(repr);
-        if value.is_none().into() {
-            None
-        } else {
-            Some(value.unwrap())
+
+        pallas::Base::from_repr(repr).into()
+    }
+
+    fn del(&mut self, key: &BigUint) -> ContractResult {
+        if let Err(e) = self.overlay.remove(self.tree_key, &key.to_bytes_le()) {
+            error!(
+                target: "runtime::smt::SledStorage::del",
+                "[WASM] SledStorage::del(): Removing key {:?} from DB tree: {:?}: {}",
+                key, self.tree_key, e,
+            );
+            return Err(ContractError::SmtDelFailed)
         }
+
+        Ok(())
     }
 }
 

src/sdk/src/crypto/smt/mod.rs

@@ -63,7 +63,7 @@ use std::collections::HashMap;
 // Only used for the type aliases below
 use pasta_curves::pallas;
 
-use crate::error::{ContractError, ContractResult};
+use crate::error::ContractResult;
 use util::{FieldElement, FieldHasher};
 
 mod empty;
@@ -92,12 +92,13 @@ pub type SmtMemoryFp = SparseMerkleTree<
 pub type PathFp = Path<SMT_FP_DEPTH, pallas::Base, PoseidonFp>;
 
 /// Pluggable storage backend for the SMT.
-/// Has a minimal interface to simply put and get objects from the store.
+/// Has a minimal interface to put, get, and delete objects from the store.
 pub trait StorageAdapter {
     type Value;
 
-    fn put(&mut self, key: BigUint, value: Self::Value) -> bool;
+    fn put(&mut self, key: BigUint, value: Self::Value) -> ContractResult;
     fn get(&self, key: &BigUint) -> Option<Self::Value>;
+    fn del(&mut self, key: &BigUint) -> ContractResult;
 }
 
 /// An in-memory storage, useful for unit tests and smaller trees.
@@ -115,14 +116,19 @@ impl<F: FieldElement> MemoryStorage<F> {
 impl<F: FieldElement> StorageAdapter for MemoryStorage<F> {
     type Value = F;
 
-    fn put(&mut self, key: BigUint, value: F) -> bool {
+    fn put(&mut self, key: BigUint, value: F) -> ContractResult {
         self.tree.insert(key, value);
-        true
+        Ok(())
     }
 
     fn get(&self, key: &BigUint) -> Option<F> {
         self.tree.get(key).copied()
     }
+
+    fn del(&mut self, key: &BigUint) -> ContractResult {
+        self.tree.remove(key);
+        Ok(())
+    }
 }
 
 /// The Sparse Merkle Tree struct.
@@ -212,6 +218,50 @@ impl<
         Ok(())
     }
 
+    pub fn remove_leaves(&mut self, leaves: Vec<(F, F)>) -> ContractResult {
+        if leaves.is_empty() {
+            return Ok(())
+        }
+
+        let mut dirty_idxs = Vec::new();
+        for (pos, _leaf) in leaves {
+            let idx = util::leaf_pos_to_index::<N, _>(&pos);
+            self.remove_node(&idx)?;
+
+            // Mark node parent as dirty
+            let parent_idx = util::parent(&idx).unwrap();
+            dirty_idxs.push(parent_idx);
+        }
+
+        // Depth first from the bottom of the tree
+        for _ in 0..N + 1 {
+            let mut new_dirty_idxs = Vec::new();
+
+            for idx in dirty_idxs {
+                let left_idx = util::left_child(&idx);
+                let right_idx = util::right_child(&idx);
+                let left = self.get_node(&left_idx);
+                let right = self.get_node(&right_idx);
+                // Recalculate the node
+                let node = self.hasher.hash([left, right]);
+
+                self.put_node(idx.clone(), node)?;
+
+                // Add this node's parent to the update list
+                let parent_idx = match util::parent(&idx) {
+                    Some(idx) => idx,
+                    // We are at the root node so no parents exist
+                    None => break,
+                };
+                new_dirty_idxs.push(parent_idx);
+            }
+
+            dirty_idxs = new_dirty_idxs;
+        }
+
+        Ok(())
+    }
+
     /// Returns the Merkle tree root.
     pub fn root(&self) -> F {
         self.get_node(&BigUint::from(0u32))
@@ -251,10 +301,11 @@ impl<
     }
 
     fn put_node(&mut self, key: BigUint, value: F) -> ContractResult {
-        if !self.store.put(key, value) {
-            return Err(ContractError::SmtPutFailed)
-        }
-        Ok(())
+        self.store.put(key, value)
+    }
+
+    fn remove_node(&mut self, key: &BigUint) -> ContractResult {
+        self.store.del(key)
     }
 }
 

src/sdk/src/crypto/smt/test.rs

@@ -141,6 +141,11 @@ fn poseidon_smt_incl_proof() {
         (Fp::from(1), Fp::random(&mut OsRng)),
         (Fp::from(2), Fp::random(&mut OsRng)),
         (Fp::from(3), Fp::random(&mut OsRng)),
+        /*
+        (Fp::from(1), Fp::from(111)),
+        (Fp::from(2), Fp::from(222)),
+        (Fp::from(3), Fp::from(333)),
+        */
     ];
     smt.insert_batch(leaves.clone()).unwrap();
 
@@ -149,4 +154,7 @@ fn poseidon_smt_incl_proof() {
 
     let path = smt.prove_membership(&pos);
     assert!(path.verify(&smt.root(), &leaf, &pos));
+
+    smt.remove_leaves(vec![(pos, leaf)]).unwrap();
+    assert!(!path.verify(&smt.root(), &leaf, &pos));
 }

src/sdk/src/crypto/smt/wasmdb.rs

@@ -21,9 +21,10 @@ use num_bigint::BigUint;
 use super::{PoseidonFp, SparseMerkleTree, StorageAdapter, SMT_FP_DEPTH};
 use crate::{
     crypto::pasta_prelude::*,
+    error::ContractResult,
     msg,
     pasta::pallas,
-    wasm::db::{db_get, db_set, DbHandle},
+    wasm::db::{db_del, db_get, db_set, DbHandle},
 };
 
 pub type SmtWasmFp = SparseMerkleTree<
@@ -48,8 +49,8 @@ impl SmtWasmDbStorage {
 impl StorageAdapter for SmtWasmDbStorage {
     type Value = pallas::Base;
 
-    fn put(&mut self, key: BigUint, value: pallas::Base) -> bool {
-        db_set(self.db, &key.to_bytes_le(), &value.to_repr()).is_ok()
+    fn put(&mut self, key: BigUint, value: pallas::Base) -> ContractResult {
+        db_set(self.db, &key.to_bytes_le(), &value.to_repr())
     }
 
     fn get(&self, key: &BigUint) -> Option<pallas::Base> {
@@ -65,4 +66,8 @@ impl StorageAdapter for SmtWasmDbStorage {
 
         pallas::Base::from_repr(repr).into()
     }
+
+    fn del(&mut self, key: &BigUint) -> ContractResult {
+        db_del(self.db, &key.to_bytes_le())
+    }
 }

src/sdk/src/error.rs

@@ -81,6 +81,9 @@ pub enum ContractError {
     #[error("SMT: put failed with storage backend")]
     SmtPutFailed,
 
+    #[error("SMT: rm failed with storage backend")]
+    SmtDelFailed,
+
     #[error("Error retrieving system time")]
     GetSystemTimeFailed,
 
@@ -119,9 +122,10 @@ pub const DB_CONTAINS_KEY_FAILED: i64 = to_builtin!(15);
 pub const INVALID_FUNCTION: i64 = to_builtin!(16);
 pub const DB_DEL_FAILED: i64 = to_builtin!(17);
 pub const SMT_PUT_FAILED: i64 = to_builtin!(18);
-pub const GET_SYSTEM_TIME_FAILED: i64 = to_builtin!(19);
-pub const DATA_TOO_LARGE: i64 = to_builtin!(20);
-pub const HEX_FMT_ERR: i64 = to_builtin!(21);
+pub const SMT_DEL_FAILED: i64 = to_builtin!(19);
+pub const GET_SYSTEM_TIME_FAILED: i64 = to_builtin!(20);
+pub const DATA_TOO_LARGE: i64 = to_builtin!(21);
+pub const HEX_FMT_ERR: i64 = to_builtin!(22);
 
 impl From<ContractError> for i64 {
     fn from(err: ContractError) -> Self {
@@ -143,6 +147,7 @@ impl From<ContractError> for i64 {
             ContractError::InvalidFunction => INVALID_FUNCTION,
             ContractError::DbDelFailed => DB_DEL_FAILED,
             ContractError::SmtPutFailed => SMT_PUT_FAILED,
+            ContractError::SmtDelFailed => SMT_DEL_FAILED,
             ContractError::GetSystemTimeFailed => GET_SYSTEM_TIME_FAILED,
             ContractError::DataTooLarge => DATA_TOO_LARGE,
             ContractError::HexFmtErr => HEX_FMT_ERR,
@@ -178,6 +183,7 @@ impl From<i64> for ContractError {
             INVALID_FUNCTION => Self::InvalidFunction,
             DB_DEL_FAILED => Self::DbDelFailed,
             SMT_PUT_FAILED => Self::SmtPutFailed,
+            SMT_DEL_FAILED => Self::SmtDelFailed,
             GET_SYSTEM_TIME_FAILED => Self::GetSystemTimeFailed,
             DATA_TOO_LARGE => Self::DataTooLarge,
             HEX_FMT_ERR => Self::HexFmtErr,