feat(chain): add query and query_multi methods to CheckPoint

These methods allow us to query for checkpoints contained within the
linked list by height. 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.

crates/chain/src/local_chain.rs

@@ -4,6 +4,7 @@ use core::convert::Infallible;
 
 use crate::collections::BTreeMap;
 use crate::{BlockId, ChainOracle};
+use alloc::collections::BTreeSet;
 use alloc::sync::Arc;
 use bitcoin::block::Header;
 use bitcoin::BlockHash;
@@ -148,6 +149,34 @@ impl CheckPoint {
     pub fn iter(&self) -> CheckPointIter {
         self.clone().into_iter()
     }
+
+    /// Find checkpoint at `height`.
+    ///
+    /// Returns `None` if checkpoint at `height` does not exist`. If you are querying for multiple
+    /// heights, it will be more efficient to use [`query_multi`].
+    ///
+    /// [`query_multi`]: CheckPoint::query_multi
+    pub fn query(&self, height: u32) -> Option<Self> {
+        self.iter()
+            .take_while(|cp| cp.height() >= height)
+            .find(|cp| cp.height() == height)
+    }
+
+    /// Find checkpoints of the given `heights`.
+    ///
+    /// This returns an iterator which iterates over the requested heights and whether a
+    /// corresponding checkpoint is found.
+    pub fn query_multi(&self, heights: BTreeSet<u32>) -> impl Iterator<Item = (u32, Option<Self>)> {
+        let mut cp_iter = self.iter();
+        let mut current_cp = cp_iter.next();
+        heights.into_iter().rev().map(move |h| loop {
+            match current_cp.as_ref() {
+                Some(cp) if cp.height() > h => current_cp = cp_iter.next(),
+                Some(cp) if cp.height() == h => return (h, Some(cp.clone())),
+                _ => return (h, None),
+            }
+        })
+    }
 }
 
 /// Iterates over checkpoints backwards.

crates/chain/tests/test_local_chain.rs

@@ -1,3 +1,5 @@
+use std::collections::BTreeSet;
+
 use bdk_chain::{
     local_chain::{
         AlterCheckPointError, ApplyHeaderError, CannotConnectError, ChangeSet, CheckPoint,
@@ -528,6 +530,106 @@ 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.query(h);
+            let exp_hash = t.chain.blocks().get(&h).cloned();
+            match query_result {
+                Some(cp) => {
+                    assert_eq!(Some(cp.hash()), exp_hash);
+                    assert_eq!(cp.height(), h);
+                }
+                None => assert!(query_result.is_none()),
+            }
+        }
+    }
+}
+
+#[test]
+fn checkpoint_query_multi() {
+    struct TestCase<'a> {
+        /// Name of the test.
+        name: &'a str,
+        /// The local chain which contains the [`CheckPoint`] linked list to test on.
+        chain: LocalChain,
+        /// Tuples of heights to query for to the expected [`BlockHash`]s of the [`CheckPoint`]
+        /// results.
+        queries: &'a [u32],
+    }
+
+    let test_cases = [
+        TestCase {
+            name: "query_all_existing",
+            chain: local_chain![(0, h!("_")), (1, h!("A")), (3, h!("C"))],
+            queries: &[0, 1, 3],
+        },
+        TestCase {
+            name: "query_all_non_existing",
+            chain: local_chain![(0, h!("_")), (1, h!("A")), (3, h!("C"))],
+            queries: &[2, 4, 100],
+        },
+        TestCase {
+            name: "query_mix_of_existing_and_non_existing",
+            chain: local_chain![(0, h!("_")), (3, h!("C")), (4, h!("D")), (10, h!("F"))],
+            queries: &[2, 4, 5, 10],
+        },
+    ];
+
+    for (i, t) in test_cases.into_iter().enumerate() {
+        println!("Running test case [{}]: {}", i, t.name);
+        let chain = t.chain;
+        let tip = chain.tip();
+        let query_multi_result = tip
+            .query_multi(t.queries.iter().copied().collect())
+            .collect::<Vec<_>>();
+        assert_eq!(
+            query_multi_result.len(),
+            t.queries.len(),
+            "what is queried and what is returned should have the same count"
+        );
+        assert_eq!(
+            query_multi_result
+                .iter()
+                .map(|(h, _)| *h)
+                .collect::<BTreeSet<u32>>(),
+            t.queries.iter().copied().collect(),
+            "all queried heights should have a corresponding returned height"
+        );
+        for (h, cp_result) in query_multi_result {
+            let exp_hash = chain.blocks().get(&h).copied();
+            assert_eq!(
+                cp_result.map(|cp| (cp.height(), cp.hash())),
+                exp_hash.map(|hash| (h, hash)),
+                "queried checkpoint results must have correct heights and hashes"
+            );
+        }
+    }
+}
+
 #[test]
 fn local_chain_apply_header_connected_to() {
     fn header_from_prev_blockhash(prev_blockhash: BlockHash) -> Header {