feat!: Persist utxo lock status

- Add `Wallet::lock_unspent`
- Add `Wallet::unlock_unspent`
- Add `Wallet::is_utxo_locked`
- Add pub struct `UtxoInfo`
- Extend ChangeSet with new member `utxo_info` with type
  `BTreeMap<OutPoint, UtxoInfo>`
- Test `test_lock_unspent_persist`

Author: ValuedMammal

wallet/src/wallet/changeset.rs

@@ -1,12 +1,17 @@
+use alloc::collections::BTreeMap;
+
 use bdk_chain::{
     indexed_tx_graph, keychain_txout, local_chain, tx_graph, ConfirmationBlockTime, Merge,
 };
+use bitcoin::{OutPoint, Txid};
 use miniscript::{Descriptor, DescriptorPublicKey};
 use serde::{Deserialize, Serialize};
 
 type IndexedTxGraphChangeSet =
     indexed_tx_graph::ChangeSet<ConfirmationBlockTime, keychain_txout::ChangeSet>;
 
+use crate::UtxoInfo;
+
 /// A change set for [`Wallet`]
 ///
 /// ## Definition
@@ -114,6 +119,8 @@ pub struct ChangeSet {
     pub tx_graph: tx_graph::ChangeSet<ConfirmationBlockTime>,
     /// Changes to [`KeychainTxOutIndex`](keychain_txout::KeychainTxOutIndex).
     pub indexer: keychain_txout::ChangeSet,
+    /// Changes to utxo info.
+    pub utxo_info: BTreeMap<OutPoint, UtxoInfo>,
 }
 
 impl Merge for ChangeSet {
@@ -142,6 +149,11 @@ impl Merge for ChangeSet {
             self.network = other.network;
         }
 
+        // To merge utxo_info we just extend the existing collection. If there's
+        // an existing entry for a given outpoint, then it is overwritten by the
+        // new utxo info.
+        self.utxo_info.extend(other.utxo_info);
+
         Merge::merge(&mut self.local_chain, other.local_chain);
         Merge::merge(&mut self.tx_graph, other.tx_graph);
         Merge::merge(&mut self.indexer, other.indexer);
@@ -154,6 +166,7 @@ impl Merge for ChangeSet {
             && self.local_chain.is_empty()
             && self.tx_graph.is_empty()
             && self.indexer.is_empty()
+            && self.utxo_info.is_empty()
     }
 }
 
@@ -163,6 +176,8 @@ impl ChangeSet {
     pub const WALLET_SCHEMA_NAME: &'static str = "bdk_wallet";
     /// Name of table to store wallet descriptors and network.
     pub const WALLET_TABLE_NAME: &'static str = "bdk_wallet";
+    /// Name of table to store wallet utxo info.
+    pub const WALLET_UTXO_TABLE_NAME: &'static str = "bdk_wallet_utxo_info";
 
     /// Get v0 sqlite [ChangeSet] schema
     pub fn schema_v0() -> alloc::string::String {
@@ -177,12 +192,25 @@ impl ChangeSet {
         )
     }
 
+    /// Get v1 sqlite [ChangeSet] schema. Schema v1 adds a table for wallet UTXO info.
+    pub fn schema_v1() -> alloc::string::String {
+        format!(
+            "CREATE TABLE {} ( \
+                txid TEXT NOT NULL, \
+                vout INTEGER NOT NULL, \
+                is_locked INTEGER, \
+                PRIMARY KEY(txid, vout) \
+                ) STRICT;",
+            Self::WALLET_UTXO_TABLE_NAME,
+        )
+    }
+
     /// Initialize sqlite tables for wallet tables.
     pub fn init_sqlite_tables(db_tx: &chain::rusqlite::Transaction) -> chain::rusqlite::Result<()> {
         crate::rusqlite_impl::migrate_schema(
             db_tx,
             Self::WALLET_SCHEMA_NAME,
-            &[&Self::schema_v0()],
+            &[&Self::schema_v0(), &Self::schema_v1()],
         )?;
 
         bdk_chain::local_chain::ChangeSet::init_sqlite_tables(db_tx)?;
@@ -220,6 +248,27 @@ impl ChangeSet {
             changeset.network = network.map(Impl::into_inner);
         }
 
+        // Load utxo info.
+        let mut stmt = db_tx.prepare(&format!(
+            "SELECT txid, vout, is_locked FROM {}",
+            Self::WALLET_UTXO_TABLE_NAME,
+        ))?;
+        let rows = stmt.query_map([], |row| {
+            Ok((
+                row.get::<_, Impl<Txid>>("txid")?,
+                row.get::<_, u32>("vout")?,
+                row.get::<_, bool>("is_locked")?,
+            ))
+        })?;
+        for row in rows {
+            let (Impl(txid), vout, is_locked) = row?;
+            let utxo_info = UtxoInfo {
+                outpoint: OutPoint::new(txid, vout),
+                is_locked,
+            };
+            changeset.utxo_info.insert(utxo_info.outpoint, utxo_info);
+        }
+
         changeset.local_chain = local_chain::ChangeSet::from_sqlite(db_tx)?;
         changeset.tx_graph = tx_graph::ChangeSet::<_>::from_sqlite(db_tx)?;
         changeset.indexer = keychain_txout::ChangeSet::from_sqlite(db_tx)?;
@@ -268,6 +317,20 @@ impl ChangeSet {
             })?;
         }
 
