feat!: Persist utxo lock status

clippy.toml

@@ -1,4 +1,4 @@
 # TODO fix, see: https://rust-lang.github.io/rust-clippy/master/index.html#large_enum_variant
-enum-variant-size-threshold = 1032
+enum-variant-size-threshold = 2048
 # TODO fix, see: https://rust-lang.github.io/rust-clippy/master/index.html#result_large_err
-large-error-threshold = 993
+large-error-threshold = 2048

wallet/src/wallet/changeset.rs

@@ -1,12 +1,15 @@
 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::locked_outpoints;
+
 /// A change set for [`Wallet`]
 ///
 /// ## Definition
@@ -114,6 +117,8 @@ pub struct ChangeSet {
     pub tx_graph: tx_graph::ChangeSet<ConfirmationBlockTime>,
     /// Changes to [`KeychainTxOutIndex`](keychain_txout::KeychainTxOutIndex).
     pub indexer: keychain_txout::ChangeSet,
+    /// Changes to locked outpoints.
+    pub locked_outpoints: locked_outpoints::ChangeSet,
 }
 
 impl Merge for ChangeSet {
@@ -142,6 +147,9 @@ impl Merge for ChangeSet {
             self.network = other.network;
         }
 
+        // merge locked outpoints
+        self.locked_outpoints.merge(other.locked_outpoints);
+
         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 +162,7 @@ impl Merge for ChangeSet {
             && self.local_chain.is_empty()
             && self.tx_graph.is_empty()
             && self.indexer.is_empty()
+            && self.locked_outpoints.is_empty()
     }
 }
 
@@ -163,6 +172,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 locked outpoints.
+    pub const WALLET_OUTPOINT_LOCK_TABLE_NAME: &'static str = "bdk_wallet_locked_outpoints";
 
     /// Get v0 sqlite [ChangeSet] schema
     pub fn schema_v0() -> alloc::string::String {
@@ -177,12 +188,26 @@ impl ChangeSet {
         )
     }
 
+    /// Get v1 sqlite [`ChangeSet`] schema. Schema v1 adds a table for locked outpoints.
+    pub fn schema_v1() -> alloc::string::String {
+        format!(
+            "CREATE TABLE {} ( \
+                txid TEXT NOT NULL, \
+                vout INTEGER NOT NULL, \
+                is_locked INTEGER, \
+                expiration_height INTEGER, \
+                PRIMARY KEY(txid, vout) \
+                ) STRICT;",
+            Self::WALLET_OUTPOINT_LOCK_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 +245,31 @@ impl ChangeSet {
             changeset.network = network.map(Impl::into_inner);
         }
 
+        // Select locked outpoints.
+        let mut stmt = db_tx.prepare(&format!(
+            "SELECT txid, vout, is_locked, expiration_height FROM {}",
+            Self::WALLET_OUTPOINT_LOCK_TABLE_NAME,
+        ))?;
+        let rows = stmt.query_map([], |row| {
+            Ok((
+                row.get::<_, Impl<Txid>>("txid")?,
+                row.get::<_, u32>("vout")?,
+                row.get::<_, bool>("is_locked")?,
+                row.get::<_, Option<u32>>("expiration_height")?,
+            ))
+        })?;
+        let locked_outpoints = &mut changeset.locked_outpoints;
+        for row in rows {
+            let (Impl(txid), vout, is_locked, expiration_height) = row?;
+            let outpoint = OutPoint::new(txid, vout);
+            locked_outpoints
+                .locked_outpoints
+                .insert(outpoint, is_locked);
+            locked_outpoints
+                .expiration_heights
+                .insert(outpoint, expiration_height);
+        }
+
         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 +318,35 @@ impl ChangeSet {
             })?;
         }
 
+        // Insert locked outpoints.
+        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_OUTPOINT_LOCK_TABLE_NAME,
+        ))?;
+        let locked_outpoints = &self.locked_outpoints.locked_outpoints;
+        for (&outpoint, is_locked) in locked_outpoints.iter() {
+            let OutPoint { txid, vout } = outpoint;
+            stmt.execute(named_params! {
+                ":txid": Impl(txid),
+                ":vout": vout,
+                ":is_locked": is_locked,
+            })?;
+        }
+        // Insert locked outpoints expiration heights.
+        let mut stmt = db_tx.prepare_cached(&format!(
+            "INSERT INTO {}(txid, vout, expiration_height) VALUES(:txid, :vout, :expiration_height) ON CONFLICT DO UPDATE SET expiration_height=:expiration_height",
+            Self::WALLET_OUTPOINT_LOCK_TABLE_NAME,
+        ))?;
+        let expiration_heights = &self.locked_outpoints.expiration_heights;
+        for (&outpoint, expiration_height) in expiration_heights.iter() {
+            let OutPoint { txid, vout } = outpoint;
+            stmt.execute(named_params! {
+                ":txid": Impl(txid),
+                ":vout": vout,
+                ":expiration_height": expiration_height,
+            })?;
+        }
+
         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/locked_outpoints.rs

