feat(electrum)!: use new sync/full-scan structs for ElectrumExt

`ElectrumResultExt` trait is also introduced that adds methods which can
convert the `Anchor` type for the update `TxGraph`.

Examples and tests are updated to use the new `ElectrumExt` API.

Author: evanlinjin

crates/electrum/Cargo.toml

@@ -12,7 +12,7 @@ readme = "README.md"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
-bdk_chain = { path = "../chain", version = "0.13.0", default-features = false }
+bdk_chain = { path = "../chain", version = "0.13.0" }
 electrum-client = { version = "0.19" }
 #rustls = { version = "=0.21.1", optional = true, features = ["dangerous_configuration"] }
 

crates/electrum/src/electrum_ext.rs

@@ -2,87 +2,17 @@ use bdk_chain::{
     bitcoin::{OutPoint, ScriptBuf, Transaction, Txid},
     collections::{BTreeMap, HashMap, HashSet},
     local_chain::CheckPoint,
+    spk_client::{FullScanRequest, FullScanResult, SyncRequest, SyncResult, TxCache},
     tx_graph::TxGraph,
     BlockId, ConfirmationHeightAnchor, ConfirmationTimeHeightAnchor,
 };
-use core::{fmt::Debug, str::FromStr};
+use core::str::FromStr;
 use electrum_client::{ElectrumApi, Error, HeaderNotification};
 use std::sync::Arc;
 
 /// We include a chain suffix of a certain length for the purpose of robustness.
 const CHAIN_SUFFIX_LENGTH: u32 = 8;
 
