[chain_redesign] Add LocalChain::insert_block

Author: evanlinjin

crates/chain/src/local_chain.rs

@@ -173,6 +173,31 @@ impl LocalChain {
     pub fn heights(&self) -> BTreeSet<u32> {
         self.blocks.keys().cloned().collect()
     }
+
+    /// Insert a block of [`BlockId`] into the [`LocalChain`].
+    ///
+    /// # Error
+    ///
+    /// If the insertion height already contains a block, and the block has a different blockhash,
+    /// this will result in an [`InsertBlockNotMatchingError`].
+    pub fn insert_block(
+        &mut self,
+        block_id: BlockId,
+    ) -> Result<ChangeSet, InsertBlockNotMatchingError> {
+        let mut update = Self::from_blocks(self.tip());
+
+        if let Some(original_hash) = update.blocks.insert(block_id.height, block_id.hash) {
+            if original_hash != block_id.hash {
+                return Err(InsertBlockNotMatchingError {
+                    height: block_id.height,
+                    original_hash,
+                    update_hash: block_id.hash,
+                });
+            }
+        }
+
+        Ok(self.apply_update(update).expect("should always connect"))
+    }
 }
 
 /// This is the return value of [`determine_changeset`] and represents changes to [`LocalChain`].
@@ -201,3 +226,24 @@ impl core::fmt::Display for UpdateNotConnectedError {
 
 #[cfg(feature = "std")]
 impl std::error::Error for UpdateNotConnectedError {}
+
+/// Represents a failure when trying to insert a checkpoint into [`LocalChain`].
+#[derive(Clone, Debug, PartialEq)]
+pub struct InsertBlockNotMatchingError {
+    pub height: u32,
+    pub original_hash: BlockHash,
+    pub update_hash: BlockHash,
+}
+
+impl core::fmt::Display for InsertBlockNotMatchingError {
+    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
+        write!(
+            f,
+            "failed to insert block at height {} as blockhashes conflict: original={}, update={}",
+            self.height, self.original_hash, self.update_hash
+        )
+    }
+}
+
+#[cfg(feature = "std")]
+impl std::error::Error for InsertBlockNotMatchingError {}

crates/chain/tests/test_local_chain.rs

@@ -1,4 +1,7 @@
-use bdk_chain::local_chain::{LocalChain, UpdateNotConnectedError};
+use bdk_chain::local_chain::{
+    ChangeSet, InsertBlockNotMatchingError, LocalChain, UpdateNotConnectedError,
+};
+use bitcoin::BlockHash;
 
 #[macro_use]
 mod common;
@@ -165,3 +168,61 @@ fn invalidation_but_no_connection() {
         Err(UpdateNotConnectedError(0))
     )
 }
+
+#[test]
+fn insert_block() {
+    struct TestCase {
+        original: LocalChain,
+        insert: (u32, BlockHash),
+        expected_result: Result<ChangeSet, InsertBlockNotMatchingError>,
+        expected_final: LocalChain,
+    }
+
+    let test_cases = [
+        TestCase {
+            original: local_chain![],
+            insert: (5, h!("block5")),
+            expected_result: Ok([(5, Some(h!("block5")))].into()),
+            expected_final: local_chain![(5, h!("block5"))],
+        },
+        TestCase {
+            original: local_chain![(3, h!("A"))],
+            insert: (4, h!("B")),
+            expected_result: Ok([(4, Some(h!("B")))].into()),
+            expected_final: local_chain![(3, h!("A")), (4, h!("B"))],
+        },
+        TestCase {
+            original: local_chain![(4, h!("B"))],
+            insert: (3, h!("A")),
+            expected_result: Ok([(3, Some(h!("A")))].into()),
+            expected_final: local_chain![(3, h!("A")), (4, h!("B"))],
+        },
+        TestCase {
+            original: local_chain![(2, h!("K"))],
+            insert: (2, h!("K")),
+            expected_result: Ok([].into()),
+            expected_final: local_chain![(2, h!("K"))],
+        },
+        TestCase {
+            original: local_chain![(2, h!("K"))],
+            insert: (2, h!("J")),
+            expected_result: Err(InsertBlockNotMatchingError {
+                height: 2,
+                original_hash: h!("K"),
+                update_hash: h!("J"),
+            }),
+            expected_final: local_chain![(2, h!("K"))],
+        },
+    ];
+
+    for (i, t) in test_cases.into_iter().enumerate() {
+        let mut chain = t.original;
+        assert_eq!(
+            chain.insert_block(t.insert.into()),
+            t.expected_result,
+            "[{}] unexpected result when inserting block",
+            i,
+        );
+        assert_eq!(chain, t.expected_final, "[{}] unexpected final chain", i,);
+    }
+}