@@ -0,0 +1,27 @@
+//! Module containing the locked outpoints change set.
+
+use bdk_chain::Merge;
+use bitcoin::OutPoint;
+use serde::{Deserialize, Serialize};
+
+use crate::collections::BTreeMap;
+
+/// Represents changes to locked outpoints.
+#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
+pub struct ChangeSet {
+    /// The lock status of an outpoint, `true == is_locked`.
+    pub locked_outpoints: BTreeMap<OutPoint, bool>,
+    /// The expiration height of the lock.
+    pub expiration_heights: BTreeMap<OutPoint, Option<u32>>,
+}
+
+impl Merge for ChangeSet {
+    fn merge(&mut self, other: Self) {
+        self.locked_outpoints.extend(other.locked_outpoints);
+        self.expiration_heights.extend(other.expiration_heights);
+    }
+
+    fn is_empty(&self) -> bool {
+        self.locked_outpoints.is_empty() && self.expiration_heights.is_empty()
+    }
+}

wallet/src/wallet/mod.rs

@@ -53,6 +53,7 @@ mod changeset;
 pub mod coin_selection;
 pub mod error;
 pub mod export;
+pub mod locked_outpoints;
 mod params;
 mod persisted;
 pub mod signer;
@@ -109,6 +110,7 @@ pub struct Wallet {
     stage: ChangeSet,
     network: Network,
     secp: SecpCtx,
+    locked_outpoints: BTreeMap<OutPoint, UtxoLock>,
 }
 
 /// An update to [`Wallet`].
