refactor(core): CheckPoint takes a generic

LagginTimes · LagginTimes · commit 7c1378121ec6 · 2024-09-02T10:53:36.000+08:00
diff --git a/crates/core/src/checkpoint.rs b/crates/core/src/checkpoint.rs
@@ -1,7 +1,7 @@
 use core::ops::RangeBounds;
 
 use alloc::sync::Arc;
-use bitcoin::BlockHash;
+use bitcoin::{block::Header, consensus::Encodable, hashes::Hash, BlockHash};
 
 use crate::BlockId;
 
@@ -10,29 +10,83 @@ use crate::BlockId;
 /// Checkpoints are cheaply cloneable and are useful to find the agreement point between two sparse
 /// block chains.
 #[derive(Debug, Clone)]
-pub struct CheckPoint(Arc<CPInner>);
+pub struct CheckPoint<B = BlockHash>(Arc<CPInner<B>>);
 
 /// The internal contents of [`CheckPoint`].
 #[derive(Debug, Clone)]
-struct CPInner {
-    /// Block id (hash and height).
-    block: BlockId,
+struct CPInner<B> {
+    /// Block height.
+    height: u32,
+    /// Data.
+    data: B,
     /// Previous checkpoint (if any).
-    prev: Option<Arc<CPInner>>,
+    prev: Option<Arc<CPInner<B>>>,
 }
 
-impl PartialEq for CheckPoint {
+/// TODO: ToBlockHash doc
+pub trait ToBlockHash {
+    /// TODO: to_blockhash doc
+    fn to_blockhash(&self) -> BlockHash;
+}
+
+impl<B: ToBlockHash> ToBlockHash for CheckPoint<B> {
+    fn to_blockhash(&self) -> BlockHash {
+        self.0.data.to_blockhash()
+    }
+}
+
+impl ToBlockHash for Header {
+    fn to_blockhash(&self) -> BlockHash {
+        let mut bytes = vec![];
+        self.consensus_encode(&mut bytes).unwrap_or_default();
+        BlockHash::hash(&bytes)
+    }
+}
+
+impl<B> PartialEq for CheckPoint<B>
+where
+    B: Copy + core::cmp::PartialEq,
+{
     fn eq(&self, other: &Self) -> bool {
-        let self_cps = self.iter().map(|cp| cp.block_id());
-        let other_cps = other.iter().map(|cp| cp.block_id());
+        let self_cps = self.iter().map(|cp| *cp.inner());
+        let other_cps = other.iter().map(|cp| *cp.inner());
         self_cps.eq(other_cps)
     }
 }
 
-impl CheckPoint {
-    /// Construct a new base block at the front of a linked list.
+impl CheckPoint<BlockHash> {
+    /// Construct a new [`CheckPoint`] at the front of a linked list.
     pub fn new(block: BlockId) -> Self {
-        Self(Arc::new(CPInner { block, prev: None }))
+        Self(Arc::new(CPInner {
+            height: block.height,
+            data: block.hash,
+            prev: None,
+        }))
+    }
+
+    /// Construct a checkpoint from the given `header` and block `height`.
+    ///
+    /// If `header` is of the genesis block, the checkpoint won't have a [`prev`] node. Otherwise,
+    /// we return a checkpoint linked with the previous block.
+    ///
+    /// [`prev`]: CheckPoint::prev
+    pub fn from_header(header: &bitcoin::block::Header, height: u32) -> Self {
+        let hash = header.block_hash();
+        let this_block_id = BlockId { height, hash };
+
+        let prev_height = match height.checked_sub(1) {
+            Some(h) => h,
+            None => return Self::new(this_block_id),
+        };
+
+        let prev_block_id = BlockId {
+            height: prev_height,
+            hash: header.prev_blockhash,
+        };
+
+        CheckPoint::new(prev_block_id)
+            .push(this_block_id)
+            .expect("must construct checkpoint")
     }
 
     /// Construct a checkpoint from a list of [`BlockId`]s in ascending height order.
@@ -50,36 +104,74 @@ impl CheckPoint {
         block_ids: impl IntoIterator<Item = BlockId>,
     ) -> Result<Self, Option<Self>> {
         let mut blocks = block_ids.into_iter();
-        let mut acc = CheckPoint::new(blocks.next().ok_or(None)?);
+        let block = blocks.next().ok_or(None)?;
+        let mut acc = CheckPoint::new(block);
         for id in blocks {
             acc = acc.push(id).map_err(Some)?;
         }
         Ok(acc)
     }
 
-    /// Construct a checkpoint from the given `header` and block `height`.
+    /// Extends the checkpoint linked list by a iterator of block ids.
     ///
-    /// If `header` is of the genesis block, the checkpoint won't have a [`prev`] node. Otherwise,
-    /// we return a checkpoint linked with the previous block.
+    /// Returns an `Err(self)` if there is block which does not have a greater height than the
+    /// previous one.
+    pub fn extend(self, blockdata: impl IntoIterator<Item = BlockId>) -> Result<Self, Self> {
+        let mut curr = self.clone();
+        for block in blockdata {
+            curr = curr.push(block).map_err(|_| self.clone())?;
+        }
+        Ok(curr)
+    }
+
+    /// Get the block hash of the checkpoint.
+    pub fn hash(&self) -> BlockHash {
+        self.0.data
+    }
+
+    /// Get the [`BlockId`] of the checkpoint.
+    pub fn block_id(&self) -> BlockId {
+        BlockId {
+            height: self.height(),
+            hash: self.hash(),
+        }
+    }
+
+    /// Inserts `block_id` at its height within the chain.
     ///
-    /// [`prev`]: CheckPoint::prev
-    pub fn from_header(header: &bitcoin::block::Header, height: u32) -> Self {
-        let hash = header.block_hash();
-        let this_block_id = BlockId { height, hash };
+    /// The effect of `insert` depends on whether a height already exists. If it doesn't the
+    /// `block_id` we inserted and all pre-existing blocks higher than it will be re-inserted after
+    /// it. If the height already existed and has a conflicting block hash then it will be purged
+    /// along with all block followin it. The returned chain will have a tip of the `block_id`
+    /// passed in. Of course, if the `block_id` was already present then this just returns `self`.
+    #[must_use]
+    pub fn insert(self, block_id: BlockId) -> Self {
+        assert_ne!(block_id.height, 0, "cannot insert the genesis block");
 
-        let prev_height = match height.checked_sub(1) {
-            Some(h) => h,
-            None => return Self::new(this_block_id),
-        };
+        let mut cp = self.clone();
+        let mut tail = vec![];
+        let base = loop {
+            if cp.height() == block_id.height {
+                if cp.hash() == block_id.hash {
+                    return self;
+                }
+                // if we have a conflict we just return the inserted block because the tail is by
+                // implication invalid.
+                tail = vec![];
+                break cp.prev().expect("can't be called on genesis block");
+            }
 
-        let prev_block_id = BlockId {
-            height: prev_height,
-            hash: header.prev_blockhash,
+            if cp.height() < block_id.height {
+                break cp;
+            }
+
+            tail.push(cp.block_id());
+            cp = cp.prev().expect("will break before genesis block");
         };
 
-        CheckPoint::new(prev_block_id)
-            .push(this_block_id)
-            .expect("must construct checkpoint")
+        let new_cp = core::iter::once(block_id).chain(tail.into_iter().rev());
+
+        base.extend(new_cp).expect("tail is in order")
     }
 
     /// Puts another checkpoint onto the linked list representing the blockchain.
@@ -89,48 +181,46 @@ impl CheckPoint {
     pub fn push(self, block: BlockId) -> Result<Self, Self> {
         if self.height() < block.height {
             Ok(Self(Arc::new(CPInner {
-                block,
+                height: block.height,
+                data: block.hash,
                 prev: Some(self.0),
             })))
         } else {
             Err(self)
         }
     }
+}
 
-    /// Extends the checkpoint linked list by a iterator of block ids.
-    ///
-    /// Returns an `Err(self)` if there is block which does not have a greater height than the
-    /// previous one.
-    pub fn extend(self, blocks: impl IntoIterator<Item = BlockId>) -> Result<Self, Self> {
-        let mut curr = self.clone();
-        for block in blocks {
-            curr = curr.push(block).map_err(|_| self.clone())?;
-        }
-        Ok(curr)
+impl<B> CheckPoint<B>
+where
+    B: Copy,
+{
+    /// Construct a new [`CheckPoint`] at the front of a linked list.
+    pub fn from_data(height: u32, data: B) -> Self {
+        Self(Arc::new(CPInner {
+            height,
+            data,
+            prev: None,
+        }))
     }
 
-    /// Get the [`BlockId`] of the checkpoint.
-    pub fn block_id(&self) -> BlockId {
-        self.0.block
+    /// Get reference to the inner type.
+    pub fn inner(&self) -> &B {
+        &self.0.data
     }
 
     /// Get the height of the checkpoint.
     pub fn height(&self) -> u32 {
-        self.0.block.height
-    }
-
-    /// Get the block hash of the checkpoint.
-    pub fn hash(&self) -> BlockHash {
-        self.0.block.hash
+        self.0.height
     }
 
     /// Get the previous checkpoint in the chain
-    pub fn prev(&self) -> Option<CheckPoint> {
+    pub fn prev(&self) -> Option<CheckPoint<B>> {
         self.0.prev.clone().map(CheckPoint)
     }
 
     /// Iterate from this checkpoint in descending height.
-    pub fn iter(&self) -> CheckPointIter {
+    pub fn iter(&self) -> CheckPointIter<B> {
         self.clone().into_iter()
     }
 
@@ -145,7 +235,7 @@ impl CheckPoint {
     ///
     /// Note that we always iterate checkpoints in reverse height order (iteration starts at tip
     /// height).
-    pub fn range<R>(&self, range: R) -> impl Iterator<Item = CheckPoint>
+    pub fn range<R>(&self, range: R) -> impl Iterator<Item = CheckPoint<B>>
     where
         R: RangeBounds<u32>,
     {
@@ -164,6 +254,28 @@ impl CheckPoint {
             })
     }
 
+    /// This method tests for `self` and `other` to have equal internal pointers.
+    pub fn eq_ptr(&self, other: &Self) -> bool {
+        Arc::as_ptr(&self.0) == Arc::as_ptr(&other.0)
+    }
+}
+
+impl<B> CheckPoint<B>
+where
+    B: Copy + core::fmt::Debug + ToBlockHash + From<BlockHash>,
+{
+    /// Extends the checkpoint linked list by a iterator of block ids.
+    ///
+    /// Returns an `Err(self)` if there is block which does not have a greater height than the
+    /// previous one.
+    pub fn extend_data(self, blockdata: impl IntoIterator<Item = (u32, B)>) -> Result<Self, Self> {
+        let mut curr = self.clone();
+        for (height, data) in blockdata {
+            curr = curr.push_data(height, data).map_err(|_| self.clone())?;
+        }
+        Ok(curr)
+    }
+
     /// Inserts `block_id` at its height within the chain.
     ///
     /// The effect of `insert` depends on whether a height already exists. If it doesn't the
@@ -172,14 +284,14 @@ impl CheckPoint {
     /// along with all block followin it. The returned chain will have a tip of the `block_id`
     /// passed in. Of course, if the `block_id` was already present then this just returns `self`.
     #[must_use]
-    pub fn insert(self, block_id: BlockId) -> Self {
-        assert_ne!(block_id.height, 0, "cannot insert the genesis block");
+    pub fn insert_data(self, height: u32, data: B) -> Self {
+        assert_ne!(height, 0, "cannot insert the genesis block");
 
         let mut cp = self.clone();
         let mut tail = vec![];
         let base = loop {
-            if cp.height() == block_id.height {
-                if cp.hash() == block_id.hash {
+            if cp.height() == height {
+                if cp.to_blockhash() == data.to_blockhash() {
                     return self;
                 }
                 // if we have a conflict we just return the inserted block because the tail is by
@@ -188,31 +300,50 @@ impl CheckPoint {
                 break cp.prev().expect("can't be called on genesis block");
             }
 
-            if cp.height() < block_id.height {
+            if cp.height() < height {
                 break cp;
             }
 
-            tail.push(cp.block_id());
+            tail.push(BlockId {
+                height,
+                hash: data.to_blockhash(),
+            });
             cp = cp.prev().expect("will break before genesis block");
         };
 
-        base.extend(core::iter::once(block_id).chain(tail.into_iter().rev()))
-            .expect("tail is in order")
+        let new_cp = core::iter::once((height, data)).chain(
+            tail.into_iter()
+                .rev()
+                .map(|block| (block.height, B::from(block.hash))),
+        );
+
+        base.extend_data(new_cp).expect("tail is in order")
     }
 
-    /// This method tests for `self` and `other` to have equal internal pointers.
-    pub fn eq_ptr(&self, other: &Self) -> bool {
-        Arc::as_ptr(&self.0) == Arc::as_ptr(&other.0)
+    /// Puts another checkpoint onto the linked list representing the blockchain.
+    ///
+    /// Returns an `Err(self)` if the block you are pushing on is not at a greater height that the one you
+    /// are pushing on to.
+    pub fn push_data(self, height: u32, data: B) -> Result<Self, Self> {
+        if self.height() < height {
+            Ok(Self(Arc::new(CPInner {
+                height,
+                data,
+                prev: Some(self.0),
+            })))
+        } else {
+            Err(self)
+        }
     }
 }
 
 /// Iterates over checkpoints backwards.
-pub struct CheckPointIter {
-    current: Option<Arc<CPInner>>,
+pub struct CheckPointIter<B = BlockHash> {
+    current: Option<Arc<CPInner<B>>>,
 }
 
-impl Iterator for CheckPointIter {
-    type Item = CheckPoint;
+impl<B> Iterator for CheckPointIter<B> {
+    type Item = CheckPoint<B>;
 
     fn next(&mut self) -> Option<Self::Item> {
         let current = self.current.clone()?;
@@ -221,9 +352,9 @@ impl Iterator for CheckPointIter {
     }
 }
 
-impl IntoIterator for CheckPoint {
-    type Item = CheckPoint;
-    type IntoIter = CheckPointIter;
+impl<B> IntoIterator for CheckPoint<B> {
+    type Item = CheckPoint<B>;
+    type IntoIter = CheckPointIter<B>;
 
     fn into_iter(self) -> Self::IntoIter {
         CheckPointIter {
