feat: add multi_keychain module

ValuedMammal · ValuedMammal · commit 3afb324e4812 · 2025-06-20T15:18:44.000-04:00
- Add rusqlite support for `Wallet&lt;DescriptorId&gt;`.
diff --git a/wallet/src/lib.rs b/wallet/src/lib.rs
@@ -33,6 +33,8 @@ pub mod test_utils;
 mod types;
 mod wallet;
 
+pub mod multi_keychain;
+
 pub(crate) use bdk_chain::collections;
 #[cfg(feature = "rusqlite")]
 pub use bdk_chain::rusqlite;
diff --git a/wallet/src/multi_keychain.rs b/wallet/src/multi_keychain.rs
@@ -0,0 +1,12 @@
+//! Module containing the multi-keychain [`Wallet`].
+
+mod changeset;
+pub mod keyring;
+mod wallet;
+
+pub use changeset::*;
+pub use keyring::KeyRing;
+pub use wallet::*;
+
+/// Alias for [`DescriptorId`](bdk_chain::DescriptorId).
+pub(crate) type Did = bdk_chain::DescriptorId;
diff --git a/wallet/src/multi_keychain/changeset.rs b/wallet/src/multi_keychain/changeset.rs
@@ -0,0 +1,227 @@
+use bdk_chain::{
+    indexed_tx_graph, keychain_txout, local_chain, tx_graph, ConfirmationBlockTime, DescriptorId,
+    Merge,
+};
+use bitcoin::Network;
+use miniscript::{Descriptor, DescriptorPublicKey};
+use serde::{Deserialize, Serialize};
+
+use crate::multi_keychain::keyring;
+
+/// Change set.
+#[derive(Debug, Clone, PartialEq, Deserialize, Serialize)]
+pub struct ChangeSet<K: Ord> {
+    /// Keyring changeset.
+    pub keyring: keyring::ChangeSet<K>,
+    /// Changes to the [`LocalChain`](local_chain::LocalChain).
+    pub local_chain: local_chain::ChangeSet,
+    /// Changes to [`TxGraph`](tx_graph::TxGraph).
+    pub tx_graph: tx_graph::ChangeSet<ConfirmationBlockTime>,
+    /// Changes to [`KeychainTxOutIndex`](keychain_txout::KeychainTxOutIndex).
+    pub indexer: keychain_txout::ChangeSet,
+}
+
+impl<K: Ord> Default for ChangeSet<K> {
+    fn default() -> Self {
+        Self {
+            keyring: Default::default(),
+            local_chain: Default::default(),
+            tx_graph: Default::default(),
+            indexer: Default::default(),
+        }
+    }
+}
+
+impl<K: Ord> Merge for ChangeSet<K> {
+    fn merge(&mut self, other: Self) {
+        // merge keyring
+        self.keyring.merge(other.keyring);
+
+        // merge local chain, tx-graph, indexer
+        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);
+    }
+
+    fn is_empty(&self) -> bool {
+        self.keyring.is_empty()
+            && self.local_chain.is_empty()
+            && self.tx_graph.is_empty()
+            && self.indexer.is_empty()
+    }
+}
+
+#[cfg(feature = "rusqlite")]
+use bdk_chain::rusqlite;
+
+#[cfg(feature = "rusqlite")]
+impl ChangeSet<DescriptorId> {
+    /// Schema name for wallet.
+    pub const WALLET_SCHEMA_NAME: &'static str = "bdk_wallet";
+    /// Name of table to store wallet metainformation.
+    pub const WALLET_TABLE_NAME: &'static str = "bdk_wallet";
+    /// Name of table to store wallet descriptors.
+    pub const DESCRIPTORS_TABLE_NAME: &'static str = "bdk_descriptor";
+
+    /// Get v0 sqlite [ChangeSet] schema.
+    pub fn schema_v0() -> alloc::string::String {
+        format!(
+            "CREATE TABLE {} ( \
+                id INTEGER PRIMARY KEY NOT NULL, \
+                network TEXT NOT NULL \
+            ); \
+            CREATE TABLE {} ( \
+                descriptor_id TEXT PRIMARY KEY NOT NULL, \
+                descriptor BLOB NOT NULL \
+            );",
+            Self::WALLET_TABLE_NAME,
+            Self::DESCRIPTORS_TABLE_NAME,
+        )
+    }
+
+    /// Initializes tables and returns the aggregate data if the database is non-empty
+    /// otherwise returns `Ok(None)`.
+    pub fn initialize(db_tx: &rusqlite::Transaction) -> rusqlite::Result<Option<Self>> {
+        Self::init_sqlite_tables(db_tx)?;
+        let changeset = Self::from_sqlite(db_tx)?;
+
+        if changeset.is_empty() {
+            Ok(None)
+        } else {
+            Ok(Some(changeset))
+        }
+    }
+
+    /// Initialize SQLite tables.
+    fn init_sqlite_tables(db_tx: &rusqlite::Transaction) -> rusqlite::Result<()> {
+        crate::rusqlite_impl::migrate_schema(
+            db_tx,
+            Self::WALLET_SCHEMA_NAME,
+            &[&Self::schema_v0()],
+        )?;
+
+        local_chain::ChangeSet::init_sqlite_tables(db_tx)?;
+        tx_graph::ChangeSet::<ConfirmationBlockTime>::init_sqlite_tables(db_tx)?;
+        keychain_txout::ChangeSet::init_sqlite_tables(db_tx)?;
+
+        Ok(())
+    }
+
+    /// Construct self by reading all of the SQLite data. This should succeed
+    /// even if attempting to read an empty database.
+    fn from_sqlite(db_tx: &rusqlite::Transaction) -> rusqlite::Result<Self> {
+        use bdk_chain::Impl;
+        use rusqlite::OptionalExtension;
+        let mut changeset = Self::default();
+
+        let mut keyring = keyring::ChangeSet::default();
+
+        // Read network
+        let mut network_stmt = db_tx.prepare(&format!(
+            "SELECT network FROM {} WHERE id = 0",
+            Self::WALLET_TABLE_NAME,
+        ))?;
+        let row = network_stmt
+            .query_row([], |row| row.get::<_, Impl<Network>>("network"))
+            .optional()?;
+        if let Some(Impl(network)) = row {
+            keyring.network = Some(network);
+        }
+
+        // Read descriptors
+        let mut descriptor_stmt = db_tx.prepare(&format!(
+            "SELECT descriptor_id, descriptor FROM {}",
+            Self::DESCRIPTORS_TABLE_NAME
+        ))?;
+        let rows = descriptor_stmt.query_map([], |row| {
+            Ok((
+                row.get::<_, Impl<DescriptorId>>("descriptor_id")?,
+                row.get::<_, Impl<Descriptor<DescriptorPublicKey>>>("descriptor")?,
+            ))
+        })?;
+        for row in rows {
+            let (Impl(did), Impl(descriptor)) = row?;
+            keyring.descriptors.insert(did, descriptor);
+        }
+
+        changeset.keyring = keyring;
+        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)?;
+
+        Ok(changeset)
+    }
+
+    /// Persist self to SQLite.
+    pub fn persist_to_sqlite(&self, db_tx: &rusqlite::Transaction) -> rusqlite::Result<()> {
+        use chain::rusqlite::named_params;
+        use chain::Impl;
+
+        let keyring = &self.keyring;
+
+        // Write network
+        let mut network_stmt = db_tx.prepare_cached(&format!(
+            "REPLACE INTO {}(id, network) VALUES(:id, :network)",
+            Self::WALLET_TABLE_NAME,
+        ))?;
+        if let Some(network) = keyring.network {
+            network_stmt.execute(named_params! {
+                ":id": 0,
+                ":network": Impl(network),
+            })?;
+        }
+
+        // Write descriptors
+        let mut descriptor_stmt = db_tx.prepare_cached(&format!(
+            "INSERT OR IGNORE INTO {}(descriptor_id, descriptor) VALUES(:descriptor_id, :descriptor)",
+            Self::DESCRIPTORS_TABLE_NAME,
+        ))?;
+        for (&did, descriptor) in &keyring.descriptors {
+            descriptor_stmt.execute(named_params! {
+                ":descriptor_id": Impl(did),
+                ":descriptor": Impl(descriptor.clone()),
+            })?;
+        }
+
+        self.local_chain.persist_to_sqlite(db_tx)?;
+        self.tx_graph.persist_to_sqlite(db_tx)?;
+        self.indexer.persist_to_sqlite(db_tx)?;
+
+        Ok(())
+    }
+}
+
+impl<K: Ord> From<local_chain::ChangeSet> for ChangeSet<K> {
+    fn from(local_chain: local_chain::ChangeSet) -> Self {
+        Self {
+            local_chain,
+            ..Default::default()
+        }
+    }
+}
+
+impl<K: Ord> From<indexed_tx_graph::ChangeSet<ConfirmationBlockTime, keychain_txout::ChangeSet>>
+    for ChangeSet<K>
+{
+    fn from(
+        indexed_tx_graph: indexed_tx_graph::ChangeSet<
+            ConfirmationBlockTime,
+            keychain_txout::ChangeSet,
+        >,
+    ) -> Self {
+        Self {
+            tx_graph: indexed_tx_graph.tx_graph,
+            indexer: indexed_tx_graph.indexer,
+            ..Default::default()
+        }
+    }
+}
+
+impl<K: Ord> From<keychain_txout::ChangeSet> for ChangeSet<K> {
+    fn from(indexer: keychain_txout::ChangeSet) -> Self {
+        Self {
+            indexer,
+            ..Default::default()
+        }
+    }
+}
diff --git a/wallet/src/multi_keychain/keyring.rs b/wallet/src/multi_keychain/keyring.rs
@@ -0,0 +1,121 @@
+//! [`KeyRing`].
+
+use bdk_chain::{DescriptorExt, Merge};
+use bitcoin::{
+    secp256k1::{All, Secp256k1},
+    Network,
+};
+use miniscript::{Descriptor, DescriptorPublicKey};
+use serde::{Deserialize, Serialize};
+
+use crate::collections::BTreeMap;
+use crate::descriptor::{check_wallet_descriptor, IntoWalletDescriptor};
+use crate::multi_keychain::Did;
+
+/// KeyRing.
+#[derive(Debug, Clone)]
+pub struct KeyRing<K> {
+    pub(crate) secp: Secp256k1<All>,
+    pub(crate) network: Network,
+    pub(crate) descriptors: BTreeMap<K, Descriptor<DescriptorPublicKey>>,
+}
+
+impl<K> KeyRing<K>
+where
+    K: Ord + Clone,
+{
+    /// Construct new [`KeyRing`] with the provided `network`.
+    pub fn new(network: Network) -> Self {
+        Self {
+            secp: Secp256k1::new(),
+            network,
+            descriptors: BTreeMap::default(),
+        }
+    }
+
+    /// Add descriptor, must not be [multipath](miniscript::Descriptor::is_multipath).
+    pub fn add_descriptor(&mut self, keychain: K, descriptor: impl IntoWalletDescriptor) {
+        let descriptor = descriptor
+            .into_wallet_descriptor(&self.secp, self.network)
+            .expect("err: invalid descriptor")
+            .0;
+        assert!(
+            !descriptor.is_multipath(),
+            "err: Use `add_multipath_descriptor` instead"
+        );
+        check_wallet_descriptor(&descriptor).expect("err: failed `check_wallet_descriptor`");
+
+        self.descriptors.insert(keychain, descriptor);
+    }
+
+    /// Initial changeset.
+    pub fn initial_changeset(&self) -> ChangeSet<K> {
+        ChangeSet {
+            network: Some(self.network),
+            descriptors: self.descriptors.clone(),
+        }
+    }
+
+    /// Construct from changeset.
+    pub fn from_changeset(changeset: ChangeSet<K>) -> Option<Self> {
+        Some(Self {
+            secp: Secp256k1::new(),
+            network: changeset.network?,
+            descriptors: changeset.descriptors,
+        })
+    }
+}
+
+impl KeyRing<Did> {
+    /// Add multipath descriptor.
+    pub fn add_multipath_descriptor(&mut self, descriptor: impl IntoWalletDescriptor) {
+        let descriptor = descriptor
+            .into_wallet_descriptor(&self.secp, self.network)
+            .expect("err: invalid descriptor")
+            .0;
+        assert!(
+            descriptor.is_multipath(),
+            "err: Use `add_descriptor` instead"
+        );
+        let descriptors = descriptor
+            .into_single_descriptors()
+            .expect("err: invalid descriptor");
+        for descriptor in descriptors {
+            let did = descriptor.descriptor_id();
+            self.descriptors.insert(did, descriptor);
+        }
+    }
+}
+
+/// Represents changes to the `KeyRing`.
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+pub struct ChangeSet<K: Ord> {
+    /// Network.
+    pub network: Option<Network>,
+    /// Added descriptors.
+    pub descriptors: BTreeMap<K, Descriptor<DescriptorPublicKey>>,
+}
+
+impl<K: Ord> Default for ChangeSet<K> {
+    fn default() -> Self {
+        Self {
+            network: None,
+            descriptors: Default::default(),
+        }
+    }
+}
+
+impl<K: Ord> Merge for ChangeSet<K> {
+    fn merge(&mut self, other: Self) {
+        // merge network
+        if other.network.is_some() && self.network.is_none() {
+            self.network = other.network;
+        }
+        // merge descriptors
+        self.descriptors.extend(other.descriptors);
+    }
+
+    fn is_empty(&self) -> bool {
+        self.network.is_none() && self.descriptors.is_empty()
+    }
+}
diff --git a/wallet/src/multi_keychain/wallet.rs b/wallet/src/multi_keychain/wallet.rs