@@ -449,6 +451,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 locked_outpoints = BTreeMap::new();
 
         let stage = ChangeSet {
             descriptor,
@@ -457,6 +460,7 @@ impl Wallet {
             tx_graph: indexed_graph_changeset.tx_graph,
             indexer: indexed_graph_changeset.indexer,
             network: Some(network),
+            ..Default::default()
         };
 
         Ok(Wallet {
@@ -467,6 +471,7 @@ impl Wallet {
             indexed_graph,
             stage,
             secp,
+            locked_outpoints,
         })
     }
 
@@ -654,6 +659,23 @@ impl Wallet {
         indexed_graph.apply_changeset(changeset.indexer.into());
         indexed_graph.apply_changeset(changeset.tx_graph.into());
 
+        // Apply locked outpoints
+        let mut locked_outpoints = BTreeMap::new();
+        let locked_outpoints::ChangeSet {
+            locked_outpoints: locked_utxos,
+            expiration_heights,
+        } = changeset.locked_outpoints;
+        for (outpoint, is_locked) in locked_utxos {
+            locked_outpoints.insert(
+                outpoint,
+                UtxoLock {
+                    outpoint,
+                    is_locked,
+                    expiration_height: expiration_heights.get(&outpoint).cloned().flatten(),
+                },
+            );
+        }
+
         let stage = ChangeSet::default();
 
         Ok(Some(Wallet {
@@ -664,6 +686,7 @@ impl Wallet {
             stage,
             network,
             secp,
+            locked_outpoints,
         }))
     }
 
@@ -2110,6 +2133,8 @@ impl Wallet {
                     CanonicalizationParams::default(),
                     self.indexed_graph.index.outpoints().iter().cloned(),
                 )
+                // Filter out locked outpoints
+                .filter(|(_, txo)| !self.is_outpoint_locked(txo.outpoint))
                 // only create LocalOutput if UTxO is mature
                 .filter_map(move |((k, i), full_txo)| {
                     full_txo
@@ -2377,6 +2402,104 @@ impl Wallet {
         &self.chain
     }
 
+    /// Get a reference to the locked outpoints.
+    pub fn locked_outpoints(&self) -> &BTreeMap<OutPoint, UtxoLock> {
+        &self.locked_outpoints
+    }
+
+    /// List unspent outpoints that are currently locked.
+    pub fn list_locked_unspent(&self) -> impl Iterator<Item = OutPoint> + '_ {
+        self.list_unspent()
+            .filter(|output| self.is_outpoint_locked(output.outpoint))
+            .map(|output| output.outpoint)
+    }
+
+    /// Whether the `outpoint` is currently locked. See [`Wallet::lock_outpoint`] for more.
+    pub fn is_outpoint_locked(&self, outpoint: OutPoint) -> bool {
+        if let Some(utxo_lock) = self.locked_outpoints.get(&outpoint) {
+            if utxo_lock.is_locked {
+                return utxo_lock
+                    .expiration_height
+                    .map_or(true, |height| self.chain.tip().height() < height);
+            }
+        }
+        false
+    }
+
+    /// Lock a wallet output identified by the given `outpoint`.
+    ///
+    /// A locked UTXO will not be selected as an input to fund a transaction. This is useful
+    /// for excluding or reserving candidate inputs during transaction creation. You can optionally
+    /// specify the `expiration_height` of the lock that defines the height of the local chain at
+    /// which the outpoint becomes spendable.
+    ///
+    /// You must persist the staged change for the lock status to be persistent. To unlock a
+    /// previously locked outpoint, see [`Wallet::unlock_outpoint`].
+    pub fn lock_outpoint(&mut self, outpoint: OutPoint, expiration_height: Option<u32>) {
+        use alloc::collections::btree_map;
+        let lock_value = true;
+        let mut changeset = locked_outpoints::ChangeSet::default();
+
+        // If the lock status changed, update the entry and record the change
+        // in the changeset.
+        match self.locked_outpoints.entry(outpoint) {
+            btree_map::Entry::Occupied(mut e) => {
+                let utxo = e.get_mut();
+                if !utxo.is_locked {
+                    utxo.is_locked = lock_value;
+                    changeset.locked_outpoints.insert(outpoint, lock_value);
+                }
+                if utxo.expiration_height != expiration_height {
+                    utxo.expiration_height = expiration_height;
+                    changeset
+                        .expiration_heights
+                        .insert(outpoint, expiration_height);
+                }
+            }
+            btree_map::Entry::Vacant(e) => {
+                e.insert(UtxoLock {
+                    outpoint,
+                    is_locked: lock_value,
+                    expiration_height,
+                });
+                changeset.locked_outpoints.insert(outpoint, lock_value);
+                changeset
+                    .expiration_heights
+                    .insert(outpoint, expiration_height);
+            }
+        };
+
+        self.stage.merge(ChangeSet {
+            locked_outpoints: changeset,
+            ..Default::default()
+        });
+    }
+
+    /// Unlock the wallet output of the specified `outpoint`.
+    pub fn unlock_outpoint(&mut self, outpoint: OutPoint) {
+        use alloc::collections::btree_map;
+        let lock_value = false;
+
+        match self.locked_outpoints.entry(outpoint) {
+            btree_map::Entry::Vacant(..) => {}
+            btree_map::Entry::Occupied(mut e) => {
+                // If the utxo is currently locked, update the lock value and stage
+                // the change.
+                let utxo = e.get_mut();
+                if utxo.is_locked {
+                    utxo.is_locked = lock_value;
+                    self.stage.merge(ChangeSet {
+                        locked_outpoints: locked_outpoints::ChangeSet {
+                            locked_outpoints: [(outpoint, lock_value)].into(),
+                            ..Default::default()
+                        },
+                        ..Default::default()
+                    });
+                }
+            }
+        }
+    }
+
     /// Introduces a `block` of `height` to the wallet, and tries to connect it to the
     /// `prev_blockhash` of the block's header.
     ///
@@ -2580,6 +2703,24 @@ impl AsRef<bdk_chain::tx_graph::TxGraph<ConfirmationBlockTime>> for Wallet {
     }
 }
 
+/// Records the lock status of a wallet UTXO (outpoint).
+///
+/// Only 1 [`UtxoLock`] may be applied to a particular outpoint at a time. Refer to the docs
+/// for [`Wallet::lock_outpoint`]. Note that the lock status is user-defined and does not take
+/// into account any timelocks directly encoded by the descriptor.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
+#[non_exhaustive]
+pub struct UtxoLock {
+    /// Outpoint.
+    pub outpoint: OutPoint,
+    /// Whether the outpoint is currently locked.
+    pub is_locked: bool,
+    /// Height at which the UTXO lock expires. The outpoint may be selected when the
+    /// wallet chain tip is at or above this height. If `None`, the lock remains in
+    /// effect unless explicitly unlocked.
+    pub expiration_height: Option<u32>,
+}
+
 /// Deterministically generate a unique name given the descriptors defining the wallet
 ///
 /// Compatible with [`wallet_name_from_descriptor`]