Merge bitcoindevkit/bdk#1369: feat(chain): add get and range methods to CheckPoint

53942cced492138174638a8087b08f643a8d41ad chore(chain)!: rm `From<LocalChain> for BTreeMap<u32, BlockHash>` (志宇)
2d1d95a6850a5bab9c5bf369a498fe090852cd93 feat(chain): impl `PartialEq` on `CheckPoint` (志宇)
9a62d56900a33a519dd0165ccdb91711917772f3 feat(chain): add `get` and `range` methods to `CheckPoint` (志宇)

Pull request description:

  Partially fixes #1354

  ### Description

  These methods allow us to query for checkpoints contained within the linked list by height and height range. This is useful to determine checkpoints to fetch for chain sources without having to refer back to the `LocalChain`.

  Currently this is not implemented efficiently, but in the future, we will change `CheckPoint` to use a skip list structure.

  ### Notes to the reviewers

  Please refer to the conversation in #1354 for more details. In summary, both `TxGraph::missing_heights` and `tx_graph::ChangeSet::missing_heights_from` are problematic in their own ways. Additionally, it's a better API for chain sources if we only need to lock our receiving structures twice (once to obtain initial state and once for applying the update). Instead of relying on methods such as `EsploraExt::update_local_chain` to get relevant checkpoints, we can use these query methods instead. This allows up to get rid of `::missing_heights`-esc methods and remove the need to lock receiving structures in the middle of the sync.

  ### Changelog notice

  * Added `get` and `range` methods to `CheckPoint` (and in turn, `LocalChain`). This simulates an API where we have implemented a skip list of checkpoints (to implement in the future). This is a better API because we can query for any height or height range with just a checkpoint tip instead of relying on a separate checkpoint index (which needs to live in `LocalChain`).
  * Changed `LocalChain` to have a faster `Eq` implementation. We now maintain an xor value of all checkpoint block hashes. We compare this xor value to determine whether two chains are equal.
  * Added `PartialEq` implementation for `CheckPoint` and `local_chain::Update`.

  ### Checklists

  #### All Submissions:

  * [x] I've signed all my commits
  * [x] I followed the [contribution guidelines](https://github.com/bitcoindevkit/bdk/blob/master/CONTRIBUTING.md)
  * [x] I ran `cargo fmt` and `cargo clippy` before committing

  #### New Features:

  * [x] I've added tests for the new feature
  * [x] I've added docs for the new feature

ACKs for top commit:
  LLFourn:
    ACK 53942cced492138174638a8087b08f643a8d41ad

Tree-SHA512: 90ef8047fe1265daa54c9dfe8a8c520685c898a50d18efd6e803707fcb529d0790d20373c9e439b9c7ff301db32b47453020cae7db4da2ea64eba895aa047f30

crates/bdk/src/wallet/mod.rs

@@ -1127,18 +1127,14 @@ impl<D> Wallet<D> {
                 // anchor tx to checkpoint with lowest height that is >= position's height
                 let anchor = self
                     .chain
-                    .blocks()
                     .range(height..)
-                    .next()
+                    .last()
                     .ok_or(InsertTxError::ConfirmationHeightCannotBeGreaterThanTip {
                         tip_height: self.chain.tip().height(),
                         tx_height: height,
                     })
-                    .map(|(&anchor_height, &hash)| ConfirmationTimeHeightAnchor {
+                    .map(|anchor_cp| ConfirmationTimeHeightAnchor {
-                        anchor_block: BlockId {
+                        anchor_block: anchor_cp.block_id(),
-                            height: anchor_height,
-                            hash,
-                        },
                         confirmation_height: height,
                         confirmation_time: time,
                     })?;

crates/bitcoind_rpc/tests/test_emitter.rs

@@ -57,12 +57,15 @@ pub fn test_sync_local_chain() -> anyhow::Result<()> {
     }
 
     assert_eq!(
-        local_chain.blocks(),
+        local_chain
-        &exp_hashes
+            .iter_checkpoints()
+            .map(|cp| (cp.height(), cp.hash()))
+            .collect::<BTreeSet<_>>(),
+        exp_hashes
             .iter()
             .enumerate()
             .map(|(i, hash)| (i as u32, *hash))
-            .collect(),
+            .collect::<BTreeSet<_>>(),
         "final local_chain state is unexpected",
     );
 
@@ -110,12 +113,15 @@ pub fn test_sync_local_chain() -> anyhow::Result<()> {
     }
 
     assert_eq!(
-        local_chain.blocks(),
+        local_chain
-        &exp_hashes
+            .iter_checkpoints()
+            .map(|cp| (cp.height(), cp.hash()))
+            .collect::<BTreeSet<_>>(),
+        exp_hashes
             .iter()
             .enumerate()
             .map(|(i, hash)| (i as u32, *hash))
-            .collect(),
+            .collect::<BTreeSet<_>>(),
         "final local_chain state is unexpected after reorg",
     );
 

crates/chain/src/local_chain.rs

@@ -1,6 +1,7 @@
 //! The [`LocalChain`] is a local implementation of [`ChainOracle`].
 
 use core::convert::Infallible;
+use core::ops::RangeBounds;
 
 use crate::collections::BTreeMap;
 use crate::{BlockId, ChainOracle};
@@ -34,6 +35,14 @@ struct CPInner {
     prev: Option<Arc<CPInner>>,
 }
 
+impl PartialEq for CheckPoint {
+    fn eq(&self, other: &Self) -> bool {
+        let self_cps = self.iter().map(|cp| cp.block_id());
+        let other_cps = other.iter().map(|cp| cp.block_id());
+        self_cps.eq(other_cps)
+    }
+}
+
 impl CheckPoint {
     /// Construct a new base block at the front of a linked list.
     pub fn new(block: BlockId) -> Self {
@@ -148,6 +157,36 @@ impl CheckPoint {
     pub fn iter(&self) -> CheckPointIter {
         self.clone().into_iter()
     }
+
+    /// Get checkpoint at `height`.
+    ///
+    /// Returns `None` if checkpoint at `height` does not exist`.
+    pub fn get(&self, height: u32) -> Option<Self> {
+        self.range(height..=height).next()
+    }
+
+    /// Iterate checkpoints over a height range.
+    ///
+    /// Note that we always iterate checkpoints in reverse height order (iteration starts at tip
+    /// height).
+    pub fn range<R>(&self, range: R) -> impl Iterator<Item = CheckPoint>
+    where
+        R: RangeBounds<u32>,
+    {
+        let start_bound = range.start_bound().cloned();
+        let end_bound = range.end_bound().cloned();
+        self.iter()
+            .skip_while(move |cp| match end_bound {
+                core::ops::Bound::Included(inc_bound) => cp.height() > inc_bound,
+                core::ops::Bound::Excluded(exc_bound) => cp.height() >= exc_bound,
+                core::ops::Bound::Unbounded => false,
+            })
+            .take_while(move |cp| match start_bound {
+                core::ops::Bound::Included(inc_bound) => cp.height() >= inc_bound,
+                core::ops::Bound::Excluded(exc_bound) => cp.height() > exc_bound,
+                core::ops::Bound::Unbounded => true,
+            })
+    }
 }
 
 /// Iterates over checkpoints backwards.
@@ -188,7 +227,7 @@ impl IntoIterator for CheckPoint {
 /// Script-pubkey based syncing mechanisms may not introduce transactions in a chronological order
 /// so some updates require introducing older blocks (to anchor older transactions). For
 /// script-pubkey based syncing, `introduce_older_blocks` would typically be `true`.
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, PartialEq)]
 pub struct Update {
     /// The update chain's new tip.
     pub tip: CheckPoint,
@@ -202,22 +241,9 @@ pub struct Update {
 }
 
 /// This is a local implementation of [`ChainOracle`].
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, PartialEq)]
 pub struct LocalChain {
     tip: CheckPoint,
-    index: BTreeMap<u32, BlockHash>,
-}
-
-impl PartialEq for LocalChain {
-    fn eq(&self, other: &Self) -> bool {
-        self.index == other.index
-    }
-}
-
-impl From<LocalChain> for BTreeMap<u32, BlockHash> {
-    fn from(value: LocalChain) -> Self {
-        value.index
-    }
 }
 
 impl ChainOracle for LocalChain {
@@ -228,18 +254,16 @@ impl ChainOracle for LocalChain {
         block: BlockId,
         chain_tip: BlockId,
     ) -> Result<Option<bool>, Self::Error> {
-        if block.height > chain_tip.height {
+        let chain_tip_cp = match self.tip.get(chain_tip.height) {
-            return Ok(None);
+            // we can only determine whether `block` is in chain of `chain_tip` if `chain_tip` can
+            // be identified in chain
+            Some(cp) if cp.hash() == chain_tip.hash => cp,
+            _ => return Ok(None),
+        };
+        match chain_tip_cp.get(block.height) {
+            Some(cp) => Ok(Some(cp.hash() == block.hash)),
+            None => Ok(None),
         }
-        Ok(
-            match (
-                self.index.get(&block.height),
-                self.index.get(&chain_tip.height),
-            ) {
-                (Some(cp), Some(tip_cp)) => Some(*cp == block.hash && *tip_cp == chain_tip.hash),
-                _ => None,
-            },
-        )
     }
 
     fn get_chain_tip(&self) -> Result<BlockId, Self::Error> {
@@ -250,7 +274,7 @@ impl ChainOracle for LocalChain {
 impl LocalChain {
     /// Get the genesis hash.
     pub fn genesis_hash(&self) -> BlockHash {
-        self.index.get(&0).copied().expect("must have genesis hash")
+        self.tip.get(0).expect("genesis must exist").hash()
     }
 
     /// Construct [`LocalChain`] from genesis `hash`.
@@ -259,7 +283,6 @@ impl LocalChain {
         let height = 0;
         let chain = Self {
             tip: CheckPoint::new(BlockId { height, hash }),
-            index: core::iter::once((height, hash)).collect(),
         };
         let changeset = chain.initial_changeset();
         (chain, changeset)
@@ -276,7 +299,6 @@ impl LocalChain {
         let (mut chain, _) = Self::from_genesis_hash(genesis_hash);
         chain.apply_changeset(&changeset)?;
 
-        debug_assert!(chain._check_index_is_consistent_with_tip());
         debug_assert!(chain._check_changeset_is_applied(&changeset));
 
         Ok(chain)
@@ -284,18 +306,11 @@ impl LocalChain {
 
     /// Construct a [`LocalChain`] from a given `checkpoint` tip.
     pub fn from_tip(tip: CheckPoint) -> Result<Self, MissingGenesisError> {
-        let mut chain = Self {
+        let genesis_cp = tip.iter().last().expect("must have at least one element");
-            tip,
+        if genesis_cp.height() != 0 {
-            index: BTreeMap::new(),
-        };
-        chain.reindex(0);
-
-        if chain.index.get(&0).copied().is_none() {
             return Err(MissingGenesisError);
         }
-
+        Ok(Self { tip })
-        debug_assert!(chain._check_index_is_consistent_with_tip());
-        Ok(chain)
     }
 
     /// Constructs a [`LocalChain`] from a [`BTreeMap`] of height to [`BlockHash`].
@@ -303,12 +318,11 @@ impl LocalChain {
     /// The [`BTreeMap`] enforces the height order. However, the caller must ensure the blocks are
     /// all of the same chain.
     pub fn from_blocks(blocks: BTreeMap<u32, BlockHash>) -> Result<Self, MissingGenesisError> {
-        if !blocks.contains_key(&0) {
+        if blocks.get(&0).is_none() {
             return Err(MissingGenesisError);
         }
 
         let mut tip: Option<CheckPoint> = None;
-
         for block in &blocks {
             match tip {
                 Some(curr) => {
@@ -321,13 +335,9 @@ impl LocalChain {
             }
         }
 
-        let chain = Self {
+        Ok(Self {
-            index: blocks,
             tip: tip.expect("already checked to have genesis"),
-        };
+        })
-
-        debug_assert!(chain._check_index_is_consistent_with_tip());
-        Ok(chain)
     }
 
     /// Get the highest checkpoint.
@@ -494,9 +504,7 @@ impl LocalChain {
                 None => LocalChain::from_blocks(extension)?.tip(),
             };
             self.tip = new_tip;
-            self.reindex(start_height);
 
-            debug_assert!(self._check_index_is_consistent_with_tip());
             debug_assert!(self._check_changeset_is_applied(changeset));
         }
 
@@ -509,16 +517,16 @@ impl LocalChain {
     ///
     /// Replacing the block hash of an existing checkpoint will result in an error.
     pub fn insert_block(&mut self, block_id: BlockId) -> Result<ChangeSet, AlterCheckPointError> {
-        if let Some(&original_hash) = self.index.get(&block_id.height) {
+        if let Some(original_cp) = self.tip.get(block_id.height) {
+            let original_hash = original_cp.hash();
             if original_hash != block_id.hash {
                 return Err(AlterCheckPointError {
                     height: block_id.height,
                     original_hash,
                     update_hash: Some(block_id.hash),
                 });
-            } else {
-                return Ok(ChangeSet::default());
             }
+            return Ok(ChangeSet::default());
         }
 
         let mut changeset = ChangeSet::default();
@@ -542,33 +550,41 @@ impl LocalChain {
     /// This will fail with [`MissingGenesisError`] if the caller attempts to disconnect from the
     /// genesis block.
     pub fn disconnect_from(&mut self, block_id: BlockId) -> Result<ChangeSet, MissingGenesisError> {
-        if self.index.get(&block_id.height) != Some(&block_id.hash) {
+        let mut remove_from = Option::<CheckPoint>::None;
-            return Ok(ChangeSet::default());
+        let mut changeset = ChangeSet::default();
-        }
+        for cp in self.tip().iter() {
-
+            let cp_id = cp.block_id();
-        let changeset = self
+            if cp_id.height < block_id.height {
-            .index
-            .range(block_id.height..)
-            .map(|(&height, _)| (height, None))
-            .collect::<ChangeSet>();
-        self.apply_changeset(&changeset).map(|_| changeset)
-    }
-
-    /// Reindex the heights in the chain from (and including) `from` height
-    fn reindex(&mut self, from: u32) {
-        let _ = self.index.split_off(&from);
-        for cp in self.iter_checkpoints() {
-            if cp.height() < from {
                 break;
             }
-            self.index.insert(cp.height(), cp.hash());
+            changeset.insert(cp_id.height, None);
+            if cp_id == block_id {
+                remove_from = Some(cp);
+            }
         }
+        self.tip = match remove_from.map(|cp| cp.prev()) {
+            // The checkpoint below the earliest checkpoint to remove will be the new tip.
+            Some(Some(new_tip)) => new_tip,
+            // If there is no checkpoint below the earliest checkpoint to remove, it means the
+            // "earliest checkpoint to remove" is the genesis block. We disallow removing the
+            // genesis block.
+            Some(None) => return Err(MissingGenesisError),
+            // If there is nothing to remove, we return an empty changeset.
+            None => return Ok(ChangeSet::default()),
+        };
+        Ok(changeset)
     }
 
     /// Derives an initial [`ChangeSet`], meaning that it can be applied to an empty chain to
     /// recover the current chain.
     pub fn initial_changeset(&self) -> ChangeSet {
-        self.index.iter().map(|(k, v)| (*k, Some(*v))).collect()
+        self.tip
+            .iter()
+            .map(|cp| {
+                let block_id = cp.block_id();
+                (block_id.height, Some(block_id.hash))
+            })
+            .collect()
     }
 
     /// Iterate over checkpoints in descending height order.
@@ -578,28 +594,49 @@ impl LocalChain {
         }
     }
 
-    /// Get a reference to the internal index mapping the height to block hash.
-    pub fn blocks(&self) -> &BTreeMap<u32, BlockHash> {
-        &self.index
-    }
-
-    fn _check_index_is_consistent_with_tip(&self) -> bool {
-        let tip_history = self
-            .tip
-            .iter()
-            .map(|cp| (cp.height(), cp.hash()))
-            .collect::<BTreeMap<_, _>>();
-        self.index == tip_history
-    }
-
     fn _check_changeset_is_applied(&self, changeset: &ChangeSet) -> bool {
-        for (height, exp_hash) in changeset {
+        let mut curr_cp = self.tip.clone();
-            if self.index.get(height) != exp_hash.as_ref() {
+        for (height, exp_hash) in changeset.iter().rev() {
-                return false;
+            match curr_cp.get(*height) {
+                Some(query_cp) => {
+                    if query_cp.height() != *height || Some(query_cp.hash()) != *exp_hash {
+                        return false;
+                    }
+                    curr_cp = query_cp;
+                }
+                None => {
+                    if exp_hash.is_some() {
+                        return false;
+                    }
+                }
             }
         }
         true
     }
+
+    /// Get checkpoint at given `height` (if it exists).
+    ///
+    /// This is a shorthand for calling [`CheckPoint::get`] on the [`tip`].
+    ///
+    /// [`tip`]: LocalChain::tip
+    pub fn get(&self, height: u32) -> Option<CheckPoint> {
+        self.tip.get(height)
+    }
+
+    /// Iterate checkpoints over a height range.
+    ///
+    /// Note that we always iterate checkpoints in reverse height order (iteration starts at tip
+    /// height).
+    ///
+    /// This is a shorthand for calling [`CheckPoint::range`] on the [`tip`].
+    ///
+    /// [`tip`]: LocalChain::tip
+    pub fn range<R>(&self, range: R) -> impl Iterator<Item = CheckPoint>
+    where
+        R: RangeBounds<u32>,
+    {
+        self.tip.range(range)
+    }
 }
 
 /// An error which occurs when a [`LocalChain`] is constructed without a genesis checkpoint.

crates/chain/src/tx_graph.rs

@@ -725,13 +725,13 @@ impl<A: Anchor> TxGraph<A> {
                     };
                     let mut has_missing_height = false;
                     for anchor_block in tx_anchors.iter().map(Anchor::anchor_block) {
-                        match chain.blocks().get(&anchor_block.height) {
+                        match chain.get(anchor_block.height) {
                             None => {
                                 has_missing_height = true;
                                 continue;
                             }
-                            Some(chain_hash) => {
+                            Some(chain_cp) => {
-                                if chain_hash == &anchor_block.hash {
+                                if chain_cp.hash() == anchor_block.hash {
                                     return true;
                                 }
                             }
@@ -749,7 +749,7 @@ impl<A: Anchor> TxGraph<A> {
             .filter_map(move |(a, _)| {
                 let anchor_block = a.anchor_block();
                 if Some(anchor_block.height) != last_height_emitted
-                    && !chain.blocks().contains_key(&anchor_block.height)
+                    && chain.get(anchor_block.height).is_none()
                 {
                     last_height_emitted = Some(anchor_block.height);
                     Some(anchor_block.height)
@@ -1299,7 +1299,7 @@ impl<A> ChangeSet<A> {
         A: Anchor,
     {
         self.anchor_heights()
-            .filter(move |height| !local_chain.blocks().contains_key(height))
+            .filter(move |&height| local_chain.get(height).is_none())
     }
 }
 

crates/chain/tests/test_indexed_tx_graph.rs

@@ -7,7 +7,7 @@ use bdk_chain::{
     indexed_tx_graph::{self, IndexedTxGraph},
     keychain::{self, Balance, KeychainTxOutIndex},
     local_chain::LocalChain,
-    tx_graph, BlockId, ChainPosition, ConfirmationHeightAnchor,
+    tx_graph, ChainPosition, ConfirmationHeightAnchor,
 };
 use bitcoin::{secp256k1::Secp256k1, OutPoint, Script, ScriptBuf, Transaction, TxIn, TxOut};
 use miniscript::Descriptor;
@@ -212,10 +212,8 @@ fn test_list_owned_txouts() {
             (
                 *tx,
                 local_chain
-                    .blocks()
+                    .get(height)
-                    .get(&height)
+                    .map(|cp| cp.block_id())
-                    .cloned()
-                    .map(|hash| BlockId { height, hash })
                     .map(|anchor_block| ConfirmationHeightAnchor {
                         anchor_block,
                         confirmation_height: anchor_block.height,
@@ -230,9 +228,8 @@ fn test_list_owned_txouts() {
         |height: u32,
          graph: &IndexedTxGraph<ConfirmationHeightAnchor, KeychainTxOutIndex<String>>| {
             let chain_tip = local_chain
-                .blocks()
+                .get(height)
-                .get(&height)
+                .map(|cp| cp.block_id())
-                .map(|&hash| BlockId { height, hash })
                 .unwrap_or_else(|| panic!("block must exist at {}", height));
             let txouts = graph
                 .graph()

crates/chain/tests/test_local_chain.rs

@@ -528,6 +528,52 @@ fn checkpoint_from_block_ids() {
     }
 }
 
+#[test]
+fn checkpoint_query() {
+    struct TestCase {
+        chain: LocalChain,
+        /// The heights we want to call [`CheckPoint::query`] with, represented as an inclusive
+        /// range.
+        ///
+        /// If a [`CheckPoint`] exists at that height, we expect [`CheckPoint::query`] to return
+        /// it. If not, [`CheckPoint::query`] should return `None`.
+        query_range: (u32, u32),
+    }
+
+    let test_cases = [
+        TestCase {
+            chain: local_chain![(0, h!("_")), (1, h!("A"))],
+            query_range: (0, 2),
+        },
+        TestCase {
+            chain: local_chain![(0, h!("_")), (2, h!("B")), (3, h!("C"))],
+            query_range: (0, 3),
+        },
+    ];
+
+    for t in test_cases.into_iter() {
+        let tip = t.chain.tip();
+        for h in t.query_range.0..=t.query_range.1 {
+            let query_result = tip.get(h);
+
+            // perform an exhausitive search for the checkpoint at height `h`
+            let exp_hash = t
+                .chain
+                .iter_checkpoints()
+                .find(|cp| cp.height() == h)
+                .map(|cp| cp.hash());
+
+            match query_result {
+                Some(cp) => {
+                    assert_eq!(Some(cp.hash()), exp_hash);
+                    assert_eq!(cp.height(), h);
+                }
+                None => assert!(exp_hash.is_none()),
+            }
+        }
+    }
+}
+
 #[test]
 fn local_chain_apply_header_connected_to() {
     fn header_from_prev_blockhash(prev_blockhash: BlockHash) -> Header {

crates/esplora/tests/blocking_ext.rs

@@ -360,8 +360,8 @@ fn update_local_chain() -> anyhow::Result<()> {
         for height in t.request_heights {
             let exp_blockhash = blocks.get(height).expect("block must exist in bitcoind");
             assert_eq!(
-                chain.blocks().get(height),
+                chain.get(*height).map(|cp| cp.hash()),
-                Some(exp_blockhash),
+                Some(*exp_blockhash),
                 "[{}:{}] block {}:{} must exist in final chain",
                 i,
                 t.name,