-/// Type that maintains a cache of [`Arc`]-wrapped transactions.
-pub type TxCache = HashMap<Txid, Arc<Transaction>>;
-
-/// Combination of chain and transactions updates from electrum
-///
-/// We have to update the chain and the txids at the same time since we anchor the txids to
-/// the same chain tip that we check before and after we gather the txids.
-#[derive(Debug)]
-pub struct ElectrumUpdate<A = ConfirmationHeightAnchor> {
-    /// Chain update
-    pub chain_update: CheckPoint,
-    /// Tracks electrum updates in TxGraph
-    pub graph_update: TxGraph<A>,
-}
-
-impl<A: Clone + Ord> ElectrumUpdate<A> {
-    /// Transform the [`ElectrumUpdate`] to have [`bdk_chain::Anchor`]s of another type.
-    ///
-    /// Refer to [`TxGraph::map_anchors`].
-    pub fn map_anchors<A2: Clone + Ord, F>(self, f: F) -> ElectrumUpdate<A2>
-    where
-        F: FnMut(A) -> A2,
-    {
-        ElectrumUpdate {
-            chain_update: self.chain_update,
-            graph_update: self.graph_update.map_anchors(f),
-        }
-    }
-}
-
-impl ElectrumUpdate {
-    /// Transforms the [`TxGraph`]'s [`bdk_chain::Anchor`] type to [`ConfirmationTimeHeightAnchor`].
-    pub fn into_confirmation_time_update(
-        self,
-        client: &impl ElectrumApi,
-    ) -> Result<ElectrumUpdate<ConfirmationTimeHeightAnchor>, Error> {
-        let relevant_heights = self
-            .graph_update
-            .all_anchors()
-            .iter()
-            .map(|(a, _)| a.confirmation_height)
-            .collect::<HashSet<_>>();
-
-        let height_to_time = relevant_heights
-            .clone()
-            .into_iter()
-            .zip(
-                client
-                    .batch_block_header(relevant_heights)?
-                    .into_iter()
-                    .map(|bh| bh.time as u64),
-            )
-            .collect::<HashMap<u32, u64>>();
-
-        let chain_update = self.chain_update;
-        let graph_update =
-            self.graph_update
-                .clone()
-                .map_anchors(|a| ConfirmationTimeHeightAnchor {
-                    anchor_block: a.anchor_block,
-                    confirmation_height: a.confirmation_height,
-                    confirmation_time: height_to_time[&a.confirmation_height],
-                });
-
-        Ok(ElectrumUpdate {
-            chain_update,
-            graph_update,
-        })
-    }
-}
-
 /// Trait to extend [`electrum_client::Client`] functionality.
 pub trait ElectrumExt {
     /// Full scan the keychain scripts specified with the blockchain (via an Electrum client) and
@@ -97,12 +27,10 @@ pub trait ElectrumExt {
     /// single batch request.
     fn full_scan<K: Ord + Clone>(
         &self,
-        tx_cache: &mut TxCache,
-        prev_tip: CheckPoint,
-        keychain_spks: BTreeMap<K, impl IntoIterator<Item = (u32, ScriptBuf)>>,
+        request: FullScanRequest<K>,
         stop_gap: usize,
         batch_size: usize,
-    ) -> Result<(ElectrumUpdate, BTreeMap<K, u32>), Error>;
+    ) -> Result<FullScanResult<K, ConfirmationHeightAnchor>, Error>;
 
     /// Sync a set of scripts with the blockchain (via an Electrum client) for the data specified
     /// and returns updates for [`bdk_chain`] data structures.
@@ -123,28 +51,19 @@ pub trait ElectrumExt {
     /// [`full_scan`]: ElectrumExt::full_scan
     fn sync(
         &self,
-        tx_cache: &mut TxCache,
-        prev_tip: CheckPoint,
-        misc_spks: impl IntoIterator<Item = ScriptBuf>,
-        txids: impl IntoIterator<Item = Txid>,
-        outpoints: impl IntoIterator<Item = OutPoint>,
+        request: SyncRequest,
         batch_size: usize,
-    ) -> Result<ElectrumUpdate, Error>;
+    ) -> Result<SyncResult<ConfirmationHeightAnchor>, Error>;
 }
 
 impl<E: ElectrumApi> ElectrumExt for E {
     fn full_scan<K: Ord + Clone>(
         &self,
-        tx_cache: &mut TxCache,
-        prev_tip: CheckPoint,
-        keychain_spks: BTreeMap<K, impl IntoIterator<Item = (u32, ScriptBuf)>>,
+        mut request: FullScanRequest<K>,
         stop_gap: usize,
         batch_size: usize,
-    ) -> Result<(ElectrumUpdate, BTreeMap<K, u32>), Error> {
-        let mut request_spks = keychain_spks
-            .into_iter()
-            .map(|(k, s)| (k, s.into_iter()))
-            .collect::<BTreeMap<K, _>>();
+    ) -> Result<FullScanResult<K, ConfirmationHeightAnchor>, Error> {
+        let mut request_spks = request.spks_by_keychain;
 
         // We keep track of already-scanned spks just in case a reorg happens and we need to do a
         // rescan. We need to keep track of this as iterators in `keychain_spks` are "unbounded" so
@@ -154,8 +73,8 @@ impl<E: ElectrumApi> ElectrumExt for E {
         // * val: (script_pubkey, has_tx_history).
         let mut scanned_spks = BTreeMap::<(K, u32), (ScriptBuf, bool)>::new();
 
-        let (electrum_update, keychain_update) = loop {
-            let (tip, _) = construct_update_tip(self, prev_tip.clone())?;
+        let update = loop {
+            let (tip, _) = construct_update_tip(self, request.chain_tip.clone())?;
             let mut graph_update = TxGraph::<ConfirmationHeightAnchor>::default();
             let cps = tip
                 .iter()
@@ -168,7 +87,7 @@ impl<E: ElectrumApi> ElectrumExt for E {
                     scanned_spks.append(&mut populate_with_spks(
                         self,
                         &cps,
-                        tx_cache,
+                        &mut request.tx_cache,
                         &mut graph_update,
                         &mut scanned_spks
                             .iter()
@@ -182,7 +101,7 @@ impl<E: ElectrumApi> ElectrumExt for E {
                         populate_with_spks(
                             self,
                             &cps,
-                            tx_cache,
+                            &mut request.tx_cache,
                             &mut graph_update,
                             keychain_spks,
                             stop_gap,
@@ -213,55 +132,118 @@ impl<E: ElectrumApi> ElectrumExt for E {
                 })
                 .collect::<BTreeMap<_, _>>();
 
-            break (
-                ElectrumUpdate {
-                    chain_update,
-                    graph_update,
-                },
-                keychain_update,
-            );
+            break FullScanResult {
+                graph_update,
+                chain_update,
+                last_active_indices: keychain_update,
+            };
         };
 
-        Ok((electrum_update, keychain_update))
+        Ok(update)
     }
 
     fn sync(
         &self,
-        tx_cache: &mut TxCache,
-        prev_tip: CheckPoint,
-        misc_spks: impl IntoIterator<Item = ScriptBuf>,
-        txids: impl IntoIterator<Item = Txid>,
-        outpoints: impl IntoIterator<Item = OutPoint>,
+        request: SyncRequest,
         batch_size: usize,
-    ) -> Result<ElectrumUpdate, Error> {
-        let spk_iter = misc_spks
-            .into_iter()
-            .enumerate()
-            .map(|(i, spk)| (i as u32, spk));
-
-        let (electrum_update, _) = self.full_scan(
-            tx_cache,
-            prev_tip.clone(),
-            [((), spk_iter)].into(),
-            usize::MAX,
-            batch_size,
-        )?;
-
-        let (tip, _) = construct_update_tip(self, prev_tip)?;
+    ) -> Result<SyncResult<ConfirmationHeightAnchor>, Error> {
+        let mut tx_cache = request.tx_cache.clone();
+
+        let full_scan_req = FullScanRequest::from_chain_tip(request.chain_tip.clone())
+            .cache_txs(request.tx_cache)
+            .set_spks_for_keychain((), request.spks.enumerate().map(|(i, spk)| (i as u32, spk)));
+        let full_scan_res = self.full_scan(full_scan_req, usize::MAX, batch_size)?;
+
+        let (tip, _) = construct_update_tip(self, request.chain_tip)?;
         let cps = tip
             .iter()
             .take(10)
             .map(|cp| (cp.height(), cp))
             .collect::<BTreeMap<u32, CheckPoint>>();
 
         let mut tx_graph = TxGraph::<ConfirmationHeightAnchor>::default();
-        populate_with_txids(self, &cps, tx_cache, &mut tx_graph, txids)?;
-        populate_with_outpoints(self, &cps, &mut tx_graph, outpoints)?;
+        populate_with_txids(self, &cps, &mut tx_cache, &mut tx_graph, request.txids)?;
+        populate_with_outpoints(self, &cps, &mut tx_cache, &mut tx_graph, request.outpoints)?;
+
+        Ok(SyncResult {
+            chain_update: full_scan_res.chain_update,
+            graph_update: full_scan_res.graph_update,
+        })
+    }
+}
+
+/// Trait that extends [`SyncResult`] and [`FullScanResult`] functionality.
+///
+/// Currently, only a single method exists that converts the update [`TxGraph`] to have an anchor
+/// type of [`ConfirmationTimeHeightAnchor`].
+pub trait ElectrumResultExt {
+    /// New result type with a [`TxGraph`] that contains the [`ConfirmationTimeHeightAnchor`].
+    type NewResult;
+
+    /// Convert result type to have an update [`TxGraph`]  that contains the [`ConfirmationTimeHeightAnchor`] .
+    fn try_into_confirmation_time_result(
+        self,
+        client: &impl ElectrumApi,
+    ) -> Result<Self::NewResult, Error>;
+}
+
+impl<K> ElectrumResultExt for FullScanResult<K, ConfirmationHeightAnchor> {
+    type NewResult = FullScanResult<K, ConfirmationTimeHeightAnchor>;
+
+    fn try_into_confirmation_time_result(
+        self,
+        client: &impl ElectrumApi,
+    ) -> Result<Self::NewResult, Error> {
+        Ok(FullScanResult::<K, ConfirmationTimeHeightAnchor> {
+            graph_update: try_into_confirmation_time_result(self.graph_update, client)?,
+            chain_update: self.chain_update,
+            last_active_indices: self.last_active_indices,
+        })
+    }
+}
+
+impl ElectrumResultExt for SyncResult<ConfirmationHeightAnchor> {
+    type NewResult = SyncResult<ConfirmationTimeHeightAnchor>;
 
-        Ok(electrum_update)
+    fn try_into_confirmation_time_result(
+        self,
+        client: &impl ElectrumApi,
+    ) -> Result<Self::NewResult, Error> {
+        Ok(SyncResult {
+            graph_update: try_into_confirmation_time_result(self.graph_update, client)?,
+            chain_update: self.chain_update,
+        })
     }
 }
 
+fn try_into_confirmation_time_result(
+    graph_update: TxGraph<ConfirmationHeightAnchor>,
+    client: &impl ElectrumApi,
+) -> Result<TxGraph<ConfirmationTimeHeightAnchor>, Error> {
+    let relevant_heights = graph_update
+        .all_anchors()
+        .iter()
+        .map(|(a, _)| a.confirmation_height)
+        .collect::<HashSet<_>>();
+
+    let height_to_time = relevant_heights
+        .clone()
+        .into_iter()
+        .zip(
+            client
+                .batch_block_header(relevant_heights)?
+                .into_iter()
+                .map(|bh| bh.time as u64),
+        )
+        .collect::<HashMap<u32, u64>>();
+
+    Ok(graph_update.map_anchors(|a| ConfirmationTimeHeightAnchor {
+        anchor_block: a.anchor_block,
+        confirmation_height: a.confirmation_height,
+        confirmation_time: height_to_time[&a.confirmation_height],
+    }))
+}
+
 /// Return a [`CheckPoint`] of the latest tip, that connects with `prev_tip`.
 fn construct_update_tip(
     client: &impl ElectrumApi,
@@ -380,6 +362,7 @@ fn determine_tx_anchor(
 fn populate_with_outpoints(
     client: &impl ElectrumApi,
     cps: &BTreeMap<u32, CheckPoint>,
+    tx_cache: &mut TxCache,
     tx_graph: &mut TxGraph<ConfirmationHeightAnchor>,
     outpoints: impl IntoIterator<Item = OutPoint>,
 ) -> Result<(), Error> {
@@ -415,9 +398,9 @@ fn populate_with_outpoints(
                 let res_tx = match tx_graph.get_tx(res.tx_hash) {
                     Some(tx) => tx,
                     None => {
-                        let res_tx = client.transaction_get(&res.tx_hash)?;
-                        let _ = tx_graph.insert_tx(res_tx);
-                        tx_graph.get_tx(res.tx_hash).expect("just inserted")
+                        let res_tx = fetch_tx(client, tx_cache, res.tx_hash)?;
+                        let _ = tx_graph.insert_tx(Arc::clone(&res_tx));
+                        res_tx
                     }
                 };
                 has_spending = res_tx