feat(wallet-sdk): Mempool input(outgoing) scanning

wallet/cli/src/core/console.rs

@@ -134,21 +134,30 @@ pub async fn start_console(wallet: Wallet) {
                 let spendable_balance = wallet.spendable_balance().await;
                 let total_balance = wallet.total_balance().await;
                 let unconfirmed_balance = wallet.unconfirmed_balance().await;
-
-                info!(
-                    "Balance: {} NYKS ({} UTXOs; {} spendable, {} timelocked, {} unconfirmed).",
+                let outgoing_balance = wallet.outgoing_balance().await;
+
+                println!(
+                    "\n\
+                        Balance\n\
+                        ├─ Total:       {} NYKS\n\
+                        ├─ Spendable:   {} NYKS\n\
+                        ├─ Timelocked:  {} NYKS\n\
+                        ├─ Unconfirmed: {} NYKS\n\
+                        ├─ Outgoing:    {} NYKS\n\
+                        └─ UTXOs:       {}\n",
                     total_balance,
-                    utxo_count,
                     spendable_balance,
                     total_balance.checked_sub(&spendable_balance).unwrap(),
                     unconfirmed_balance,
+                    outgoing_balance,
+                    utxo_count,
                 );
             }
             Command::Address(key_type) => {
                 let key_type = key_type.unwrap_or(KeyType::Generation);
                 let address = wallet.address(key_type).await;
 
-                println!("\n{}", address.to_bech32m(wallet.network));
+                println!("\n{}\n", address.to_bech32m(wallet.network));
             }
             Command::Send {
                 recipient,

wallet/cli/src/core/storage.rs

@@ -5,8 +5,8 @@ use std::path::PathBuf;
 
 use nyks_consensus::block::block_height::BlockHeight;
 use nyks_standards::wallet::keys::key::KeyType;
-use nyks_wallet_sdk::state::utxos::utxo::MonitoredUtxo;
-use nyks_wallet_sdk::state::utxos::utxo::UtxoKey;
+use nyks_wallet_sdk::state::utxos::MonitoredUtxo;
+use nyks_wallet_sdk::state::utxos::UtxoKey;
 use serde::Deserialize;
 use serde::Serialize;
 

wallet/cli/src/main.rs

@@ -1,3 +1,5 @@
+pub mod core;
+
 use std::panic;
 use std::time::Duration;
 
@@ -14,8 +16,6 @@ use tracing_subscriber::EnvFilter;
 
 use crate::core::storage::Storage;
 
-pub mod core;
-
 #[derive(Parser)]
 #[command(name = "nyks-wallet")]
 #[command(about = "A nyks daemon wallet")]
@@ -96,6 +96,9 @@ async fn main() -> Result<()> {
 
                     storage.utxos.put(&key, &utxo);
                 }
+                WalletEvent::UtxosOutgoing { id, utxos } => {
+                    info!("{} UTXOs being spent on transaction {}.", utxos.len(), id);
+                }
             }
         }
 

wallet/sdk/src/scanners/chain.rs

@@ -19,7 +19,7 @@ use nyks_standards::wallet::keys::viewing_key::Decryptor;
 use nyks_standards::wallet::keys::viewing_key::ViewingKey;
 use thiserror::Error;
 
