From a703c49fd1d51184b74b0b9ab9d69458c74d1bb3 Mon Sep 17 00:00:00 2001
From: Rob N <rob.netzke@gmail.com>
Date: Thu, 12 Dec 2024 14:59:21 -1000
Subject: [PATCH] feat: add `RelevantScripts` to extend `Wallet` for
 ease-of-use

---
 src/builder.rs | 24 ++++++------------------
 src/lib.rs     | 40 +++++++++++++++++++++++++++++++++++++++-
 2 files changed, 45 insertions(+), 19 deletions(-)

diff --git a/src/builder.rs b/src/builder.rs
index 4e859f5..94f7a66 100644
--- 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
index d7beb72..0fd3c2f 100644
--- 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;
 }