feat: add RelevantScripts to extend Wallet for ease-of-use

bitcoindevkit · Dec 13, 2024 · a703c49 · a703c49
1 parent 9d81919
commit a703c49
Show file tree

Hide file tree

Showing 2 changed files with 45 additions and 19 deletions.
diff --git a/src/builder.rs b/src/builder.rs
@@ -45,17 +45,17 @@
 //! }
 //! ```
 
-use std::{collections::HashSet, path::PathBuf, time::Duration};
+use std::{path::PathBuf, time::Duration};
 
 use bdk_chain::local_chain::MissingGenesisError;
-use bdk_wallet::{KeychainKind, Wallet};
+use bdk_wallet::Wallet;
 use kyoto::NodeBuilder;
 pub use kyoto::{
     db::error::SqlInitializationError, AddrV2, HeaderCheckpoint, ScriptBuf, ServiceFlags,
     TrustedPeer,
 };
 
-use crate::{EventReceiver, LightClient};
+use crate::{EventReceiver, LightClient, RelevantScripts};
 
 const RECOMMENDED_PEERS: u8 = 2;
 
@@ -155,21 +155,9 @@ impl<'a> LightClientBuilder<'a> {
         }
         node_builder =
             node_builder.num_required_peers(self.connections.unwrap_or(RECOMMENDED_PEERS));
-        let mut spks: HashSet<ScriptBuf> = HashSet::new();
-        for keychain in [KeychainKind::External, KeychainKind::Internal] {
-            // The user may choose to recover a wallet with lookahead scripts
-            // or use the last revealed index plus some padding to find new transactions
-            let last_revealed = self
-                .wallet
-                .spk_index()
-                .last_revealed_index(keychain)
-                .unwrap_or(0);
-            let lookahead_index = last_revealed + self.wallet.spk_index().lookahead();
-            for index in 0..=lookahead_index {
-                spks.insert(self.wallet.peek_address(keychain, index).script_pubkey());
-            }
-        }
-        let (node, kyoto_client) = node_builder.add_scripts(spks).build_node()?;
+        let (node, kyoto_client) = node_builder
+            .add_scripts(self.wallet.current_relevant_scripts().collect())
+            .build_node()?;
         let (sender, receiver) = kyoto_client.split();
         let event_receiver = EventReceiver::from_index(
             self.wallet.local_chain().tip(),

diff --git a/src/lib.rs b/src/lib.rs
@@ -160,6 +160,7 @@
 #![warn(missing_docs)]
 use core::fmt;
 use std::collections::BTreeMap;
+use std::collections::HashSet;
 
 use bdk_chain::{
     keychain_txout::KeychainTxOutIndex,
@@ -183,7 +184,9 @@ pub use kyoto::{DisconnectedHeader, FailurePayload};
 
 pub use kyoto::ClientSender as EventSender;
 use kyoto::{IndexedBlock, NodeMessage, RejectReason};
-pub use kyoto::{NodeState, Receiver, SyncUpdate, TxBroadcast, TxBroadcastPolicy, Txid, Warning};
+pub use kyoto::{
+    NodeState, Receiver, ScriptBuf, SyncUpdate, TxBroadcast, TxBroadcastPolicy, Txid, Warning,
+};
 
 #[cfg(all(feature = "wallet", feature = "rusqlite"))]
 pub mod builder;
@@ -487,6 +490,41 @@ pub enum LogLevel {
     Warning,
 }
 
+trait RelevantScripts {
+    fn current_relevant_scripts(&self) -> Box<dyn Iterator<Item = ScriptBuf>>;
+}
+
+#[cfg(feature = "wallet")]
+impl RelevantScripts for bdk_wallet::Wallet {
+    fn current_relevant_scripts(&self) -> Box<dyn Iterator<Item = ScriptBuf>> {
+        let mut spks: HashSet<ScriptBuf> = HashSet::new();
+        for keychain in [KeychainKind::External, KeychainKind::Internal] {
+            let last_revealed = self.spk_index().last_revealed_index(keychain).unwrap_or(0);
+            let lookahead_index = last_revealed + self.spk_index().lookahead();
+            for index in 0..=lookahead_index {
+                spks.insert(self.peek_address(keychain, index).script_pubkey());
+            }
+        }
+        Box::new(spks.into_iter())
+    }
+}
+
+impl<K: fmt::Debug + Clone + Ord> RelevantScripts for KeychainTxOutIndex<K> {
+    fn current_relevant_scripts(&self) -> Box<dyn Iterator<Item = ScriptBuf>> {
+        let mut spks = HashSet::new();
+        for (keychain, _) in self.keychains() {
+            let last_revealed = self.last_revealed_index(keychain.clone()).unwrap_or(0);
+            let lookahead_index = last_revealed + self.lookahead();
+            for index in 0..=lookahead_index {
+                if let Some(script) = self.spk_at_index(keychain.clone(), index) {
+                    spks.insert(script);
+                }
+            }
+        }
+        Box::new(spks.into_iter())
+    }
+}
+
 trait StringExt {
     fn into_string(self) -> String;
 }