drk: introduced scan cache to minimize disk IO

bin/drk/src/cache.rs

@@ -16,7 +16,9 @@
  * along with this program.  If not, see <https://www.gnu.org/licenses/>.
  */
 
-use darkfi::{blockchain::HeaderHash, Result};
+use std::collections::HashMap;
+
+use darkfi::{blockchain::HeaderHash, Error, Result};
 use darkfi_sdk::{
     crypto::{
         pasta_prelude::PrimeField,
@@ -77,7 +79,7 @@ impl Cache {
 }
 
 /// Overlay structure over a [`Cache`] instance.
-pub struct CacheOverlay(SledDbOverlay);
+pub struct CacheOverlay(pub SledDbOverlay);
 
 impl CacheOverlay {
     /// Instantiate a new `CacheOverlay` over the given [`Cache`] instance.
@@ -135,25 +137,41 @@ pub type CacheSmt = SparseMerkleTree<
     { SMT_FP_DEPTH + 1 },
     pallas::Base,
     PoseidonFp,
-    CacheOverlay,
+    CacheSmtStorage,
 >;
 
-pub struct CacheSmtStorage<'a> {
-    overlay: &'a mut CacheOverlay,
-    tree: &'a [u8],
+pub struct CacheSmtStorage {
+    pub overlay: CacheOverlay,
+    tree: Vec<u8>,
 }
 
-impl<'a> CacheSmtStorage<'a> {
-    pub fn new(overlay: &'a mut CacheOverlay, tree: &'a [u8]) -> Self {
-        Self { overlay, tree }
+impl CacheSmtStorage {
+    pub fn new(overlay: CacheOverlay, tree: &[u8]) -> Self {
+        Self { overlay, tree: tree.to_vec() }
+    }
+
+    pub fn snapshot(&self) -> Result<HashMap<BigUint, pallas::Base>> {
+        let mut smt = HashMap::new();
+        for record in self.overlay.0.iter(&self.tree)? {
+            let (key, value) = record?;
+            let mut repr = [0; 32];
+            repr.copy_from_slice(&value);
+            let Some(value) = pallas::Base::from_repr(repr).into() else {
+                return Err(Error::ParseFailed(
+                    "[cache::CacheSmtStorage::snapshot] Value conversion failed",
+                ))
+            };
+            smt.insert(BigUint::from_bytes_le(&key), value);
+        }
+        Ok(smt)
     }
 }
 
-impl StorageAdapter for CacheSmtStorage<'_> {
+impl StorageAdapter for CacheSmtStorage {
     type Value = pallas::Base;
 
     fn put(&mut self, key: BigUint, value: pallas::Base) -> ContractResult {
-        if let Err(e) = self.overlay.0.insert(self.tree, &key.to_bytes_le(), &value.to_repr()) {
+        if let Err(e) = self.overlay.0.insert(&self.tree, &key.to_bytes_le(), &value.to_repr()) {
             error!(target: "cache::StorageAdapter::put", "Inserting key {key:?}, value {value:?} into DB failed: {e:?}");
             return Err(ContractError::SmtPutFailed)
         }
@@ -161,7 +179,7 @@ impl StorageAdapter for CacheSmtStorage<'_> {
     }
 
     fn get(&self, key: &BigUint) -> Option<pallas::Base> {
-        let value = match self.overlay.0.get(self.tree, &key.to_bytes_le()) {
+        let value = match self.overlay.0.get(&self.tree, &key.to_bytes_le()) {
             Ok(v) => v,
             Err(e) => {
                 error!(target: "cache::StorageAdapter::get", "Fetching key {key:?} from DB failed: {e:?}");
@@ -178,7 +196,7 @@ impl StorageAdapter for CacheSmtStorage<'_> {
     }
 
     fn del(&mut self, key: &BigUint) -> ContractResult {
-        if let Err(e) = self.overlay.0.remove(self.tree, &key.to_bytes_le()) {
+        if let Err(e) = self.overlay.0.remove(&self.tree, &key.to_bytes_le()) {
             error!(target: "cache::StorageAdapter::del", "Removing key {key:?} from DB failed: {e:?}");
             return Err(ContractError::SmtDelFailed)
         }
@@ -203,14 +221,14 @@ mod tests {
         // Setup cache and its overlay
         let sled_db = sled::Config::new().temporary(true).open()?;
         let cache = Cache::new(&sled_db)?;
-        let mut overlay = CacheOverlay::new(&cache)?;
+        let overlay = CacheOverlay::new(&cache)?;
 
         // Setup SMT
         const HEIGHT: usize = 3;
         let hasher = PoseidonFp::new();
         let empty_leaf = pallas::Base::ZERO;
         let empty_nodes = gen_empty_nodes::<{ HEIGHT + 1 }, _, _>(&hasher, empty_leaf);
-        let store = CacheSmtStorage::new(&mut overlay, SLED_MONEY_SMT_TREE);
+        let store = CacheSmtStorage::new(overlay, SLED_MONEY_SMT_TREE);
         let mut smt = SparseMerkleTree::<HEIGHT, { HEIGHT + 1 }, _, _, _>::new(
             store,
             hasher.clone(),
@@ -254,16 +272,16 @@ mod tests {
         assert!(path.verify(&root, &hash3, &pos));
 
         // Grab the overlay diff
-        let diff = overlay.0.diff(&[])?;
+        let diff = smt.store.overlay.0.diff(&[])?;
 
         // Apply the overlay
-        overlay.0.apply_diff(&diff)?;
+        smt.store.overlay.0.apply_diff(&diff)?;
 
         // Verify database contains keys
         assert!(!cache.money_smt.is_empty());
 
         // We are now going to rollback the changes
-        overlay.0.apply_diff(&diff.inverse())?;
+        smt.store.overlay.0.apply_diff(&diff.inverse())?;
 
         // Verify database is empty again
         assert!(cache.money_smt.is_empty());

bin/drk/src/dao.rs

@@ -72,13 +72,18 @@ use darkfi_serial::{
 };
 
 use crate::{
+    cache::{CacheOverlay, CacheSmt, CacheSmtStorage, SLED_MONEY_SMT_TREE},
     convert_named_params,
     error::{WalletDbError, WalletDbResult},
-    money::{BALANCE_BASE10_DECIMALS, MONEY_SMT_COL_KEY, MONEY_SMT_COL_VALUE, MONEY_SMT_TABLE},
-    walletdb::{WalletSmt, WalletStorage},
+    money::BALANCE_BASE10_DECIMALS,
+    rpc::ScanCache,
     Drk,
 };
 
+// DAO Merkle trees Sled keys
+pub const SLED_MERKLE_TREES_DAO_DAOS: &[u8] = b"_dao_daos";
+pub const SLED_MERKLE_TREES_DAO_PROPOSALS: &[u8] = b"_dao_proposals";
+
 // Wallet SQL table constant names. These have to represent the `dao.sql`
 // SQL schema. Table names are prefixed with the contract ID to avoid collisions.
 lazy_static! {
@@ -980,19 +985,6 @@ impl Drk {
         }
     }
 
-    /// Fetch all DAO notes secret keys from the wallet.
-    pub async fn get_dao_notes_secrets(&self) -> Result<Vec<SecretKey>> {
-        let daos = self.get_daos().await?;
-        let mut ret = Vec::with_capacity(daos.len());
-        for dao in daos {
-            if let Some(secret_key) = dao.params.notes_secret_key {
-                ret.push(secret_key);
-            }
-        }
-
-        Ok(ret)
-    }
-
     /// Auxiliary function to parse a `DAO_DAOS_TABLE` record.
     async fn parse_dao_record(&self, row: &[Value]) -> Result<DaoRecord> {
         let Value::Text(ref name) = row[1] else {
@@ -1191,49 +1183,34 @@ impl Drk {
     /// Returns a flag indicating if the provided call refers to our own wallet.
     async fn apply_dao_mint_data(
         &self,
-        new_bulla: DaoBulla,
-        tx_hash: TransactionHash,
-        call_index: u8,
+        scan_cache: &mut ScanCache,
+        new_bulla: &DaoBulla,
+        tx_hash: &TransactionHash,
+        call_index: &u8,
     ) -> Result<bool> {
-        let daos = self.get_daos().await?;
-        let (mut daos_tree, proposals_tree) = self.get_dao_trees().await?;
-        daos_tree.append(MerkleNode::from(new_bulla.inner()));
-
-        let mut wallet_tx = false;
-        for dao in &daos {
-            if dao.bulla() != new_bulla {
-                continue
-            }
-
-            println!(
-                "[apply_dao_mint_data] Found minted DAO {new_bulla}, noting down for wallet update"
-            );
-
-            // We have this DAO imported in our wallet. Add the metadata:
-            let mut dao_to_confirm = dao.clone();
-            dao_to_confirm.leaf_position = daos_tree.mark();
-            dao_to_confirm.tx_hash = Some(tx_hash);
-            dao_to_confirm.call_index = Some(call_index);
-
-            // Confirm it
-            if let Err(e) = self.confirm_dao(&dao_to_confirm).await {
-                return Err(Error::DatabaseError(format!(
-                    "[apply_dao_mint_data] Confirm DAO failed: {e:?}"
-                )))
-            }
+        // Append the new dao bulla to the Merkle tree.
+        // Every dao bulla has to be added.
+        scan_cache.dao_daos_tree.append(MerkleNode::from(new_bulla.inner()));
 
-            wallet_tx = true;
-            break
+        // Check if we have the DAO
+        if !scan_cache.own_daos.contains_key(new_bulla) {
+            return Ok(false)
         }
 
-        // Update wallet data
-        if let Err(e) = self.put_dao_trees(&daos_tree, &proposals_tree).await {
+        // Confirm it
+        println!(
+            "[apply_dao_mint_data] Found minted DAO {new_bulla}, noting down for wallet update"
+        );
+        if let Err(e) = self
+            .confirm_dao(new_bulla, &scan_cache.dao_daos_tree.mark().unwrap(), tx_hash, call_index)
+            .await
+        {
             return Err(Error::DatabaseError(format!(
-                "[apply_dao_mint_data] Put DAO tree failed: {e:?}"
+                "[apply_dao_mint_data] Confirm DAO failed: {e:?}"
             )))
         }
 
-        Ok(wallet_tx)
+        Ok(true)
     }
 
     /// Auxiliary function to apply `DaoFunction::Propose` call data to the wallet,
@@ -1241,54 +1218,53 @@ impl Drk {
     /// Returns a flag indicating if the provided call refers to our own wallet.
     async fn apply_dao_propose_data(
         &self,
-        params: DaoProposeParams,
-        tx_hash: TransactionHash,
-        call_index: u8,
+        scan_cache: &mut ScanCache,
+        params: &DaoProposeParams,
+        tx_hash: &TransactionHash,
+        call_index: &u8,
     ) -> Result<bool> {
-        let daos = self.get_daos().await?;
-        let (daos_tree, mut proposals_tree) = self.get_dao_trees().await?;
-        proposals_tree.append(MerkleNode::from(params.proposal_bulla.inner()));
+        // Append the new proposal bulla to the Merkle tree.
+        // Every proposal bulla has to be added.
+        scan_cache.dao_proposals_tree.append(MerkleNode::from(params.proposal_bulla.inner()));
 
         // If we're able to decrypt this note, that's the way to link it
         // to a specific DAO.
-        let mut wallet_tx = false;
-        for dao in &daos {
+        for (dao, (proposals_secret_key, _)) in &scan_cache.own_daos {
             // Check if we have the proposals key
-            let Some(proposals_secret_key) = dao.params.proposals_secret_key else { continue };
+            let Some(proposals_secret_key) = proposals_secret_key else { continue };
 
             // Try to decrypt the proposal note
-            let Ok(note) = params.note.decrypt::<DaoProposal>(&proposals_secret_key) else {
+            let Ok(note) = params.note.decrypt::<DaoProposal>(proposals_secret_key) else {
                 continue
             };
 
             // We managed to decrypt it. Let's place this in a proper ProposalRecord object
-            println!("[apply_dao_propose_data] Managed to decrypt DAO proposal note");
-
-            // We need to clone the trees here for reproducing the snapshot Merkle roots
-            let money_tree = self.get_money_tree().await?;
-            let nullifiers_smt = self.get_nullifiers_smt().await?;
+            println!("[apply_dao_propose_data] Managed to decrypt proposal note for DAO: {dao}");
 
             // Check if we already got the record
-            let our_proposal = match self.get_dao_proposal_by_bulla(&params.proposal_bulla).await {
-                Ok(p) => {
-                    let mut our_proposal = p;
-                    our_proposal.leaf_position = proposals_tree.mark();
-                    our_proposal.money_snapshot_tree = Some(money_tree);
-                    our_proposal.nullifiers_smt_snapshot = Some(nullifiers_smt);
-                    our_proposal.tx_hash = Some(tx_hash);
-                    our_proposal.call_index = Some(call_index);
-                    our_proposal
-                }
-                Err(_) => ProposalRecord {
+            let our_proposal = if scan_cache.own_proposals.contains_key(&params.proposal_bulla) {
+                // Grab the record from the db
+                let mut our_proposal =
+                    self.get_dao_proposal_by_bulla(&params.proposal_bulla).await?;
+                our_proposal.leaf_position = scan_cache.dao_proposals_tree.mark();
+                our_proposal.money_snapshot_tree = Some(scan_cache.money_tree.clone());
+                our_proposal.nullifiers_smt_snapshot = Some(scan_cache.money_smt.store.snapshot()?);
+                our_proposal.tx_hash = Some(*tx_hash);
+                our_proposal.call_index = Some(*call_index);
+                our_proposal
+            } else {
+                let our_proposal = ProposalRecord {
                     proposal: note,
                     data: None,
-                    leaf_position: proposals_tree.mark(),
-                    money_snapshot_tree: Some(money_tree),
-                    nullifiers_smt_snapshot: Some(nullifiers_smt),
-                    tx_hash: Some(tx_hash),
-                    call_index: Some(call_index),
+                    leaf_position: scan_cache.dao_proposals_tree.mark(),
+                    money_snapshot_tree: Some(scan_cache.money_tree.clone()),
+                    nullifiers_smt_snapshot: Some(scan_cache.money_smt.store.snapshot()?),
+                    tx_hash: Some(*tx_hash),
+                    call_index: Some(*call_index),
                     exec_tx_hash: None,
-                },
+                };
+                scan_cache.own_proposals.insert(params.proposal_bulla, *dao);
+                our_proposal
             };
 
             // Update/store our record
@@ -1298,18 +1274,10 @@ impl Drk {
                 )))
             }
 
-            wallet_tx = true;
-            break
-        }
-
-        // Update wallet data
-        if let Err(e) = self.put_dao_trees(&daos_tree, &proposals_tree).await {
-            return Err(Error::DatabaseError(format!(
-                "[apply_dao_propose_data] Put DAO tree failed: {e:?}"
-            )))
+            return Ok(true)
         }
 
-        Ok(wallet_tx)
+        Ok(false)
     }
 
     /// Auxiliary function to apply `DaoFunction::Vote` call data to the wallet,
@@ -1317,38 +1285,34 @@ impl Drk {
     /// Returns a flag indicating if the provided call refers to our own wallet.
     async fn apply_dao_vote_data(
         &self,
-        params: DaoVoteParams,
-        tx_hash: TransactionHash,
-        call_index: u8,
+        scan_cache: &ScanCache,
+        params: &DaoVoteParams,
+        tx_hash: &TransactionHash,
+        call_index: &u8,
     ) -> Result<bool> {
         // Check if we got the corresponding proposal
-        let Ok(proposal) = self.get_dao_proposal_by_bulla(&params.proposal_bulla).await else {
+        let Some(dao_bulla) = scan_cache.own_proposals.get(&params.proposal_bulla) else {
             return Ok(false)
         };
 
-        // Grab the proposal DAO
-        let dao = match self.get_dao_by_bulla(&proposal.proposal.dao_bulla).await {
-            Ok(d) => d,
-            Err(e) => {
-                return Err(Error::DatabaseError(format!(
-                    "[apply_dao_vote_data] Couldn't find proposal {} DAO {}: {e}",
-                    proposal.bulla(),
-                    proposal.proposal.dao_bulla,
-                )))
-            }
+        // Grab the proposal DAO votes key
+        let Some((_, votes_secret_key)) = scan_cache.own_daos.get(dao_bulla) else {
+            return Err(Error::DatabaseError(format!(
+                "[apply_dao_vote_data] Couldn't find proposal {} DAO {}",
+                params.proposal_bulla, dao_bulla,
+            )))
         };
 
-        // Check if we have the votes key
-        let Some(votes_secret_key) = dao.params.votes_secret_key else { return Ok(false) };
+        // Check if we actually have the votes key
+        let Some(votes_secret_key) = votes_secret_key else { return Ok(false) };
 
         // Decrypt the vote note
-        let note = match params.note.decrypt_unsafe(&votes_secret_key) {
+        let note = match params.note.decrypt_unsafe(votes_secret_key) {
             Ok(n) => n,
             Err(e) => {
                 return Err(Error::DatabaseError(format!(
                     "[apply_dao_vote_data] Couldn't decrypt proposal {} vote with DAO {} keys: {e}",
-                    proposal.bulla(),
-                    proposal.proposal.dao_bulla,
+                    params.proposal_bulla, dao_bulla,
                 )))
             }
         };
@@ -1358,7 +1322,7 @@ impl Drk {
         if vote_option > 1 {
             return Err(Error::DatabaseError(format!(
                 "[apply_dao_vote_data] Malformed vote for proposal {}: {vote_option}",
-                proposal.bulla(),
+                params.proposal_bulla,
             )))
         }
         let vote_option = vote_option != 0;
@@ -1373,8 +1337,8 @@ impl Drk {
             yes_vote_blind,
             all_vote_value,
             all_vote_blind,
-            tx_hash,
-            call_index,
+            tx_hash: *tx_hash,
+            call_index: *call_index,
             nullifiers: params.inputs.iter().map(|i| i.vote_nullifier).collect(),
         };
 
@@ -1392,13 +1356,14 @@ impl Drk {
     /// Returns a flag indicating if the provided call refers to our own wallet.
     async fn apply_dao_exec_data(
         &self,
-        params: DaoExecParams,
-        tx_hash: TransactionHash,
+        scan_cache: &ScanCache,
+        params: &DaoExecParams,
+        tx_hash: &TransactionHash,
     ) -> Result<bool> {
         // Check if we got the corresponding proposal
-        if self.get_dao_proposal_by_bulla(&params.proposal_bulla).await.is_err() {
+        if !scan_cache.own_proposals.contains_key(&params.proposal_bulla) {
             return Ok(false)
-        };
+        }
 
         // Grab proposal record key
         let key = serialize_async(&params.proposal_bulla).await;
@@ -1428,7 +1393,7 @@ impl Drk {
         // Execute the query
         if let Err(e) = self
             .wallet
-            .exec_sql(&query, rusqlite::params![Some(serialize_async(&tx_hash).await), key])
+            .exec_sql(&query, rusqlite::params![Some(serialize_async(tx_hash).await), key])
         {
             return Err(Error::DatabaseError(format!(
                 "[apply_dao_exec_data] Update DAO proposal failed: {e:?}"
@@ -1447,39 +1412,51 @@ impl Drk {
 
     /// Append data related to DAO contract transactions into the wallet database,
     /// and store their inverse queries into the cache.
-    /// Returns a flag indicating if the provided data refer to our own wallet.
+    /// Returns a flag indicating if the daos tree should be updated,
+    /// one indicating if the proposals tree should be updated and
+    /// another one indicating if provided data refer to our own
+    /// wallet.
     pub async fn apply_tx_dao_data(
         &self,
+        scan_cache: &mut ScanCache,
         data: &[u8],
-        tx_hash: TransactionHash,
-        call_idx: u8,
-    ) -> Result<bool> {
+        tx_hash: &TransactionHash,
+        call_idx: &u8,
+    ) -> Result<(bool, bool, bool)> {
         // Run through the transaction call data and see what we got:
         match DaoFunction::try_from(data[0])? {
             DaoFunction::Mint => {
                 println!("[apply_tx_dao_data] Found Dao::Mint call");
                 let params: DaoMintParams = deserialize_async(&data[1..]).await?;
-                self.apply_dao_mint_data(params.dao_bulla, tx_hash, call_idx).await
+                let own_tx = self
+                    .apply_dao_mint_data(scan_cache, &params.dao_bulla, tx_hash, call_idx)
+                    .await?;
+                Ok((true, false, own_tx))
             }
             DaoFunction::Propose => {
                 println!("[apply_tx_dao_data] Found Dao::Propose call");
                 let params: DaoProposeParams = deserialize_async(&data[1..]).await?;
-                self.apply_dao_propose_data(params, tx_hash, call_idx).await
+                let own_tx =
+                    self.apply_dao_propose_data(scan_cache, &params, tx_hash, call_idx).await?;
+                Ok((false, true, own_tx))
             }
             DaoFunction::Vote => {
                 println!("[apply_tx_dao_data] Found Dao::Vote call");
                 let params: DaoVoteParams = deserialize_async(&data[1..]).await?;
-                self.apply_dao_vote_data(params, tx_hash, call_idx).await
+                let own_tx =
+                    self.apply_dao_vote_data(scan_cache, &params, tx_hash, call_idx).await?;
+                Ok((false, false, own_tx))
             }
             DaoFunction::Exec => {
                 println!("[apply_tx_dao_data] Found Dao::Exec call");
                 let params: DaoExecParams = deserialize_async(&data[1..]).await?;
-                self.apply_dao_exec_data(params, tx_hash).await
+                let own_tx = self.apply_dao_exec_data(scan_cache, &params, tx_hash).await?;
+                Ok((false, false, own_tx))
             }
             DaoFunction::AuthMoneyTransfer => {
                 println!("[apply_tx_dao_data] Found Dao::AuthMoneyTransfer call");
                 // Does nothing, just verifies the other calls are correct
-                Ok(false)
+                Ok((false, false, false))
             }
         }
     }
@@ -1488,9 +1465,15 @@ impl Drk {
     /// and store its inverse query into the cache.
     /// Here we just write the leaf position, tx hash, and call index.
     /// Panics if the fields are None.
-    pub async fn confirm_dao(&self, dao: &DaoRecord) -> WalletDbResult<()> {
+    pub async fn confirm_dao(
+        &self,
+        dao: &DaoBulla,
+        leaf_position: &bridgetree::Position,
+        tx_hash: &TransactionHash,
+        call_index: &u8,
+    ) -> WalletDbResult<()> {
         // Grab dao record key
-        let key = serialize_async(&dao.bulla()).await;
+        let key = serialize_async(dao).await;
 
         // Create an SQL `UPDATE` query
         let query = format!(
@@ -1504,9 +1487,9 @@ impl Drk {
 
         // Create its params
         let params = rusqlite::params![
-            serialize_async(&dao.leaf_position.unwrap()).await,
-            serialize_async(&dao.tx_hash.unwrap()).await,
-            dao.call_index.unwrap(),
+            serialize_async(leaf_position).await,
+            serialize_async(tx_hash).await,
+            call_index,
             key,
         ];
 
@@ -2491,13 +2474,8 @@ impl Drk {
         };
 
         // Generate the Money nullifiers Sparse Merkle Tree
-        let store = WalletStorage::new(
-            &self.wallet,
-            &MONEY_SMT_TABLE,
-            MONEY_SMT_COL_KEY,
-            MONEY_SMT_COL_VALUE,
-        );
-        let money_null_smt = WalletSmt::new(store, PoseidonFp::new(), &EMPTY_NODES_FP);
+        let store = CacheSmtStorage::new(CacheOverlay::new(&self.cache)?, SLED_MONEY_SMT_TREE);
+        let money_null_smt = CacheSmt::new(store, PoseidonFp::new(), &EMPTY_NODES_FP);
 
         // Create the proposal call
         let call = DaoProposeCall {
@@ -2679,13 +2657,8 @@ impl Drk {
         };
 
         // Generate the Money nullifiers Sparse Merkle Tree
-        let store = WalletStorage::new(
-            &self.wallet,
-            &MONEY_SMT_TABLE,
-            MONEY_SMT_COL_KEY,
-            MONEY_SMT_COL_VALUE,
-        );
-        let money_null_smt = WalletSmt::new(store, PoseidonFp::new(), &EMPTY_NODES_FP);
+        let store = CacheSmtStorage::new(CacheOverlay::new(&self.cache)?, SLED_MONEY_SMT_TREE);
+        let money_null_smt = CacheSmt::new(store, PoseidonFp::new(), &EMPTY_NODES_FP);
 
         // Create the proposal call
         let call = DaoProposeCall {

bin/drk/src/money.rs

@@ -16,7 +16,10 @@
  * along with this program.  If not, see <https://www.gnu.org/licenses/>.
  */
 
-use std::{collections::HashMap, str::FromStr};
+use std::{
+    collections::{BTreeMap, HashMap},
+    str::FromStr,
+};
 
 use lazy_static::lazy_static;
 use num_bigint::BigUint;
@@ -46,11 +49,8 @@ use darkfi_money_contract::{
 use darkfi_sdk::{
     bridgetree,
     crypto::{
-        note::AeadEncryptedNote,
-        pasta_prelude::PrimeField,
-        smt::{PoseidonFp, EMPTY_NODES_FP},
-        BaseBlind, FuncId, Keypair, MerkleNode, MerkleTree, PublicKey, ScalarBlind, SecretKey,
-        MONEY_CONTRACT_ID,
+        note::AeadEncryptedNote, pasta_prelude::PrimeField, BaseBlind, FuncId, Keypair, MerkleNode,
+        MerkleTree, PublicKey, ScalarBlind, SecretKey, MONEY_CONTRACT_ID,
     },
     dark_tree::DarkLeaf,
     pasta::pallas,
@@ -59,13 +59,13 @@ use darkfi_sdk::{
 use darkfi_serial::{deserialize_async, serialize_async, AsyncEncodable};
 
 use crate::{
-    cli_util::kaching,
-    convert_named_params,
-    error::WalletDbResult,
-    walletdb::{WalletSmt, WalletStorage},
-    Drk,
+    cache::CacheSmt, cli_util::kaching, convert_named_params, error::WalletDbResult,
+    rpc::ScanCache, Drk,
 };
 
+// Money Merkle tree Sled key
+pub const SLED_MERKLE_TREES_MONEY: &[u8] = b"_money_tree";
+
 // Wallet SQL table constant names. These have to represent the `money.sql`
 // SQL schema. Table names are prefixed with the contract ID to avoid collisions.
 lazy_static! {
@@ -405,31 +405,6 @@ impl Drk {
         Ok(owncoins)
     }
 
-    /// Fetch provided transaction coins from the wallet.
-    pub async fn get_transaction_coins(&self, spent_tx_hash: &String) -> Result<Vec<OwnCoin>> {
-        let query = self.wallet.query_multiple(
-            &MONEY_COINS_TABLE,
-            &[],
-            convert_named_params! {(MONEY_COINS_COL_SPENT_TX_HASH, spent_tx_hash)},
-        );
-
-        let rows = match query {
-            Ok(r) => r,
-            Err(e) => {
-                return Err(Error::DatabaseError(format!(
-                    "[get_transaction_coins] Coins retrieval failed: {e:?}"
-                )))
-            }
-        };
-
-        let mut owncoins = Vec::with_capacity(rows.len());
-        for row in rows {
-            owncoins.push(self.parse_coin_record(&row).await?.0)
-        }
-
-        Ok(owncoins)
-    }
-
     /// Fetch provided token unspend balances from the wallet.
     pub async fn get_token_coins(&self, token_id: &TokenId) -> Result<Vec<OwnCoin>> {
         let query = self.wallet.query_multiple(
@@ -742,41 +717,37 @@ impl Drk {
         Ok(smt)
     }
 
-    /// Auxiliary function to grab all the nullifiers, coins, notes and freezes from
-    /// a transaction money call.
+    /// Auxiliary function to grab all the nullifiers, coins with their
+    /// notes and freezes from a transaction money call.
     async fn parse_money_call(
         &self,
-        call_idx: usize,
+        call_idx: &usize,
         calls: &[DarkLeaf<ContractCall>],
-    ) -> Result<(Vec<Nullifier>, Vec<Coin>, Vec<AeadEncryptedNote>, Vec<TokenId>)> {
+    ) -> Result<(Vec<Nullifier>, Vec<(Coin, AeadEncryptedNote)>, Vec<TokenId>)> {
         let mut nullifiers: Vec<Nullifier> = vec![];
-        let mut coins: Vec<Coin> = vec![];
-        let mut notes: Vec<AeadEncryptedNote> = vec![];
+        let mut coins: Vec<(Coin, AeadEncryptedNote)> = vec![];
         let mut freezes: Vec<TokenId> = vec![];
 
-        let call = &calls[call_idx];
+        let call = &calls[*call_idx];
         let data = &call.data.data;
         match MoneyFunction::try_from(data[0])? {
             MoneyFunction::FeeV1 => {
                 println!("[parse_money_call] Found Money::FeeV1 call");
                 let params: MoneyFeeParamsV1 = deserialize_async(&data[9..]).await?;
                 nullifiers.push(params.input.nullifier);
-                coins.push(params.output.coin);
-                notes.push(params.output.note);
+                coins.push((params.output.coin, params.output.note));
             }
             MoneyFunction::GenesisMintV1 => {
                 println!("[parse_money_call] Found Money::GenesisMintV1 call");
                 let params: MoneyGenesisMintParamsV1 = deserialize_async(&data[1..]).await?;
                 for output in params.outputs {
-                    coins.push(output.coin);
-                    notes.push(output.note);
+                    coins.push((output.coin, output.note));
                 }
             }
             MoneyFunction::PoWRewardV1 => {
                 println!("[parse_money_call] Found Money::PoWRewardV1 call");
                 let params: MoneyPoWRewardParamsV1 = deserialize_async(&data[1..]).await?;
-                coins.push(params.output.coin);
-                notes.push(params.output.note);
+                coins.push((params.output.coin, params.output.note));
             }
             MoneyFunction::TransferV1 => {
                 println!("[parse_money_call] Found Money::TransferV1 call");
@@ -787,8 +758,7 @@ impl Drk {
                 }
 
                 for output in params.outputs {
-                    coins.push(output.coin);
-                    notes.push(output.note);
+                    coins.push((output.coin, output.note));
                 }
             }
             MoneyFunction::OtcSwapV1 => {
@@ -800,8 +770,7 @@ impl Drk {
                 }
 
                 for output in params.outputs {
-                    coins.push(output.coin);
-                    notes.push(output.note);
+                    coins.push((output.coin, output.note));
                 }
             }
             MoneyFunction::AuthTokenMintV1 => {
@@ -816,61 +785,69 @@ impl Drk {
             MoneyFunction::TokenMintV1 => {
                 println!("[parse_money_call] Found Money::TokenMintV1 call");
                 let params: MoneyTokenMintParamsV1 = deserialize_async(&data[1..]).await?;
-                coins.push(params.coin);
                 // Grab the note from the child auth call
                 let child_idx = call.children_indexes[0];
                 let child_call = &calls[child_idx];
-                let params: MoneyAuthTokenMintParamsV1 =
+                let child_params: MoneyAuthTokenMintParamsV1 =
                     deserialize_async(&child_call.data.data[1..]).await?;
-                notes.push(params.enc_note);
+                coins.push((params.coin, child_params.enc_note));
             }
         }
 
-        Ok((nullifiers, coins, notes, freezes))
+        Ok((nullifiers, coins, freezes))
     }
 
-    /// Append data related to Money contract transactions into the wallet database,
-    /// and store their inverse queries into the cache.
-    /// Returns a flag indicating if the provided data refer to our own wallet.
-    pub async fn apply_tx_money_data(
+    /// Auxiliary function to handle coins with their notes from a
+    /// transaction money call.
+    /// Returns a flag indicating if the money tree should be updated,
+    /// along with found own coins.
+    fn handle_money_call_coins(
         &self,
-        call_idx: usize,
-        calls: &[DarkLeaf<ContractCall>],
-        tx_hash: &String,
-    ) -> Result<bool> {
-        let (nullifiers, coins, notes, freezes) = self.parse_money_call(call_idx, calls).await?;
-        let secrets = self.get_money_secrets().await?;
-        let dao_notes_secrets = self.get_dao_notes_secrets().await?;
-        let mut tree = self.get_money_tree().await?;
-
+        tree: &mut MerkleTree,
+        secrets: &[SecretKey],
+        coins: &[(Coin, AeadEncryptedNote)],
+    ) -> (bool, Vec<OwnCoin>) {
+        // Keep track of our own coins found in the vec
         let mut owncoins = vec![];
 
-        for (coin, note) in coins.iter().zip(notes.iter()) {
-            // Append the new coin to the Merkle tree. Every coin has to be added.
+        // Check if provided coins vec is empty
+        if coins.is_empty() {
+            return (false, owncoins)
+        }
+
+        // Handle provided coins vector and grab our own
+        for (coin, note) in coins {
+            // Append the new coin to the Merkle tree.
+            // Every coin has to be added.
             tree.append(MerkleNode::from(coin.inner()));
 
             // Attempt to decrypt the note
-            for secret in secrets.iter().chain(dao_notes_secrets.iter()) {
-                if let Ok(note) = note.decrypt::<MoneyNote>(secret) {
-                    println!("[apply_tx_money_data] Successfully decrypted a Money Note");
-                    println!("[apply_tx_money_data] Witnessing coin in Merkle tree");
-                    let leaf_position = tree.mark().unwrap();
-
-                    let owncoin =
-                        OwnCoin { coin: *coin, note: note.clone(), secret: *secret, leaf_position };
-
-                    owncoins.push(owncoin);
-                }
+            for secret in secrets {
+                let Ok(note) = note.decrypt::<MoneyNote>(secret) else { continue };
+                println!("[handle_money_call_coins] Successfully decrypted a Money Note");
+                println!("[handle_money_call_coins] Witnessing coin in Merkle tree");
+                let leaf_position = tree.mark().unwrap();
+                let owncoin = OwnCoin { coin: *coin, note, secret: *secret, leaf_position };
+                owncoins.push(owncoin);
             }
         }
 
-        if let Err(e) = self.put_money_tree(&tree).await {
-            return Err(Error::DatabaseError(format!(
-                "[apply_tx_money_data] Put Money tree failed: {e:?}"
-            )))
+        (true, owncoins)
+    }
+
+    /// Auxiliary function to handle own coins from a transaction money
+    /// call.
+    async fn handle_money_call_owncoins(
+        &self,
+        owncoins_nullifiers: &mut BTreeMap<[u8; 32], [u8; 32]>,
+        coins: &[OwnCoin],
+    ) -> Result<()> {
+        println!("Found {} OwnCoin(s) in transaction", coins.len());
+
+        // Check if we have any owncoins to process
+        if coins.is_empty() {
+            return Ok(())
         }
-        self.smt_insert(&nullifiers)?;
-        let wallet_spent_coins = self.mark_spent_coins(&nullifiers, tx_hash).await?;
 
         // This is the SQL query we'll be executing to insert new coins into the wallet
         let query = format!(
@@ -894,11 +871,14 @@ impl Drk {
         let inverse_query =
             format!("DELETE FROM {} WHERE {} = ?1;", *MONEY_COINS_TABLE, MONEY_COINS_COL_COIN);
 
-        println!("Found {} OwnCoin(s) in transaction", owncoins.len());
-        for owncoin in &owncoins {
-            println!("OwnCoin: {:?}", owncoin.coin);
+        // Handle our own coins
+        for coin in coins {
+            println!("OwnCoin: {:?}", coin.coin);
             // Grab coin record key
-            let key = serialize_async(&owncoin.coin).await;
+            let key = coin.coin.to_bytes();
+
+            // Push to our own coins nullifiers cache
+            owncoins_nullifiers.insert(coin.nullifier().to_bytes(), key);
 
             // Create its inverse query
             let inverse =
@@ -907,7 +887,7 @@ impl Drk {
                     Ok(q) => q,
                     Err(e) => {
                         return Err(Error::DatabaseError(format!(
-                    "[apply_tx_money_data] Creating Money coin insert inverse query failed: {e:?}"
+                    "[handle_money_call_owncoins] Creating Money coin insert inverse query failed: {e:?}"
                 )))
                     }
                 };
@@ -916,32 +896,43 @@ impl Drk {
             let params = rusqlite::params![
                 key,
                 0, // <-- is_spent
-                serialize_async(&owncoin.note.value).await,
-                serialize_async(&owncoin.note.token_id).await,
-                serialize_async(&owncoin.note.spend_hook).await,
-                serialize_async(&owncoin.note.user_data).await,
-                serialize_async(&owncoin.note.coin_blind).await,
-                serialize_async(&owncoin.note.value_blind).await,
-                serialize_async(&owncoin.note.token_blind).await,
-                serialize_async(&owncoin.secret).await,
-                serialize_async(&owncoin.leaf_position).await,
-                serialize_async(&owncoin.note.memo).await,
+                serialize_async(&coin.note.value).await,
+                serialize_async(&coin.note.token_id).await,
+                serialize_async(&coin.note.spend_hook).await,
+                serialize_async(&coin.note.user_data).await,
+                serialize_async(&coin.note.coin_blind).await,
+                serialize_async(&coin.note.value_blind).await,
+                serialize_async(&coin.note.token_blind).await,
+                serialize_async(&coin.secret).await,
+                serialize_async(&coin.leaf_position).await,
+                serialize_async(&coin.note.memo).await,
             ];
 
             if let Err(e) = self.wallet.exec_sql(&query, params) {
                 return Err(Error::DatabaseError(format!(
-                    "[apply_tx_money_data] Inserting Money coin failed: {e:?}"
+                    "[handle_money_call_owncoins] Inserting Money coin failed: {e:?}"
                 )))
             }
 
             // Store its inverse
             if let Err(e) = self.wallet.cache_inverse(inverse) {
                 return Err(Error::DatabaseError(format!(
-                    "[apply_tx_money_data] Inserting inverse query into cache failed: {e:?}"
+                    "[handle_money_call_owncoins] Inserting inverse query into cache failed: {e:?}"
                 )))
             }
         }
 
+        Ok(())
+    }
+
+    /// Auxiliary function to handle freezes from a transaction money
+    /// call.
+    async fn handle_money_call_freezes(&self, freezes: &[TokenId]) -> Result<()> {
+        // Check if we have any freezes to process
+        if freezes.is_empty() {
+            return Ok(())
+        }
+
         // This is the SQL query we'll be executing to update frozen tokens into the wallet
         let query = format!(
             "UPDATE {} SET {} = 1 WHERE {} = ?1;",
@@ -954,7 +945,7 @@ impl Drk {
             *MONEY_TOKENS_TABLE, MONEY_TOKENS_COL_IS_FROZEN, MONEY_TOKENS_COL_TOKEN_ID,
         );
 
-        for token_id in &freezes {
+        for token_id in freezes {
             // Grab token record key
             let key = serialize_async(token_id).await;
 
@@ -965,7 +956,7 @@ impl Drk {
                     Ok(q) => q,
                     Err(e) => {
                         return Err(Error::DatabaseError(format!(
-                    "[apply_tx_money_data] Creating Money token freeze inverse query failed: {e:?}"
+                    "[handle_money_call_freezes] Creating Money token freeze inverse query failed: {e:?}"
                 )))
                     }
                 };
@@ -973,23 +964,62 @@ impl Drk {
             // Execute the query
             if let Err(e) = self.wallet.exec_sql(&query, rusqlite::params![key]) {
                 return Err(Error::DatabaseError(format!(
-                    "[apply_tx_money_data] Update Money token freeze failed: {e:?}"
+                    "[handle_money_call_freezes] Update Money token freeze failed: {e:?}"
                 )))
             }
 
             // Store its inverse
             if let Err(e) = self.wallet.cache_inverse(inverse) {
                 return Err(Error::DatabaseError(format!(
-                    "[apply_tx_money_data] Inserting inverse query into cache failed: {e:?}"
+                    "[handle_money_call_freezes] Inserting inverse query into cache failed: {e:?}"
                 )))
             }
         }
 
+        Ok(())
+    }
+
+    /// Append data related to Money contract transactions into the
+    /// wallet database, and store their inverse queries into the
+    /// cache.
+    /// Returns a flag indicating if the money tree should be updated
+    /// and one indicating if provided data refer to our own wallet.
+    pub async fn apply_tx_money_data(
+        &self,
+        scan_cache: &mut ScanCache,
+        call_idx: &usize,
+        calls: &[DarkLeaf<ContractCall>],
+        tx_hash: &String,
+    ) -> Result<(bool, bool)> {
+        // Parse the call
+        let (nullifiers, coins, freezes) = self.parse_money_call(call_idx, calls).await?;
+
+        // Parse call coins and grab our own
+        let (update_tree, owncoins) = self.handle_money_call_coins(
+            &mut scan_cache.money_tree,
+            &scan_cache.notes_secrets,
+            &coins,
+        );
+
+        // Update nullifiers smt
+        self.smt_insert(&mut scan_cache.money_smt, &nullifiers)?;
+
+        // Check if we have any spent coins
+        let wallet_spent_coins =
+            self.mark_spent_coins(&scan_cache.owncoins_nullifiers, &nullifiers, tx_hash)?;
+
+        // Handle our own coins
+        self.handle_money_call_owncoins(&mut scan_cache.owncoins_nullifiers, &owncoins).await?;
+
+        // Handle freezes
+        // TODO: this should return flag if we have frozen tokens indeed
+        self.handle_money_call_freezes(&freezes).await?;
+
         if self.fun && !owncoins.is_empty() {
             kaching().await;
         }
 
-        Ok(wallet_spent_coins || !owncoins.is_empty() || !freezes.is_empty())
+        Ok((update_tree, wallet_spent_coins || !owncoins.is_empty() || !freezes.is_empty()))
     }
 
     /// Auxiliary function to  grab all the nullifiers from a transaction money call.
@@ -1024,6 +1054,12 @@ impl Drk {
 
     /// Mark provided transaction input coins as spent.
     pub async fn mark_tx_spend(&self, tx: &Transaction) -> Result<()> {
+        // Create a cache of all our own nullifiers
+        let mut owncoins_nullifiers = BTreeMap::new();
+        for coin in self.get_coins(true).await? {
+            owncoins_nullifiers.insert(coin.0.nullifier().to_bytes(), coin.0.coin.to_bytes());
+        }
+
         let tx_hash = tx.hash().to_string();
         println!("[mark_tx_spend] Processing transaction: {tx_hash}");
         for (i, call) in tx.calls.iter().enumerate() {
@@ -1033,16 +1069,34 @@ impl Drk {
 
             println!("[mark_tx_spend] Found Money contract in call {i}");
             let nullifiers = self.money_call_nullifiers(call).await?;
-            self.mark_spent_coins(&nullifiers, &tx_hash).await?;
+            self.mark_spent_coins(&owncoins_nullifiers, &nullifiers, &tx_hash)?;
         }
 
         Ok(())
     }
 
-    /// Mark a coin in the wallet as spent, and store its inverse query into the cache.
-    pub async fn mark_spent_coin(&self, coin: &Coin, spent_tx_hash: &String) -> WalletDbResult<()> {
-        // Grab coin record key
-        let key = serialize_async(&coin.inner()).await;
+    /// Marks all coins in the wallet as spent, if their nullifier is in the given set.
+    /// Returns a flag indicating if any of the provided nullifiers refer to our own wallet.
+    pub fn mark_spent_coins(
+        &self,
+        owncoins_nullifiers: &BTreeMap<[u8; 32], [u8; 32]>,
+        nullifiers: &[Nullifier],
+        spent_tx_hash: &String,
+    ) -> Result<bool> {
+        if nullifiers.is_empty() {
+            return Ok(false)
+        }
+
+        // Find our owncoins that where spent
+        let mut spent_owncoins = Vec::new();
+        for nullifier in nullifiers {
+            if let Some(coin_key) = owncoins_nullifiers.get(&nullifier.to_bytes()) {
+                spent_owncoins.push(coin_key);
+            }
+        }
+        if spent_owncoins.is_empty() {
+            return Ok(false)
+        }
 
         // Create an SQL `UPDATE` query to mark rows as spent(1)
         let query = format!(
@@ -1054,76 +1108,51 @@ impl Drk {
         );
 
         // Create its inverse query
-        let inverse = self.wallet.create_prepared_statement(
-            &format!(
-                "UPDATE {} SET {} = 0, {} = '-' WHERE {} = ?1;",
-                *MONEY_COINS_TABLE,
-                MONEY_COINS_COL_IS_SPENT,
-                MONEY_COINS_COL_SPENT_TX_HASH,
-                MONEY_COINS_COL_COIN
-            ),
-            rusqlite::params![key],
-        )?;
-
-        // Execute the query
-        self.wallet.exec_sql(&query, rusqlite::params![spent_tx_hash, key])?;
-
-        // Store its inverse
-        self.wallet.cache_inverse(inverse)
-    }
+        let inverse_query = format!(
+            "UPDATE {} SET {} = 0, {} = '-' WHERE {} = ?1;",
+            *MONEY_COINS_TABLE,
+            MONEY_COINS_COL_IS_SPENT,
+            MONEY_COINS_COL_SPENT_TX_HASH,
+            MONEY_COINS_COL_COIN
+        );
 
-    /// Marks all coins in the wallet as spent, if their nullifier is in the given set.
-    /// Returns a flag indicating if any of the provided nullifiers refer to our own wallet.
-    pub async fn mark_spent_coins(
-        &self,
-        nullifiers: &[Nullifier],
-        spent_tx_hash: &String,
-    ) -> Result<bool> {
-        if nullifiers.is_empty() {
-            return Ok(false)
-        }
+        // Mark spent own coins
+        for ownoin in spent_owncoins {
+            // Create its inverse query
+            let inverse = match self
+                .wallet
+                .create_prepared_statement(&inverse_query, rusqlite::params![ownoin])
+            {
+                Ok(i) => i,
+                Err(e) => {
+                    return Err(Error::DatabaseError(format!(
+                        "[mark_spent_coins] Creating inverse query failed: {e:?}"
+                    )))
+                }
+            };
 
-        // First we remark transaction spent coins
-        let mut wallet_spent_coins = false;
-        for coin in self.get_transaction_coins(spent_tx_hash).await? {
-            if let Err(e) = self.mark_spent_coin(&coin.coin, spent_tx_hash).await {
+            // Execute the query
+            if let Err(e) = self.wallet.exec_sql(&query, rusqlite::params![spent_tx_hash, ownoin]) {
                 return Err(Error::DatabaseError(format!(
                     "[mark_spent_coins] Marking spent coin failed: {e:?}"
                 )))
             }
-            wallet_spent_coins = true;
-        }
 
-        // Then we mark transaction unspent coins
-        for (coin, _, _) in self.get_coins(false).await? {
-            if !nullifiers.contains(&coin.nullifier()) {
-                continue
-            }
-            if let Err(e) = self.mark_spent_coin(&coin.coin, spent_tx_hash).await {
+            // Store its inverse
+            if let Err(e) = self.wallet.cache_inverse(inverse) {
                 return Err(Error::DatabaseError(format!(
-                    "[mark_spent_coins] Marking spent coin failed: {e:?}"
+                    "[mark_spent_coins] Storing inverse query failed: {e:?}"
                 )))
             }
-            wallet_spent_coins = true;
         }
 
-        Ok(wallet_spent_coins)
+        Ok(true)
     }
 
     /// Inserts given slice to the wallets nullifiers Sparse Merkle Tree.
-    pub fn smt_insert(&self, nullifiers: &[Nullifier]) -> Result<()> {
-        let store = WalletStorage::new(
-            &self.wallet,
-            &MONEY_SMT_TABLE,
-            MONEY_SMT_COL_KEY,
-            MONEY_SMT_COL_VALUE,
-        );
-        let mut smt = WalletSmt::new(store, PoseidonFp::new(), &EMPTY_NODES_FP);
-
+    pub fn smt_insert(&self, smt: &mut CacheSmt, nullifiers: &[Nullifier]) -> Result<()> {
         let leaves: Vec<_> = nullifiers.iter().map(|x| (x.inner(), x.inner())).collect();
-        smt.insert_batch(leaves)?;
-
-        Ok(())
+        Ok(smt.insert_batch(leaves)?)
     }
 
     /// Reset the Money Merkle tree in the wallet.

bin/drk/src/rpc.rs

@@ -16,7 +16,11 @@
  * along with this program.  If not, see <https://www.gnu.org/licenses/>.
  */
 
-use std::{sync::Arc, time::Instant};
+use std::{
+    collections::{BTreeMap, HashMap},
+    sync::Arc,
+    time::Instant,
+};
 
 use url::Url;
 
@@ -32,18 +36,85 @@ use darkfi::{
     util::encoding::base64,
     Error, Result,
 };
+use darkfi_dao_contract::model::{DaoBulla, DaoProposalBulla};
 use darkfi_sdk::{
-    crypto::{ContractId, DAO_CONTRACT_ID, DEPLOYOOOR_CONTRACT_ID, MONEY_CONTRACT_ID},
+    crypto::{
+        smt::{PoseidonFp, EMPTY_NODES_FP},
+        ContractId, MerkleTree, SecretKey, DAO_CONTRACT_ID, DEPLOYOOOR_CONTRACT_ID,
+        MONEY_CONTRACT_ID,
+    },
     tx::TransactionHash,
 };
 use darkfi_serial::{deserialize_async, serialize_async};
 
 use crate::{
+    cache::{CacheOverlay, CacheSmt, CacheSmtStorage, SLED_MONEY_SMT_TREE},
+    dao::{SLED_MERKLE_TREES_DAO_DAOS, SLED_MERKLE_TREES_DAO_PROPOSALS},
     error::{WalletDbError, WalletDbResult},
+    money::SLED_MERKLE_TREES_MONEY,
     Drk,
 };
 
+/// Auxiliary structure holding various in memory caches to use during scan
+pub struct ScanCache {
+    /// The Money Merkle tree containing coins
+    pub money_tree: MerkleTree,
+    /// The Money Sparse Merkle tree containing coins nullifiers
+    pub money_smt: CacheSmt,
+    /// All our known secrets to decrypt coin notes
+    pub notes_secrets: Vec<SecretKey>,
+    /// Our own coins nullifiers
+    pub owncoins_nullifiers: BTreeMap<[u8; 32], [u8; 32]>,
+    /// The DAO Merkle tree containing DAO bullas
+    pub dao_daos_tree: MerkleTree,
+    /// The DAO Merkle tree containing proposals bullas
+    pub dao_proposals_tree: MerkleTree,
+    /// Our own DAOs with their proposals and votes keys
+    pub own_daos: HashMap<DaoBulla, (Option<SecretKey>, Option<SecretKey>)>,
+    /// Our own DAOs proposals with their corresponding DAO reference
+    pub own_proposals: HashMap<DaoProposalBulla, DaoBulla>,
+}
+
 impl Drk {
+    /// Auxiliarry function to generate a new [`ScanCache`] for the
+    /// wallet.
+    pub async fn scan_cache(&self) -> Result<ScanCache> {
+        let money_tree = self.get_money_tree().await?;
+        let smt_store = CacheSmtStorage::new(CacheOverlay::new(&self.cache)?, SLED_MONEY_SMT_TREE);
+        let money_smt = CacheSmt::new(smt_store, PoseidonFp::new(), &EMPTY_NODES_FP);
+        let mut notes_secrets = self.get_money_secrets().await?;
+        let mut owncoins_nullifiers = BTreeMap::new();
+        for coin in self.get_coins(true).await? {
+            owncoins_nullifiers.insert(coin.0.nullifier().to_bytes(), coin.0.coin.to_bytes());
+        }
+        let (dao_daos_tree, dao_proposals_tree) = self.get_dao_trees().await?;
+        let mut own_daos = HashMap::new();
+        for dao in self.get_daos().await? {
+            own_daos.insert(
+                dao.bulla(),
+                (dao.params.proposals_secret_key, dao.params.votes_secret_key),
+            );
+            if let Some(secret_key) = dao.params.notes_secret_key {
+                notes_secrets.push(secret_key);
+            }
+        }
+        let mut own_proposals = HashMap::new();
+        for proposal in self.get_proposals().await? {
+            own_proposals.insert(proposal.bulla(), proposal.proposal.dao_bulla);
+        }
+
+        Ok(ScanCache {
+            money_tree,
+            money_smt,
+            notes_secrets,
+            owncoins_nullifiers,
+            dao_daos_tree,
+            dao_proposals_tree,
+            own_daos,
+            own_proposals,
+        })
+    }
+
     /// Subscribes to darkfid's JSON-RPC notification endpoint that serves
     /// new confirmed blocks. Upon receiving them, all the transactions are
     /// scanned and we check if any of them call the money contract, and if
@@ -175,7 +246,9 @@ impl Drk {
                                         )))
                                     }
                                 };
-                                if let Err(e) = self.scan_block(&genesis).await {
+                                if let Err(e) =
+                                    self.scan_block(&mut self.scan_cache().await?, &genesis).await
+                                {
                                     return Err(Error::DatabaseError(format!(
                                         "[subscribe_blocks] Scanning block failed: {e:?}"
                                     )))
@@ -183,7 +256,8 @@ impl Drk {
                             }
                         }
 
-                        if let Err(e) = self.scan_block(&block).await {
+                        if let Err(e) = self.scan_block(&mut self.scan_cache().await?, &block).await
+                        {
                             return Err(Error::DatabaseError(format!(
                                 "[subscribe_blocks] Scanning block failed: {e:?}"
                             )))
@@ -214,12 +288,18 @@ impl Drk {
     /// `scan_block` will go over over transactions in a block and handle their calls
     /// based on the called contract. Additionally, will update `last_scanned_block` to
     /// the provided block height and will store its height, hash and inverse query.
-    async fn scan_block(&self, block: &BlockInfo) -> Result<()> {
+    async fn scan_block(&self, scan_cache: &mut ScanCache, block: &BlockInfo) -> Result<()> {
         // Reset wallet inverse cache state
         self.reset_inverse_cache().await?;
 
-        // Keep track of our wallet transactions
+        // Keep track of the trees we need to update and our wallet
+        // transactions.
+        let mut update_money_tree = false;
+        let mut update_dao_daos_tree = false;
+        let mut update_dao_proposals_tree = false;
         let mut wallet_txs = vec![];
+
+        // Scan the block
         println!("=======================================");
         println!("{}", block.header);
         println!("=======================================");
@@ -232,17 +312,32 @@ impl Drk {
             for (i, call) in tx.calls.iter().enumerate() {
                 if call.data.contract_id == *MONEY_CONTRACT_ID {
                     println!("[scan_block] Found Money contract in call {i}");
-                    if self.apply_tx_money_data(i, &tx.calls, &tx_hash_string).await? {
+                    let (update_tree, own_tx) = self
+                        .apply_tx_money_data(scan_cache, &i, &tx.calls, &tx_hash_string)
+                        .await?;
+                    if update_tree {
+                        update_money_tree = true;
+                    }
+                    if own_tx {
                         wallet_tx = true;
-                    };
+                    }
                     continue
                 }
 
                 if call.data.contract_id == *DAO_CONTRACT_ID {
                     println!("[scan_block] Found DAO contract in call {i}");
-                    if self.apply_tx_dao_data(&call.data.data, tx_hash, i as u8).await? {
+                    let (update_daos_tree, update_proposals_tree, own_tx) = self
+                        .apply_tx_dao_data(scan_cache, &call.data.data, &tx_hash, &(i as u8))
+                        .await?;
+                    if update_daos_tree {
+                        update_dao_daos_tree = true;
+                    }
+                    if update_proposals_tree {
+                        update_dao_proposals_tree = true;
+                    }
+                    if own_tx {
                         wallet_tx = true;
-                    };
+                    }
                     continue
                 }
 
@@ -262,6 +357,57 @@ impl Drk {
             }
         }
 
+        // Update money merkle tree, if needed
+        if update_money_tree {
+            scan_cache
+                .money_smt
+                .store
+                .overlay
+                .insert_merkle_tree(SLED_MERKLE_TREES_MONEY, &scan_cache.money_tree)?;
+        }
+
+        // Update dao daos merkle tree, if needed
+        if update_dao_daos_tree {
+            scan_cache
+                .money_smt
+                .store
+                .overlay
+                .insert_merkle_tree(SLED_MERKLE_TREES_DAO_DAOS, &scan_cache.dao_daos_tree)?;
+        }
+
+        // Update dao proposals merkle tree, if needed
+        if update_dao_proposals_tree {
+            scan_cache.money_smt.store.overlay.insert_merkle_tree(
+                SLED_MERKLE_TREES_DAO_PROPOSALS,
+                &scan_cache.dao_proposals_tree,
+            )?;
+        }
+
+        // Insert the block record
+        scan_cache
+            .money_smt
+            .store
+            .overlay
+            .insert_scanned_block(&block.header.height, &block.header.hash())?;
+
+        // Grab the overlay current diff
+        let diff = scan_cache.money_smt.store.overlay.0.diff(&[])?;
+
+        // Insert the state inverse diff record
+        scan_cache
+            .money_smt
+            .store
+            .overlay
+            .insert_state_inverse_diff(&block.header.height, &diff.inverse())?;
+
+        // Apply the overlay current changes
+        scan_cache
+            .money_smt
+            .store
+            .overlay
+            .0
+            .apply_diff(&scan_cache.money_smt.store.overlay.0.diff(&[])?)?;
+
         // Update wallet transactions records
         if let Err(e) = self.put_tx_history_records(&wallet_txs, "Confirmed").await {
             return Err(Error::DatabaseError(format!(
@@ -334,6 +480,15 @@ impl Drk {
             height += 1;
         }
 
+        // Generate a new scan cache
+        let mut scan_cache = match self.scan_cache().await {
+            Ok(c) => c,
+            Err(e) => {
+                eprintln!("[scan_blocks] Generating scan cache failed: {e:?}");
+                return Err(WalletDbError::GenericError)
+            }
+        };
+
         loop {
             // Grab last confirmed block
             println!("Requested to scan from block number: {height}");
@@ -361,7 +516,7 @@ impl Drk {
                     }
                 };
                 println!("Block {height} received! Scanning block...");
-                if let Err(e) = self.scan_block(&block).await {
+                if let Err(e) = self.scan_block(&mut scan_cache, &block).await {
                     eprintln!("[scan_blocks] Scan block failed: {e:?}");
                     return Err(WalletDbError::GenericError)
                 };

bin/drk/src/walletdb.rs

@@ -21,16 +21,7 @@ use std::{
     sync::{Arc, Mutex},
 };
 
-use darkfi_sdk::{
-    crypto::{
-        pasta_prelude::PrimeField,
-        smt::{PoseidonFp, SparseMerkleTree, StorageAdapter, SMT_FP_DEPTH},
-    },
-    error::{ContractError, ContractResult},
-    pasta::pallas,
-};
 use log::{debug, error};
-use num_bigint::BigUint;
 use rusqlite::{
     types::{ToSql, Value},
     Connection,
@@ -411,187 +402,11 @@ macro_rules! convert_named_params {
     };
 }
 
-/// Wallet SMT definition
-pub type WalletSmt<'a> = SparseMerkleTree<
-    'static,
-    SMT_FP_DEPTH,
-    { SMT_FP_DEPTH + 1 },
-    pallas::Base,
-    PoseidonFp,
-    WalletStorage<'a>,
->;
-
-/// An SMT adapter for wallet SQLite database storage.
-pub struct WalletStorage<'a> {
-    wallet: &'a WalletPtr,
-    table: &'a str,
-    key_col: &'a str,
-    value_col: &'a str,
-}
-
-impl<'a> WalletStorage<'a> {
-    pub fn new(
-        wallet: &'a WalletPtr,
-        table: &'a str,
-        key_col: &'a str,
-        value_col: &'a str,
-    ) -> Self {
-        Self { wallet, table, key_col, value_col }
-    }
-}
-
-impl StorageAdapter for WalletStorage<'_> {
-    type Value = pallas::Base;
-
-    fn put(&mut self, key: BigUint, value: pallas::Base) -> ContractResult {
-        // Check if record already exists to create the corresponding query,
-        // its param and its inverse.
-        let (query, params, inverse) = match self.get(&key) {
-            Some(v) => {
-                // Create an SQL `UPDATE` query
-                let q = format!(
-                    "UPDATE {} SET {} = ?1 WHERE {} = ?2;",
-                    self.table, self.value_col, self.key_col
-                );
-
-                // Create its inverse query
-                let i = match self.wallet.create_prepared_statement(
-                    &format!(
-                        "UPDATE {} SET {} = ?1 WHERE {} = ?2;",
-                        self.table, self.value_col, self.key_col
-                    ),
-                    rusqlite::params![v.to_repr(), key.to_bytes_le()],
-                ) {
-                    Ok(i) => i,
-                    Err(e) => {
-                        error!(target: "walletdb::StorageAdapter::put", "Creating inverse query for key {key:?} failed: {e:?}");
-                        return Err(ContractError::SmtPutFailed)
-                    }
-                };
-
-                (q, rusqlite::params![value.to_repr(), key.to_bytes_le()], i)
-            }
-            None => {
-                // Create an SQL `INSERT` query
-                let q = format!(
-                    "INSERT INTO {} ({}, {}) VALUES (?1, ?2);",
-                    self.table, self.key_col, self.value_col
-                );
-
-                // Create its inverse query
-                let i = match self.wallet.create_prepared_statement(
-                    &format!("DELETE FROM {} WHERE {} = ?1;", self.table, self.key_col),
-                    rusqlite::params![key.to_bytes_le()],
-                ) {
-                    Ok(i) => i,
-                    Err(e) => {
-                        error!(target: "walletdb::StorageAdapter::put", "Creating inverse query for key {key:?} failed: {e:?}");
-                        return Err(ContractError::SmtPutFailed)
-                    }
-                };
-
-                (q, rusqlite::params![key.to_bytes_le(), value.to_repr()], i)
-            }
-        };
-
-        // Execute the query
-        if let Err(e) = self.wallet.exec_sql(&query, params) {
-            error!(target: "walletdb::StorageAdapter::put", "Inserting key {key:?}, value {value:?} into DB failed: {e:?}");
-            return Err(ContractError::SmtPutFailed)
-        }
-
-        // Store its inverse
-        if let Err(e) = self.wallet.cache_inverse(inverse) {
-            error!(target: "walletdb::StorageAdapter::put", "Inserting inverse query into cache failed: {e:?}");
-            return Err(ContractError::SmtPutFailed)
-        }
-
-        Ok(())
-    }
-
-    fn get(&self, key: &BigUint) -> Option<pallas::Base> {
-        let row = match self.wallet.query_single(
-            self.table,
-            &[self.value_col],
-            convert_named_params! {(self.key_col, key.to_bytes_le())},
-        ) {
-            Ok(r) => r,
-            Err(WalletDbError::RowNotFound) => return None,
-            Err(e) => {
-                error!(target: "walletdb::StorageAdapter::get", "Fetching key {key:?} from DB failed: {e:?}");
-                return None
-            }
-        };
-
-        let Value::Blob(ref value_bytes) = row[0] else {
-            error!(target: "walletdb::StorageAdapter::get", "Parsing key {key:?} value bytes");
-            return None
-        };
-
-        let mut repr = [0; 32];
-        repr.copy_from_slice(value_bytes);
-
-        pallas::Base::from_repr(repr).into()
-    }
-
-    fn del(&mut self, key: &BigUint) -> ContractResult {
-        // Check if record already exists to create the corresponding query,
-        // its param and its inverse.
-        let (query, params, inverse) = match self.get(key) {
-            Some(value) => {
-                // Create an SQL `DELETE` query
-                let q = format!("DELETE FROM {} WHERE {} = ?1;", self.table, self.key_col);
-
-                // Create its inverse query
-                let i = match self.wallet.create_prepared_statement(
-                    &format!(
-                        "INSERT INTO {} ({}, {}) VALUES (?1, ?2);",
-                        self.table, self.key_col, self.value_col
-                    ),
-                    rusqlite::params![key.to_bytes_le(), value.to_repr()],
-                ) {
-                    Ok(i) => i,
-                    Err(e) => {
-                        error!(target: "walletdb::StorageAdapter::del", "Creating inverse query for key {key:?} failed: {e:?}");
-                        return Err(ContractError::SmtDelFailed)
-                    }
-                };
-
-                (q, rusqlite::params![key.to_bytes_le()], i)
-            }
-            None => {
-                // If record doesn't exist do nothing
-                return Ok(())
-            }
-        };
-
-        // Execute the query
-        if let Err(e) = self.wallet.exec_sql(&query, params) {
-            error!(target: "walletdb::StorageAdapter::del", "Removing key {key:?} from DB failed: {e:?}");
-            return Err(ContractError::SmtDelFailed)
-        }
-
-        // Store its inverse
-        if let Err(e) = self.wallet.cache_inverse(inverse) {
-            error!(target: "walletdb::StorageAdapter::del", "Inserting inverse query into cache failed: {e:?}");
-            return Err(ContractError::SmtDelFailed)
-        }
-
-        Ok(())
-    }
-}
-
 #[cfg(test)]
 mod tests {
-    use darkfi::zk::halo2::Field;
-    use darkfi_sdk::{
-        crypto::smt::{gen_empty_nodes, util::FieldHasher, PoseidonFp, SparseMerkleTree},
-        pasta::pallas,
-    };
-    use rand::rngs::OsRng;
     use rusqlite::types::Value;
 
-    use crate::walletdb::{WalletDb, WalletStorage};
+    use crate::walletdb::WalletDb;
 
     #[test]
     fn test_mem_wallet() {
@@ -735,78 +550,4 @@ mod tests {
             assert_eq!(row[0], Value::Blob(gae.clone()));
         }
     }
-
-    #[test]
-    fn test_sqlite_smt() {
-        // Setup SQLite database
-        let table = &"smt";
-        let key_col = &"smt_key";
-        let value_col = &"smt_value";
-        let wallet = WalletDb::new(None, None).unwrap();
-        wallet.exec_batch_sql(&format!("CREATE TABLE {table} ( {key_col} BLOB INTEGER PRIMARY KEY NOT NULL, {value_col} BLOB NOT NULL);")).unwrap();
-
-        // Setup SMT
-        const HEIGHT: usize = 3;
-        let hasher = PoseidonFp::new();
-        let empty_leaf = pallas::Base::ZERO;
-        let empty_nodes = gen_empty_nodes::<{ HEIGHT + 1 }, _, _>(&hasher, empty_leaf);
-        let store = WalletStorage::new(&wallet, table, key_col, value_col);
-        let mut smt = SparseMerkleTree::<HEIGHT, { HEIGHT + 1 }, _, _, _>::new(
-            store,
-            hasher.clone(),
-            &empty_nodes,
-        );
-
-        // Verify database is empty
-        let rows = wallet.query_multiple(table, &[key_col], &[]).unwrap();
-        assert!(rows.is_empty());
-
-        let leaves = vec![
-            (pallas::Base::from(1), pallas::Base::random(&mut OsRng)),
-            (pallas::Base::from(2), pallas::Base::random(&mut OsRng)),
-            (pallas::Base::from(3), pallas::Base::random(&mut OsRng)),
-        ];
-        smt.insert_batch(leaves.clone()).unwrap();
-
-        let hash1 = leaves[0].1;
-        let hash2 = leaves[1].1;
-        let hash3 = leaves[2].1;
-
-        let hash = |l, r| hasher.hash([l, r]);
-
-        let hash01 = hash(empty_nodes[3], hash1);
-        let hash23 = hash(hash2, hash3);
-
-        let hash0123 = hash(hash01, hash23);
-        let root = hash(hash0123, empty_nodes[1]);
-        assert_eq!(root, smt.root());
-
-        // Now try to construct a membership proof for leaf 3
-        let pos = leaves[2].0;
-        let path = smt.prove_membership(&pos);
-        assert_eq!(path.path[0], empty_nodes[1]);
-        assert_eq!(path.path[1], hash01);
-        assert_eq!(path.path[2], hash2);
-
-        assert_eq!(hash23, hash(path.path[2], hash3));
-        assert_eq!(hash0123, hash(path.path[1], hash(path.path[2], hash3)));
-        assert_eq!(root, hash(hash(path.path[1], hash(path.path[2], hash3)), path.path[0]));
-
-        assert!(path.verify(&root, &hash3, &pos));
-
-        // Verify database contains keys
-        let rows = wallet.query_multiple(table, &[key_col], &[]).unwrap();
-        assert!(!rows.is_empty());
-
-        // We are now going to rollback the wallet changes
-        let rollback_query = wallet.grab_inverse_cache_block().unwrap();
-        wallet.exec_batch_sql(&rollback_query).unwrap();
-
-        // Clear cache
-        wallet.clear_inverse_cache().unwrap();
-
-        // Verify database is empty again
-        let rows = wallet.query_multiple(table, &[key_col], &[]).unwrap();
-        assert!(rows.is_empty());
-    }
 }

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

@@ -150,7 +150,7 @@ pub struct SparseMerkleTree<
     S: StorageAdapter<Value = F>,
 > {
     /// A map from leaf indices to leaf data stored as field elements.
-    store: S,
+    pub store: S,
     /// The hasher used to build the Merkle tree.
     hasher: H,
     /// An array of empty hashes hashed with themselves `N` times.