From a56d289eef0ea241bfe29b4c1069966353e4ca57 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=BF=97=E5=AE=87?= <hello@evanlinjin.me>
Date: Fri, 5 May 2023 16:55:21 +0800
Subject: [PATCH] [chain_redesign] Various `LocalChain` improvements

* Introduce `LocalChain::inner` method to get the inner map of block
  height to hash.
* Replace `LocalChain::get_block` (which outputted `BlockId`, thus able
  to return invalid representation) with `get_blockhash` that just
returns a `BlockHash`.
* Remove `TODO` comments that should be github tickets.
---
 crates/chain/src/local_chain.rs             | 20 ++++++--------
 crates/chain/tests/test_indexed_tx_graph.rs | 30 ++++++++++++++++-----
 2 files changed, 32 insertions(+), 18 deletions(-)

diff --git a/crates/chain/src/local_chain.rs b/crates/chain/src/local_chain.rs
index b070c890..74dbb5a8 100644
--- a/crates/chain/src/local_chain.rs
+++ b/crates/chain/src/local_chain.rs
@@ -6,13 +6,6 @@ use bitcoin::BlockHash;
 use crate::{BlockId, ChainOracle};
 
 /// This is a local implementation of [`ChainOracle`].
-///
-/// TODO: We need a cache/snapshot thing for chain oracle.
-/// * Minimize calls to remotes.
-/// * Can we cache it forever? Should we drop stuff?
-/// * Assume anything deeper than (i.e. 10) blocks won't be reorged.
-/// * Is this a cache on txs or block? or both?
-/// TODO: Parents of children are confirmed if children are confirmed.
 #[derive(Debug, Default, Clone, PartialEq, Eq, PartialOrd, Ord)]
 pub struct LocalChain {
     blocks: BTreeMap<u32, BlockHash>,
@@ -71,6 +64,11 @@ impl LocalChain {
         }
     }
 
+    /// Get a reference to the inner map of block height to hash.
+    pub fn inner(&self) -> &BTreeMap<u32, BlockHash> {
+        &self.blocks
+    }
+
     pub fn tip(&self) -> Option<BlockId> {
         self.blocks
             .iter()
@@ -78,11 +76,9 @@ impl LocalChain {
             .map(|(&height, &hash)| BlockId { height, hash })
     }
 
-    /// Get a block at the given height.
-    pub fn get_block(&self, height: u32) -> Option<BlockId> {
-        self.blocks
-            .get(&height)
-            .map(|&hash| BlockId { height, hash })
+    /// Get a [`BlockHash`] at the given height.
+    pub fn get_blockhash(&self, height: u32) -> Option<BlockHash> {
+        self.blocks.get(&height).cloned()
     }
 
     /// This is like the sparsechain's logic, expect we must guarantee that all invalidated heights
diff --git a/crates/chain/tests/test_indexed_tx_graph.rs b/crates/chain/tests/test_indexed_tx_graph.rs
index d88c1e39..f5011c3e 100644
--- a/crates/chain/tests/test_indexed_tx_graph.rs
+++ b/crates/chain/tests/test_indexed_tx_graph.rs
@@ -8,7 +8,7 @@ use bdk_chain::{
     keychain::{Balance, DerivationAdditions, KeychainTxOutIndex},
     local_chain::LocalChain,
     tx_graph::Additions,
-    ConfirmationHeightAnchor, ObservedAs,
+    BlockId, ConfirmationHeightAnchor, ObservedAs,
 };
 use bitcoin::{secp256k1::Secp256k1, BlockHash, OutPoint, Script, Transaction, TxIn, TxOut};
 use miniscript::Descriptor;
@@ -208,10 +208,12 @@ fn test_list_owned_txouts() {
 
     let _ = graph.insert_relevant_txs(
         [&tx1, &tx2, &tx3, &tx6].iter().enumerate().map(|(i, tx)| {
+            let height = i as u32;
             (
                 *tx,
                 local_chain
-                    .get_block(i as u32)
+                    .get_blockhash(height)
+                    .map(|hash| BlockId { height, hash })
                     .map(|anchor_block| ConfirmationHeightAnchor {
                         anchor_block,
                         confirmation_height: anchor_block.height,
@@ -225,18 +227,34 @@ fn test_list_owned_txouts() {
 
     // A helper lambda to extract and filter data from the graph.
     let fetch =
-        |ht: u32, graph: &IndexedTxGraph<ConfirmationHeightAnchor, KeychainTxOutIndex<String>>| {
+        |height: u32,
+         graph: &IndexedTxGraph<ConfirmationHeightAnchor, KeychainTxOutIndex<String>>| {
             let txouts = graph
-                .list_owned_txouts(&local_chain, local_chain.get_block(ht).unwrap())
+                .list_owned_txouts(
+                    &local_chain,
+                    local_chain
+                        .get_blockhash(height)
+                        .map(|hash| BlockId { height, hash })
+                        .unwrap(),
+                )
                 .collect::<Vec<_>>();
 
             let utxos = graph
-                .list_owned_unspents(&local_chain, local_chain.get_block(ht).unwrap())
+                .list_owned_unspents(
+                    &local_chain,
+                    local_chain
+                        .get_blockhash(height)
+                        .map(|hash| BlockId { height, hash })
+                        .unwrap(),
+                )
                 .collect::<Vec<_>>();
 
             let balance = graph.balance(
                 &local_chain,
-                local_chain.get_block(ht).unwrap(),
+                local_chain
+                    .get_blockhash(height)
+                    .map(|hash| BlockId { height, hash })
+                    .unwrap(),
                 |spk: &Script| trusted_spks.contains(spk),
             );