[local_chain] Improve update logic and tests

Introduce `LocalChain::update` that replaces `determine_changeset` and
`apply_update`. This returns a closure that actually updates `self` when
called. This method allows for efficient and elegant updating while
being able to "preview" the update before applying.

The `LocalChain` update/determine_changeset tests have been updated to
also check for the final state after applying the update (not just
looking at the changeset).

crates/bdk/Cargo.toml

@@ -49,7 +49,7 @@ dev-getrandom-wasm = ["getrandom/js"]
 lazy_static = "1.4"
 env_logger = "0.7"
 # Move back to importing from rust-bitcoin once https://github.com/rust-bitcoin/rust-bitcoin/pull/1342 is released
-base64 = "^0.13"
+base64 = "^0.21"
 assert_matches = "1.5.0"
 
 

crates/bdk/src/wallet/mod.rs

@@ -1706,10 +1706,7 @@ impl<D> Wallet<D> {
         D: PersistBackend<ChangeSet>,
     {
         let mut changeset = match update.chain.tip() {
-            Some(new_tip) => {
-                let must_include_height = update.chain.checkpoints().keys().next().cloned();
-                ChangeSet::from(self.chain.apply_update(new_tip, must_include_height)?)
-            }
+            Some(new_tip) => ChangeSet::from(self.chain.apply_update(new_tip)?),
             None => ChangeSet::default(),
         };
         let (_, index_additions) = self

crates/chain/src/local_chain.rs

@@ -188,13 +188,9 @@ impl LocalChain {
     }
 
     /// Construct a [`LocalChain`] from a given `checkpoint` tip.
-    pub fn from_checkpoint(checkpoint: CheckPoint) -> Result<Self, InvalidHeightOrderError> {
-        match Self::default().extend_from_last_agreement(checkpoint) {
-            Ok((chain, _)) => Ok(chain),
-            Err(ExtendError::InvalidHeightOrder(err)) => Err(err),
-            Err(ExtendError::CannotConnect(_)) => {
-                unreachable!("we do not need to connect to an empty chain")
-            }
+    pub fn from_checkpoint(checkpoint: CheckPoint) -> Self {
+        Self {
+            checkpoints: checkpoint.iter().map(|cp| (cp.height(), cp)).collect(),
         }
     }
 
@@ -218,10 +214,29 @@ impl LocalChain {
         self.checkpoints.values().last().cloned()
     }
 
-    /// Preview changes of updating [`Self`].
+    /// Returns whether the [`LocalChain`] is empty (has no checkpoints).
+    pub fn is_empty(&self) -> bool {
+        self.checkpoints.is_empty()
+    }
+
+    /// Previews, and optionally applies updates to [`Self`] with the given `new_tip`.
+    ///
+    /// The method returns `(apply_update, changeset)` if [`Ok`]. `apply_update` is a closure that
+    /// can be called to apply the changes represented in `changeset.
     ///
-    /// To invalidate from a given checkpoint, `update` must contain a checkpoint of the same height
-    /// but different hash.
+    /// To update, the `new_tip` must *connect* with `self`. If `self` and `new_tip` has a mutual
+    /// checkpoint (same height and hash), it can connect if:
+    /// * The mutual checkpoint is the tip of `self`.
+    /// * An ancestor of `new_tip` has a height which is of the checkpoint one higher than the
+    ///         mutual checkpoint from `self`.
+    ///
+    /// Additionally:
+    /// * If `self` is empty, `new_tip` will always connect.
+    /// * If `self` only has one checkpoint, `new_tip` must have an ancestor checkpoint with the
+    ///     same height as it.
+    ///
+    /// To invalidate from a given checkpoint, `new_tip` must contain an ancestor checkpoint with
+    /// the same height but different hash.
     ///
     /// # Errors
     ///
@@ -230,43 +245,34 @@ impl LocalChain {
     /// Refer to [module-level documentation] for more.
     ///
     /// [module-level documentation]: crate::local_chain
-    pub fn determine_changeset(
-        &self,
-        new_tip: &CheckPoint,
-        must_include_height: Option<u32>,
-    ) -> Result<ChangeSet, CannotConnectError> {
-        // The changeset if we only records additions from `update`.
-        let mut additions = ChangeSet::new();
-        // Height in which the respective checkpoints of `self` and `update` have the same block id.
+    pub fn update(
+        &mut self,
+        new_tip: CheckPoint,
+    ) -> Result<(impl FnOnce() + '_, ChangeSet), CannotConnectError> {
+        let mut updated_cps = BTreeMap::<u32, CheckPoint>::new();
         let mut agreement_height = Option::<u32>::None;
+        let mut complete_match = false;
 
-        for update_cp in new_tip.iter() {
-            let update_id = update_cp.block_id();
-            let original_id = self
-                .checkpoints
-                .get(&update_cp.height())
-                .map(CheckPoint::block_id);
+        for cp in new_tip.iter() {
+            let block = cp.block_id();
+            let original_cp = self.checkpoints.get(&block.height);
 
-            match original_id {
-                Some(original_id) if original_id.hash == update_id.hash => {}
-                _ => {
-                    additions.insert(update_id.height, Some(update_id.hash));
-                }
-            };
+            // if original block of height does not exist, or if the hash does not match we will
+            // need to update the original checkpoint at that height
+            if original_cp.map(CheckPoint::block_id) != Some(block) {
+                updated_cps.insert(block.height, cp.clone());
+            }
 
             if agreement_height.is_none() {
-                if let Some(original_id) = original_id {
-                    if update_id.hash == original_id.hash {
-                        agreement_height = Some(update_id.height);
+                if let Some(original_cp) = original_cp.cloned() {
+                    if original_cp.block_id() == block {
+                        agreement_height = Some(block.height);
                     }
-                }
-            } else {
-                if let Some(must_include_height) = must_include_height {
-                    if update_id.height > must_include_height {
-                        continue;
+                    if Arc::as_ptr(&original_cp.0) == Arc::as_ptr(&cp.0) {
+                        complete_match = true;
+                        break;
                     }
                 }
-                break;
             }
         }
 
@@ -279,28 +285,46 @@ impl LocalChain {
             Some(height) => height + 1,
         };
 
-        // Construct initial changeset of heights to invalidate in `self`.
-        let mut changeset = self
-            .checkpoints
-            .range(invalidate_lb..)
-            .map(|(&height, _)| (height, None))
-            .collect::<ChangeSet>();
+        let changeset = {
+            // Construct initial changeset of heights to invalidate in `self`.
+            let mut changeset = self
+                .checkpoints
+                .range(invalidate_lb..)
+                .map(|(&height, _)| (height, None))
+                .collect::<ChangeSet>();
+
+            // The height of the first block to invalidate (if any) must be represented in the `update`.
+            if let Some(first_invalidated_height) = changeset.keys().next() {
+                if !updated_cps.contains_key(first_invalidated_height) {
+                    return Err(CannotConnectError {
+                        try_include: self
+                            .checkpoints
+                            .get(first_invalidated_height)
+                            .expect("checkpoint already exists")
+                            .block_id(),
+                    });
+                }
+            }
 
-        // The height of the first block to invalidate (if any) must be represented in the `update`.
-        if let Some(first_invalidated_height) = changeset.keys().next() {
-            if !additions.contains_key(first_invalidated_height) {
-                return Err(CannotConnectError {
-                    try_include: self
-                        .checkpoints
-                        .get(first_invalidated_height)
-                        .expect("checkpoint already exists")
-                        .clone(),
-                });
+            changeset.extend(
+                updated_cps
+                    .iter()
+                    .map(|(height, cp)| (*height, Some(cp.hash()))),
+            );
+            changeset
+        };
+
+        let apply_update = move || {
+            if let Some(&start_height) = updated_cps.keys().next() {
+                self.checkpoints.split_off(&invalidate_lb);
+                self.checkpoints.append(&mut updated_cps);
+                if !self.is_empty() && !complete_match {
+                    self.fix_links(start_height);
+                }
             }
-        }
+        };
 
-        changeset.append(&mut additions);
-        Ok(changeset)
+        Ok((apply_update, changeset))
     }
 
     /// Attempt to extend the chain up to `new_tip` from the last common checkpoint.
@@ -360,8 +384,8 @@ impl LocalChain {
                         .checkpoints
                         .values()
                         .last()
-                        .cloned()
-                        .expect("already checked that chain is not empty"),
+                        .expect("already checked that chain is not empty")
+                        .block_id(),
                 }))?,
             )
         };
@@ -411,13 +435,9 @@ impl LocalChain {
     ///
     /// [`determine_changeset`]: Self::determine_changeset
     /// [`apply_changeset`]: Self::apply_changeset
-    pub fn apply_update(
-        &mut self,
-        new_tip: CheckPoint,
-        must_include_height: Option<u32>,
-    ) -> Result<ChangeSet, CannotConnectError> {
-        let changeset = self.determine_changeset(&new_tip, must_include_height)?;
-        self.apply_changeset(&changeset);
+    pub fn apply_update(&mut self, new_tip: CheckPoint) -> Result<ChangeSet, CannotConnectError> {
+        let (apply, changeset) = self.update(new_tip)?;
+        apply();
         Ok(changeset)
     }
 
@@ -538,16 +558,15 @@ impl std::error::Error for InsertBlockError {}
 #[derive(Clone, Debug, PartialEq)]
 pub struct CannotConnectError {
     /// The suggested checkpoint to include to connect the two chains.
-    pub try_include: CheckPoint,
+    pub try_include: BlockId,
 }
 
 impl core::fmt::Display for CannotConnectError {
     fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
         write!(
             f,
             "introduced chain cannot connect with the original chain, try include {}:{}",
-            self.try_include.height(),
-            self.try_include.hash()
+            self.try_include.height, self.try_include.hash,
         )
     }
 }