From c09cd2afce4e649caa2797628edaffae08a60628 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=BF=97=E5=AE=87?= <hello@evanlinjin.me>
Date: Fri, 31 Mar 2023 22:42:47 +0800
Subject: [PATCH] [bdk_chain_redesign] Added methods to `LocalChain`

Also made the `IndexedTxGraph::index` field public (`index()` and
`index_mut()` methods are no longer needed).
---
 crates/chain/src/indexed_tx_graph.rs | 13 ++------
 crates/chain/src/local_chain.rs      | 50 ++++++++++++++++++++++++++--
 2 files changed, 49 insertions(+), 14 deletions(-)

diff --git a/crates/chain/src/indexed_tx_graph.rs b/crates/chain/src/indexed_tx_graph.rs
index 0b27150c..79e5105c 100644
--- a/crates/chain/src/indexed_tx_graph.rs
+++ b/crates/chain/src/indexed_tx_graph.rs
@@ -58,8 +58,9 @@ impl<A: BlockAnchor, IA: AddAssign> AddAssign for IndexedAdditions<A, IA> {
 }
 
 pub struct IndexedTxGraph<A, I> {
+    /// Transaction index.
+    pub index: I,
     graph: TxGraph<A>,
-    index: I, // [TODO] Make public
     last_height: u32,
 }
 
@@ -79,16 +80,6 @@ impl<A: BlockAnchor, I: TxIndex> IndexedTxGraph<A, I> {
         &self.graph
     }
 
-    /// Get a reference of the internal transaction index.
-    pub fn index(&self) -> &I {
-        &self.index
-    }
-
-    /// Get a mutable reference to the internal transaction index.
-    pub fn index_mut(&mut self) -> &mut I {
-        &mut self.index
-    }
-
     /// Applies the [`IndexedAdditions`] to the [`IndexedTxGraph`].
     pub fn apply_additions(&mut self, additions: IndexedAdditions<A, I::Additions>) {
         let IndexedAdditions {
diff --git a/crates/chain/src/local_chain.rs b/crates/chain/src/local_chain.rs
index 5bcb524f..5d459a15 100644
--- a/crates/chain/src/local_chain.rs
+++ b/crates/chain/src/local_chain.rs
@@ -1,6 +1,9 @@
-use core::convert::Infallible;
+use core::{convert::Infallible, ops::Deref};
 
-use alloc::{collections::BTreeMap, vec::Vec};
+use alloc::{
+    collections::{BTreeMap, BTreeSet},
+    vec::Vec,
+};
 use bitcoin::BlockHash;
 
 use crate::{BlockId, ChainOracle};
@@ -104,11 +107,52 @@ impl LocalChain {
         }
         Ok(ChangeSet(changeset))
     }
+
+    /// Applies the given `changeset`.
+    pub fn apply_changeset(&mut self, mut changeset: ChangeSet) {
+        self.blocks.append(&mut changeset.0)
+    }
+
+    /// Updates [`LocalChain`] with an update [`LocalChain`].
+    ///
+    /// This is equivilant to calling [`determine_changeset`] and [`apply_changeset`] in sequence.
+    ///
+    /// [`determine_changeset`]: Self::determine_changeset
+    /// [`apply_changeset`]: Self::apply_changeset
+    pub fn apply_update(&mut self, update: Self) -> Result<ChangeSet, UpdateError> {
+        let changeset = self.determine_changeset(&update)?;
+        self.apply_changeset(changeset.clone());
+        Ok(changeset)
+    }
+
+    pub fn initial_changeset(&self) -> ChangeSet {
+        ChangeSet(self.blocks.clone())
+    }
+
+    pub fn heights(&self) -> BTreeSet<u32> {
+        self.blocks.keys().cloned().collect()
+    }
 }
 
-#[derive(Debug, Default)]
+/// This is the return value of [`determine_changeset`] and represents changes to [`LocalChain`].
+///
+/// [`determine_changeset`]: LocalChain::determine_changeset
+#[derive(Debug, Default, Clone, PartialEq)]
+#[cfg_attr(
+    feature = "serde",
+    derive(serde::Deserialize, serde::Serialize),
+    serde(crate = "serde_crate")
+)]
 pub struct ChangeSet(pub BTreeMap<u32, BlockHash>);
 
+impl Deref for ChangeSet {
+    type Target = BTreeMap<u32, BlockHash>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
 /// Represents an update failure of [`LocalChain`].j
 #[derive(Clone, Debug, PartialEq)]
 pub enum UpdateError {