-use crate::state::utxos::utxo::IncomingUtxo;
+use crate::state::utxos::IncomingUtxo;
 
 #[derive(Debug, Copy, Clone, Error)]
 pub enum AdvanceError {

wallet/sdk/src/scanners/mempool.rs

@@ -1,15 +1,94 @@
-use nyks_standards::wallet::keys::viewing_key::ViewingKey;
+use std::collections::HashMap;
+use std::collections::HashSet;
 
+use nyks_consensus::transaction::transaction_kernel_id::TransactionKernelId;
+use nyks_rpc_client::block::transaction_kernel::RpcTransactionKernel;
+
+use crate::state::utxos::UtxoKey;
+use crate::state::utxos::index::UtxoIndex;
+
+/// Scans mempool transactions against the current UTXO pool.
 pub struct MempoolScanner {
-    keys: Vec<ViewingKey>,
+    index: UtxoIndex,
+    // IDs we've already checked, relevant or not - never rescanned.
+    checked_ids: HashSet<TransactionKernelId>,
+    // UTXOs currently observed as spent by relevant mempool transactions.
+    pub pending_spends: HashMap<UtxoKey, HashSet<TransactionKernelId>>, // Kept supporting multiple txs just in case
 }
 
 impl MempoolScanner {
-    pub fn new(keys: Vec<ViewingKey>) -> Self {
-        MempoolScanner { keys }
+    pub fn new(index: UtxoIndex) -> Self {
+        MempoolScanner {
+            index,
+            checked_ids: HashSet::new(),
+            pending_spends: HashMap::new(),
+        }
     }
 
-    pub fn add_key(&mut self, key: ViewingKey) {
-        self.keys.push(key);
+    /// Returns IDs that need to be fetched, i.e. ones we haven't checked yet.
+    pub async fn ids_to_fetch(
+        &self,
+        mempool_ids: &[TransactionKernelId],
+    ) -> Vec<TransactionKernelId> {
+        mempool_ids
+            .iter()
+            .filter(|id| !self.checked_ids.contains(id))
+            .cloned()
+            .collect()
+    }
+
+    /// Returns true if the given UTXO is currently being spent by a
+    /// transaction sitting in the mempool.
+    pub fn is_pending_spend(&self, utxo_key: &UtxoKey) -> bool {
+        self.pending_spends.contains_key(utxo_key)
+    }
+
+    /// Returns the UTXO keys currently observed as spent by a transaction
+    /// sitting in the mempool.
+    pub fn pending_spend_utxos(&self) -> impl Iterator<Item = &UtxoKey> + '_ {
+        self.pending_spends.keys()
+    }
+
+    /// Checks the transactions and marks them as checked. Each transaction
+    /// is scanned at most once ever; relevant ones are returned only here.
+    pub async fn scan(
+        &mut self,
+        transactions: Vec<(TransactionKernelId, RpcTransactionKernel)>,
+    ) -> Vec<(TransactionKernelId, Vec<UtxoKey>)> {
+        let mut relevant_transactions = Vec::new();
+
+        for (id, transaction) in &transactions {
+            let mut spent_inputs = Vec::new();
+
+            for input in &transaction.inputs {
+                if let Some(utxo_key) = self.index.get(&input.absolute_indices).await {
+                    self.pending_spends.entry(utxo_key).or_default().insert(*id);
+                    spent_inputs.push(utxo_key);
+                }
+            }
+
+            self.checked_ids.insert(*id);
+
+            if !spent_inputs.is_empty() {
+                relevant_transactions.push((*id, spent_inputs));
+            }
+        }
+
+        relevant_transactions
+    }
+
+    /// Removes transactions that are no longer in the mempool.
+    pub async fn evict_stale(&mut self, current_mempool_ids: Vec<TransactionKernelId>) {
+        let current_mempool_ids: HashSet<_> = current_mempool_ids.into_iter().collect();
+
+        self.checked_ids
+            .retain(|id| current_mempool_ids.contains(id));
+
+        // Drop stale tx ids from each UTXO's spender set, and drop the
+        // UTXO entry entirely once no live tx is spending it anymore.
+        self.pending_spends.retain(|_, spender_ids| {
+            spender_ids.retain(|id| current_mempool_ids.contains(id));
+            !spender_ids.is_empty()
+        });
     }
 }

wallet/sdk/src/state/utxos/index.rs

@@ -0,0 +1,30 @@
+use std::collections::HashMap;
+use std::sync::Arc;
+use tokio::sync::RwLock;
+
+use nyks_consensus::mutator_set::removal_record::absolute_index_set::AbsoluteIndexSet;
+
+use crate::state::utxos::UtxoKey;
+
+/// Cheap, independently-lockable index from a UTXO's absolute index set to
+/// its pool key. Kept in sync with `UtxoPool` at every insert/evict.
+#[derive(Clone, Default)]
+pub struct UtxoIndex(Arc<RwLock<HashMap<AbsoluteIndexSet, UtxoKey>>>);
+
+impl UtxoIndex {
+    pub fn new() -> Self {
+        UtxoIndex(Arc::new(RwLock::new(HashMap::new())))
+    }
+
+    pub async fn get(&self, idx: &AbsoluteIndexSet) -> Option<UtxoKey> {
+        self.0.read().await.get(idx).copied()
+    }
+
+    pub(crate) async fn insert(&self, idx: AbsoluteIndexSet, key: UtxoKey) {
+        self.0.write().await.insert(idx, key);
+    }
+
+    pub(crate) async fn remove(&self, idx: &AbsoluteIndexSet) {
+        self.0.write().await.remove(idx);
+    }
+}