+        // Persist utxo info.
+        let mut stmt = db_tx.prepare_cached(&format!(
+            "INSERT INTO {}(txid, vout, is_locked) VALUES(:txid, :vout, :is_locked) ON CONFLICT DO UPDATE SET is_locked = :is_locked",
+            Self::WALLET_UTXO_TABLE_NAME,
+        ))?;
+        for (&outpoint, utxo_info) in &self.utxo_info {
+            let OutPoint { txid, vout } = outpoint;
+            stmt.execute(named_params! {
+                ":txid": Impl(txid),
+                ":vout": vout,
+                ":is_locked": utxo_info.is_locked,
+            })?;
+        }
+
         self.local_chain.persist_to_sqlite(db_tx)?;
         self.tx_graph.persist_to_sqlite(db_tx)?;
         self.indexer.persist_to_sqlite(db_tx)?;

wallet/src/wallet/mod.rs

@@ -109,6 +109,7 @@ pub struct Wallet {
     stage: ChangeSet,
     network: Network,
     secp: SecpCtx,
+    utxo_info: BTreeMap<OutPoint, UtxoInfo>,
 }
 
 /// An update to [`Wallet`].
@@ -449,6 +450,7 @@ impl Wallet {
         let change_descriptor = index.get_descriptor(KeychainKind::Internal).cloned();
         let indexed_graph = IndexedTxGraph::new(index);
         let indexed_graph_changeset = indexed_graph.initial_changeset();
+        let utxo_info = BTreeMap::new();
 
         let stage = ChangeSet {
             descriptor,
@@ -457,6 +459,7 @@ impl Wallet {
             tx_graph: indexed_graph_changeset.tx_graph,
             indexer: indexed_graph_changeset.indexer,
             network: Some(network),
+            ..Default::default()
         };
 
         Ok(Wallet {
@@ -467,6 +470,7 @@ impl Wallet {
             indexed_graph,
             stage,
             secp,
+            utxo_info,
         })
     }
 
@@ -653,18 +657,39 @@ impl Wallet {
         let mut indexed_graph = IndexedTxGraph::new(index);
         indexed_graph.apply_changeset(changeset.indexer.into());
         indexed_graph.apply_changeset(changeset.tx_graph.into());
+        let utxo_info = BTreeMap::new();
 
         let stage = ChangeSet::default();
 
-        Ok(Some(Wallet {
+        let mut wallet = Wallet {
             signers,
             change_signers,
             chain,
             indexed_graph,
             stage,
             network,
             secp,
-        }))
+            utxo_info,
+        };
+
+        // Apply lock status to wallet utxos.
+        wallet.utxo_info = wallet
+            .list_unspent()
+            .map(|output| {
+                (
+                    output.outpoint,
+                    UtxoInfo {
+                        outpoint: output.outpoint,
+                        is_locked: false,
+                    },
+                )
+            })
+            .collect();
+        for (outpoint, utxo_info) in changeset.utxo_info {
+            wallet.utxo_info.insert(outpoint, utxo_info);
+        }
+
+        Ok(Some(wallet))
     }
 
     /// Get the Bitcoin network the wallet is using.
@@ -2110,6 +2135,8 @@ impl Wallet {
                     CanonicalizationParams::default(),
                     self.indexed_graph.index.outpoints().iter().cloned(),
                 )
+                // Filter out locked utxos
+                .filter(|(_, txo)| self.is_utxo_locked(txo.outpoint) != Some(true))
                 // only create LocalOutput if UTxO is mature
                 .filter_map(move |((k, i), full_txo)| {
                     full_txo
@@ -2377,6 +2404,84 @@ impl Wallet {
         &self.chain
     }
 
+    /// Is UTXO locked. `None` if the outpoint isn't found in `self.utxo_info`.
+    pub fn is_utxo_locked(&self, outpoint: OutPoint) -> Option<bool> {
+        Some(self.utxo_info.get(&outpoint)?.is_locked)
+    }
+
+    /// Lock an unspent output by `outpoint`.
+    ///
+    /// Locked utxos are automatically filtered out during coin selection. You need to persist
+    /// the wallet in order for the lock status to persist across restarts. To unlock a previously
+    /// locked outpoint, see [`Wallet::unlock_unspent`].
+    pub fn lock_unspent(&mut self, outpoint: OutPoint) {
+        use alloc::collections::btree_map;
+        let lock_value = true;
+
+        // If the utxo is not currently locked, update the lock value
+        // and stage the change.
+        let is_changed = match self.utxo_info.entry(outpoint) {
+            btree_map::Entry::Occupied(mut e) => {
+                let mut is_changed = false;
+
+                let utxo = e.get_mut();
+
+                if !utxo.is_locked {
+                    utxo.is_locked = lock_value;
+                    is_changed = true;
+                }
+
+                is_changed
+            }
+            btree_map::Entry::Vacant(e) => {
+                e.insert(UtxoInfo {
+                    outpoint,
+                    is_locked: lock_value,
+                });
+                true
+            }
+        };
+
+        if is_changed {
+            let utxo_info = UtxoInfo {
+                outpoint,
+                is_locked: lock_value,
+            };
+            self.stage.merge(ChangeSet {
+                utxo_info: [(outpoint, utxo_info)].into(),
+                ..Default::default()
+            });
+        }
+    }
+
+    /// Unlock unspent.
+    pub fn unlock_unspent(&mut self, outpoint: OutPoint) {
+        use alloc::collections::btree_map;
+        let lock_value = false;
+
+        match self.utxo_info.entry(outpoint) {
+            btree_map::Entry::Occupied(mut e) => {
+                let utxo = e.get_mut();
+
+                // If the utxo is currently locked, update the lock value and stage
+                // the change.
+                if utxo.is_locked {
+                    utxo.is_locked = lock_value;
+                    let utxo_info = UtxoInfo {
+                        outpoint,
+                        is_locked: lock_value,
+                    };
+                    self.stage.merge(ChangeSet {
+                        utxo_info: [(outpoint, utxo_info)].into(),
+                        ..Default::default()
+                    });
+                }
+            }
+            // If there is no entry, we're done because the utxo can't be locked.
+            btree_map::Entry::Vacant(..) => {}
+        }
+    }
+
     /// Introduces a `block` of `height` to the wallet, and tries to connect it to the
     /// `prev_blockhash` of the block's header.
     ///
@@ -2580,6 +2685,16 @@ impl AsRef<bdk_chain::tx_graph::TxGraph<ConfirmationBlockTime>> for Wallet {
     }
 }
 
+/// Information about a wallet UTXO.
+#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, serde::Serialize, serde::Deserialize)]
+pub struct UtxoInfo {
+    /// Outpoint.
+    pub outpoint: bitcoin::OutPoint,
+    /// Whether the outpoint is locked by the user. This doesn't take into account
+    /// any timelocks that may be part of the actual scriptPubKey.
+    pub is_locked: bool,
+}
+
 /// Deterministically generate a unique name given the descriptors defining the wallet
 ///
 /// Compatible with [`wallet_name_from_descriptor`]