[bdk_chain_redesign] chain_oracle::Cache

Introduce `chain_oracle::Cache` which is a cache for requests to the
chain oracle. `ChainOracle` has also been moved to the `chain_oracle`
module.

Introduce `get_tip_in_best_chain` method to the `ChainOracle` trait.
This allows for guaranteeing that chain state can be consistent across
operations with `IndexedTxGraph`.

crates/chain/src/chain_oracle.rs

@@ -0,0 +1,162 @@
+use core::{convert::Infallible, marker::PhantomData};
+
+use alloc::collections::BTreeMap;
+use bitcoin::BlockHash;
+
+use crate::BlockId;
+
+/// Represents a service that tracks the best chain history.
+/// TODO: How do we ensure the chain oracle is consistent across a single call?
+/// * We need to somehow lock the data! What if the ChainOracle is remote?
+/// * Get tip method! And check the tip still exists at the end! And every internal call
+///   does not go beyond the initial tip.
+pub trait ChainOracle {
+    /// Error type.
+    type Error: core::fmt::Debug;
+
+    /// Get the height and hash of the tip in the best chain.
+    fn get_tip_in_best_chain(&self) -> Result<Option<BlockId>, Self::Error>;
+
+    /// Returns the block hash (if any) of the given `height`.
+    fn get_block_in_best_chain(&self, height: u32) -> Result<Option<BlockHash>, Self::Error>;
+
+    /// Determines whether the block of [`BlockId`] exists in the best chain.
+    fn is_block_in_best_chain(&self, block_id: BlockId) -> Result<bool, Self::Error> {
+        Ok(matches!(self.get_block_in_best_chain(block_id.height)?, Some(h) if h == block_id.hash))
+    }
+}
+
+// [TODO] We need stuff for smart pointers. Maybe? How does rust lib do this?
+// Box<dyn ChainOracle>, Arc<dyn ChainOracle> ????? I will figure it out
+impl<C: ChainOracle> ChainOracle for &C {
+    type Error = C::Error;
+
+    fn get_tip_in_best_chain(&self) -> Result<Option<BlockId>, Self::Error> {
+        <C as ChainOracle>::get_tip_in_best_chain(self)
+    }
+
+    fn get_block_in_best_chain(&self, height: u32) -> Result<Option<BlockHash>, Self::Error> {
+        <C as ChainOracle>::get_block_in_best_chain(self, height)
+    }
+
+    fn is_block_in_best_chain(&self, block_id: BlockId) -> Result<bool, Self::Error> {
+        <C as ChainOracle>::is_block_in_best_chain(self, block_id)
+    }
+}
+
+/// This structure increases the performance of getting chain data.
+#[derive(Debug)]
+pub struct Cache<C> {
+    assume_final_depth: u32,
+    tip_height: u32,
+    cache: BTreeMap<u32, BlockHash>,
+    marker: PhantomData<C>,
+}
+
+impl<C> Cache<C> {
+    /// Creates a new [`Cache`].
+    ///
+    /// `assume_final_depth` represents the minimum number of blocks above the block in question
+    /// when we can assume the block is final (reorgs cannot happen). I.e. a value of 0 means the
+    /// tip is assumed to be final. The cache only caches blocks that are assumed to be final.
+    pub fn new(assume_final_depth: u32) -> Self {
+        Self {
+            assume_final_depth,
+            tip_height: 0,
+            cache: Default::default(),
+            marker: Default::default(),
+        }
+    }
+}
+
+impl<C: ChainOracle> Cache<C> {
+    /// This is the topmost (highest) block height that we assume as final (no reorgs possible).
+    ///
+    /// Blocks higher than this height are not cached.
+    pub fn assume_final_height(&self) -> u32 {
+        self.tip_height.saturating_sub(self.assume_final_depth)
+    }
+
+    /// Update the `tip_height` with the [`ChainOracle`]'s tip.
+    ///
+    /// `tip_height` is used with `assume_final_depth` to determine whether we should cache a
+    /// certain block height (`tip_height` - `assume_final_depth`).
+    pub fn try_update_tip_height(&mut self, chain: C) -> Result<(), C::Error> {
+        let tip = chain.get_tip_in_best_chain()?;
+        if let Some(BlockId { height, .. }) = tip {
+            self.tip_height = height;
+        }
+        Ok(())
+    }
+
+    /// Get a block from the cache with the [`ChainOracle`] as fallback.
+    ///
+    /// If the block does not exist in cache, the logic fallbacks to fetching from the internal
+    /// [`ChainOracle`]. If the block is at or below the "assume final height", we will also store
+    /// the missing block in the cache.
+    pub fn try_get_block(&mut self, chain: C, height: u32) -> Result<Option<BlockHash>, C::Error> {
+        if let Some(&hash) = self.cache.get(&height) {
+            return Ok(Some(hash));
+        }
+
+        let hash = chain.get_block_in_best_chain(height)?;
+
+        if hash.is_some() && height > self.tip_height {
+            self.tip_height = height;
+        }
+
+        // only cache block if at least as deep as `assume_final_depth`
+        let assume_final_height = self.tip_height.saturating_sub(self.assume_final_depth);
+        if height <= assume_final_height {
+            if let Some(hash) = hash {
+                self.cache.insert(height, hash);
+            }
+        }
+
+        Ok(hash)
+    }
+
+    /// Determines whether the block of `block_id` is in the chain using the cache.
+    ///
+    /// This uses [`try_get_block`] internally.
+    ///
+    /// [`try_get_block`]: Self::try_get_block
+    pub fn try_is_block_in_chain(&mut self, chain: C, block_id: BlockId) -> Result<bool, C::Error> {
+        match self.try_get_block(chain, block_id.height)? {
+            Some(hash) if hash == block_id.hash => Ok(true),
+            _ => Ok(false),
+        }
+    }
+}
+
+impl<C: ChainOracle<Error = Infallible>> Cache<C> {
+    /// Updates the `tip_height` with the [`ChainOracle`]'s tip.
+    ///
+    /// This is the no-error version of [`try_update_tip_height`].
+    ///
+    /// [`try_update_tip_height`]: Self::try_update_tip_height
+    pub fn update_tip_height(&mut self, chain: C) {
+        self.try_update_tip_height(chain)
+            .expect("chain oracle error is infallible")
+    }
+
+    /// Get a block from the cache with the [`ChainOracle`] as fallback.
+    ///
+    /// This is the no-error version of [`try_get_block`].
+    ///
+    /// [`try_get_block`]: Self::try_get_block
+    pub fn get_block(&mut self, chain: C, height: u32) -> Option<BlockHash> {
+        self.try_get_block(chain, height)
+            .expect("chain oracle error is infallible")
+    }
+
+    /// Determines whether the block at `block_id` is in the chain using the cache.
+    ///
+    /// This is the no-error version of [`try_is_block_in_chain`].
+    ///
+    /// [`try_is_block_in_chain`]: Self::try_is_block_in_chain
+    pub fn is_block_in_best_chain(&mut self, chain: C, block_id: BlockId) -> bool {
+        self.try_is_block_in_chain(chain, block_id)
+            .expect("chain oracle error is infallible")
+    }
+}

crates/chain/src/lib.rs

@@ -31,6 +31,8 @@ pub mod sparse_chain;
 mod tx_data_traits;
 pub mod tx_graph;
 pub use tx_data_traits::*;
+mod chain_oracle;
+pub use chain_oracle::*;
 
 #[doc(hidden)]
 pub mod example_utils;

crates/chain/src/local_chain.rs

@@ -22,6 +22,14 @@ pub struct LocalChain {
 impl ChainOracle for LocalChain {
     type Error = Infallible;
 
+    fn get_tip_in_best_chain(&self) -> Result<Option<BlockId>, Self::Error> {
+        Ok(self
+            .blocks
+            .iter()
+            .last()
+            .map(|(&height, &hash)| BlockId { height, hash }))
+    }
+
     fn get_block_in_best_chain(&self, height: u32) -> Result<Option<BlockHash>, Self::Error> {
         Ok(self.blocks.get(&height).cloned())
     }
@@ -153,7 +161,7 @@ impl Deref for ChangeSet {
     }
 }
 
-/// Represents an update failure of [`LocalChain`].j
+/// Represents an update failure of [`LocalChain`].
 #[derive(Clone, Debug, PartialEq)]
 pub enum UpdateError {
     /// The update cannot be applied to the chain because the chain suffix it represents did not

crates/chain/src/sparse_chain.rs

@@ -307,6 +307,7 @@
 //! );
 //! ```
 use core::{
+    convert::Infallible,
     fmt::Debug,
     ops::{Bound, RangeBounds},
 };
@@ -457,7 +458,15 @@ impl<P: core::fmt::Debug> core::fmt::Display for UpdateError<P> {
 impl<P: core::fmt::Debug> std::error::Error for UpdateError<P> {}
 
 impl<P: ChainPosition> ChainOracle for SparseChain<P> {
-    type Error = ();
+    type Error = Infallible;
+
+    fn get_tip_in_best_chain(&self) -> Result<Option<BlockId>, Self::Error> {
+        Ok(self
+            .checkpoints
+            .iter()
+            .last()
+            .map(|(&height, &hash)| BlockId { height, hash }))
+    }
 
     fn get_block_in_best_chain(&self, height: u32) -> Result<Option<BlockHash>, Self::Error> {
         Ok(self.checkpoint_at(height).map(|b| b.hash))

crates/chain/src/tx_data_traits.rs

@@ -56,38 +56,6 @@ impl BlockAnchor for (u32, BlockHash) {
     }
 }
 
-/// Represents a service that tracks the best chain history.
-/// TODO: How do we ensure the chain oracle is consistent across a single call?
-/// * We need to somehow lock the data! What if the ChainOracle is remote?
-/// * Get tip method! And check the tip still exists at the end! And every internal call
-///   does not go beyond the initial tip.
-pub trait ChainOracle {
-    /// Error type.
-    type Error: core::fmt::Debug;
-
-    /// Returns the block hash (if any) of the given `height`.
-    fn get_block_in_best_chain(&self, height: u32) -> Result<Option<BlockHash>, Self::Error>;
-
-    /// Determines whether the block of [`BlockId`] exists in the best chain.
-    fn is_block_in_best_chain(&self, block_id: BlockId) -> Result<bool, Self::Error> {
-        Ok(matches!(self.get_block_in_best_chain(block_id.height)?, Some(h) if h == block_id.hash))
-    }
-}
-
-// [TODO] We need stuff for smart pointers. Maybe? How does rust lib do this?
-// Box<dyn ChainOracle>, Arc<dyn ChainOracle> ????? I will figure it out
-impl<C: ChainOracle> ChainOracle for &C {
-    type Error = C::Error;
-
-    fn get_block_in_best_chain(&self, height: u32) -> Result<Option<BlockHash>, Self::Error> {
-        <C as ChainOracle>::get_block_in_best_chain(self, height)
-    }
-
-    fn is_block_in_best_chain(&self, block_id: BlockId) -> Result<bool, Self::Error> {
-        <C as ChainOracle>::is_block_in_best_chain(self, block_id)
-    }
-}
-
 /// Represents an index of transaction data.
 pub trait TxIndex {
     /// The resultant "additions" when new transaction data is indexed.

crates/chain/src/tx_graph.rs

@@ -586,11 +586,12 @@ impl<A: Clone + Ord> TxGraph<A> {
 
 impl<A: BlockAnchor> TxGraph<A> {
     /// Get all heights that are relevant to the graph.
-    pub fn relevant_heights(&self) -> BTreeSet<u32> {
+    pub fn relevant_heights(&self) -> impl Iterator<Item = u32> + '_ {
+        let mut visited = HashSet::new();
         self.anchors
             .iter()
             .map(|(a, _)| a.anchor_block().height)
-            .collect()
+            .filter(move |&h| visited.insert(h))
     }
 
     /// Determines whether a transaction of `txid` is in the best chain.