refactor(btc-headerchain): optimize for zkvm, add docs/comments

crates/prover/bitcoin-blockspace/src/block.rs

@@ -65,7 +65,7 @@ pub fn compute_witness_commitment(
 /// Returns the block hash.
 ///
 /// Equivalent to [`compute_block_hash`](Header::block_hash)
-fn compute_block_hash(header: &Header) -> [u8; 32] {
+pub fn compute_block_hash(header: &Header) -> [u8; 32] {
     let mut vec = Vec::with_capacity(80);
     header
         .consensus_encode(&mut vec)
@@ -129,9 +129,9 @@ pub fn check_witness_commitment(block: &Block) -> bool {
 }
 
 /// Checks that the proof-of-work for the block is valid.
-pub fn check_pow(block: &Block) -> bool {
-    let target = block.header.target();
-    let block_hash = BlockHash::from_byte_array(compute_block_hash(&block.header));
+pub fn check_pow(block: &Header) -> bool {
+    let target = block.target();
+    let block_hash = BlockHash::from_byte_array(compute_block_hash(block));
     target.is_met_by(block_hash)
 }
 
@@ -174,6 +174,6 @@ mod tests {
         assert!(check_witness_commitment(&block));
 
         assert!(block.header.validate_pow(block.header.target()).is_ok());
-        assert!(check_pow(&block));
+        assert!(check_pow(&block.header));
     }
 }

crates/prover/btc-headerchain/Cargo.toml

@@ -6,10 +6,12 @@ edition = "2021"
 [dependencies]
 anyhow.workspace = true
 bitcoin.workspace = true
+bitcoin-blockspace.workspace = true
 borsh.workspace = true
 ethnum = "1.5.0"
 serde.workspace = true
 serde_json = "1.0.127"
 
 [dev-dependencies]
 alpen-test-utils.workspace = true
+rand.workspace = true

crates/prover/btc-headerchain/src/lib.rs

@@ -1,27 +1,30 @@
-use bitcoin::{block::Header, hashes::Hash, CompactTarget, Target};
+use bitcoin::{block::Header, hashes::Hash, BlockHash, CompactTarget, Target};
+use bitcoin_blockspace::block::compute_block_hash;
 use ethnum::U256;
 
-/// Difficulty recalculation interval
+/// Difficulty recalculation interval.
 /// On [MAINNET](bitcoin::consensus::params::MAINNET), it is around 2 weeks
 const POW_TARGET_TIMESPAN: u32 = 14 * 24 * 60 * 60;
 
-/// Expected amount of time to mine one block
+/// Expected amount of time to mine one block.
 /// On [MAINNET](bitcoin::consensus::params::MAINNET), it is around 10 minutes
 const POW_TARGET_SPACING: u32 = 10 * 60;
 
+/// No of blocks after which the difficulty is adjusted.
+/// [bitcoin::consensus::params::Params::difficulty_adjustment_interval].
 const DIFFICULTY_ADJUSTMENT_INTERVAL: u32 = POW_TARGET_TIMESPAN / POW_TARGET_SPACING;
 
 #[derive(Debug, Clone)]
 pub struct HeaderVerificationState {
     /// [Block number](bitcoin::Block::bip34_block_height) of the last verified block
     pub last_verified_block_num: u32,
 
-    /// [Target](bitcoin::pow::CompactTarget) of the last verified block
-    pub last_verified_block_target: u32,
-
     /// [Hash](bitcoin::block::Header::block_hash) of the last verified block
     pub last_verified_block_hash: [u8; 32],
 
+    /// [Target](bitcoin::pow::CompactTarget) for the next block to verify
+    pub next_block_target: u32,
+
     /// Timestamp of the block at the start of a [difficulty adjustment
     /// interval](bitcoin::consensus::params::Params::difficulty_adjustment_interval).
     ///
@@ -33,19 +36,13 @@ pub struct HeaderVerificationState {
     /// (e.g., block 0, 2016, 4032, etc.).
     pub interval_start_timestamp: u32,
 
-    /// Timestamp of the block at the end of a difficulty adjustment interval.
-    ///
-    /// On [MAINNET], the interval ends at blocks with heights 2015, 4031, 6047, 8063, etc.
-    ///
-    /// This field represents the timestamp of the last block in the interval
-    /// (e.g., block 2015, 4031, 6047, etc.).
-    pub interval_end_timestamp: u32,
-
     /// Total accumulated [difficulty](bitcoin::pow::Target::difficulty_float)
     /// TODO: check if using [this](bitcoin::pow::Target::difficulty) makes more sense
     pub total_accumulated_pow: f64,
 
-    /// Timestamp of the last 11 blocks
+    /// Timestamps of the last 11 blocks in descending order.
+    /// The timestamp of the most recent block is at index 0, while the timestamp of the oldest
+    /// block is at index 10.
     pub last_11_blocks_timestamps: [u32; 11],
 }
 
@@ -56,13 +53,6 @@ impl HeaderVerificationState {
         timestamps[5]
     }
 
-    fn insert_timestamp(&mut self, timestamp: u32) {
-        for i in (1..11).rev() {
-            self.last_11_blocks_timestamps[i] = self.last_11_blocks_timestamps[i - 1];
-        }
-        self.last_11_blocks_timestamps[0] = timestamp;
-    }
-
     /// Computes the [`CompactTarget`] from a difficulty adjustment.
     ///
     /// ref: <https://github.com/bitcoin/bitcoin/blob/0503cbea9aab47ec0a87d34611e5453158727169/src/pow.cpp>
@@ -86,23 +76,20 @@ impl HeaderVerificationState {
     /// # Returns
     ///
     /// The expected [`CompactTarget`] recalculation.
-    fn next_target(&mut self) -> u32 {
+    fn next_target(&mut self, timestamp: u32) -> u32 {
         if (self.last_verified_block_num + 1) % DIFFICULTY_ADJUSTMENT_INTERVAL != 0 {
-            return self.last_verified_block_target;
+            return self.next_block_target;
         }
 
-        // Comments relate to the `pow.cpp` file from Core.
-        // ref: <https://github.com/bitcoin/bitcoin/blob/0503cbea9aab47ec0a87d34611e5453158727169/src/pow.cpp>
-        let min_timespan = POW_TARGET_TIMESPAN >> 2; // Lines 56/57
-        let max_timespan = POW_TARGET_TIMESPAN << 2; // Lines 58/59
+        let min_timespan = POW_TARGET_TIMESPAN >> 2;
+        let max_timespan = POW_TARGET_TIMESPAN << 2;
 
-        let timespan = self.interval_end_timestamp - self.interval_start_timestamp;
+        let timespan = timestamp - self.interval_start_timestamp;
         let actual_timespan = timespan.clamp(min_timespan, max_timespan);
 
-        let prev_target: Target =
-            CompactTarget::from_consensus(self.last_verified_block_target).into();
+        let prev_target: Target = CompactTarget::from_consensus(self.next_block_target).into();
 
-        let mut retarget = U256::from_le_bytes(prev_target.to_le_bytes()); // bnNew
+        let mut retarget = U256::from_le_bytes(prev_target.to_le_bytes());
         retarget *= U256::from(actual_timespan);
         retarget /= U256::from(POW_TARGET_TIMESPAN);
 
@@ -117,7 +104,19 @@ impl HeaderVerificationState {
         retarget.to_compact_lossy().to_consensus()
     }
 
-    // fn insert_timestamp(&mut self)
+    fn update_timestamps(&mut self, timestamp: u32) {
+        // Shift existing timestamps to right
+        for i in (1..11).rev() {
+            self.last_11_blocks_timestamps[i] = self.last_11_blocks_timestamps[i - 1];
+        }
+        // Insert latest timestamp at index 0
+        self.last_11_blocks_timestamps[0] = timestamp;
+
+        let new_block_num = self.last_verified_block_num;
+        if new_block_num % DIFFICULTY_ADJUSTMENT_INTERVAL == 0 {
+            self.interval_start_timestamp = timestamp;
+        }
+    }
 
     pub fn check_and_update(&mut self, header: Header) {
         // Check continuity
@@ -126,66 +125,91 @@ impl HeaderVerificationState {
             self.last_verified_block_hash,
         );
 
+        let block_hash_raw = compute_block_hash(&header);
+        let block_hash = BlockHash::from_byte_array(block_hash_raw);
+
         // Check PoW
-        assert!(header.validate_pow(header.target()).is_ok());
+        assert_eq!(header.bits.to_consensus(), self.next_block_target);
+        header.target().is_met_by(block_hash);
 
         // Check timestamp
         assert!(header.time > self.get_median_timestamp());
 
-        let updated_target = self.next_target();
-        assert_eq!(header.bits.to_consensus(), updated_target);
+        // Increase the last verified block number by 1
+        self.last_verified_block_num += 1;
 
-        if (self.last_verified_block_num + 1) % DIFFICULTY_ADJUSTMENT_INTERVAL == 0 {
-            self.interval_start_timestamp = header.time;
-        } else if (self.last_verified_block_num + 1) % DIFFICULTY_ADJUSTMENT_INTERVAL
-            == DIFFICULTY_ADJUSTMENT_INTERVAL - 1
-        {
-            self.interval_end_timestamp = header.time;
-        }
+        // Set the header block hash as the last verified block hash
+        self.last_verified_block_hash = block_hash_raw;
 
-        self.last_verified_block_hash = header.block_hash().to_byte_array();
-        self.last_verified_block_target = updated_target;
-        self.insert_timestamp(header.time);
-        self.last_verified_block_num += 1;
+        // Update the timestamps
+        self.update_timestamps(header.time);
+
+        // Update the total accumulated PoW
+        self.total_accumulated_pow += header.difficulty_float();
+
+        // Set the target for the next block
+        self.next_block_target = self.next_target(header.time);
     }
 }
 
 #[cfg(test)]
 mod tests {
     use alpen_test_utils::bitcoin::get_btc_chain;
     use bitcoin::hashes::Hash;
+    use rand::Rng;
 
     use super::HeaderVerificationState;
+    use crate::DIFFICULTY_ADJUSTMENT_INTERVAL;
+
+    /// Calculates the height at which a specific difficulty adjustment occurs relative to a
+    /// starting height.
+    ///
+    /// # Arguments
+    ///
+    /// * `idx` - The index of the difficulty adjustment (1-based). 1 for the first adjustment, 2
+    ///   for the second, and so on.
+    /// * `start` - The starting height from which to calculate.
+    fn get_difficulty_adjustment_height(idx: u32, start: u32) -> u32 {
+        ((start / DIFFICULTY_ADJUSTMENT_INTERVAL) + idx) * DIFFICULTY_ADJUSTMENT_INTERVAL
+    }
 
     #[test]
     fn test_blocks() {
         let chain = get_btc_chain();
 
-        let height1 = ((chain.start / 2016) + 1) * 2016;
-        let height2 = ((chain.start / 2016) + 2) * 2016;
+        // Start from the first difficulty adjustment block after `chain.start`
+        // This ensures we have a known difficulty adjustment point and a set
+        // `interval_start_timestamp`
+        let h1 = get_difficulty_adjustment_height(1, chain.start);
+
+        // Get the second difficulty adjustment block after `chain.start`
+        let h2 = get_difficulty_adjustment_height(2, chain.start);
+
+        // Set the random block between h1 and h2 as the last verified block
+        let r1 = rand::thread_rng().gen_range(h1..h2 - 1);
+        let last_verified_block = chain.get_block(r1);
 
+        // Fetch the previous timestamps of block from `r1`
+        // This fetches timestamps of `r1`, `r1-1`, `r1-2`, ...
         let recent_block_timestamp: [u32; 11] =
-            chain.get_last_timestamps(height2, 11).try_into().unwrap();
+            chain.get_last_timestamps(r1, 11).try_into().unwrap();
 
         let mut verification_state = HeaderVerificationState {
-            last_verified_block_num: height2 - 1,
-            last_verified_block_hash: chain
-                .get_block(height2 - 1)
+            last_verified_block_num: r1,
+            last_verified_block_hash: last_verified_block
                 .block_hash()
                 .as_raw_hash()
                 .to_byte_array(),
-            last_verified_block_target: chain
-                .get_block(height2 - 1)
+            next_block_target: last_verified_block
                 .target()
                 .to_compact_lossy()
                 .to_consensus(),
-            interval_start_timestamp: chain.get_block(height1).time,
-            interval_end_timestamp: chain.get_block(height2 - 1).time,
+            interval_start_timestamp: chain.get_block(h1).time,
             total_accumulated_pow: 0f64,
             last_11_blocks_timestamps: recent_block_timestamp,
         };
 
-        for header_idx in height2..chain.end {
+        for header_idx in (r1 + 1)..chain.end {
             verification_state.check_and_update(chain.get_block(header_idx))
         }
     }

crates/test-utils/src/bitcoin.rs

@@ -40,6 +40,7 @@ pub struct BtcChain {
 }
 
 impl BtcChain {
+    /// Retrieves the block header at the specified height.
     pub fn get_block(&self, height: u32) -> Header {
         if height < self.start {
             panic!("height must be greater than that");
@@ -51,10 +52,12 @@ impl BtcChain {
         self.headers[idx as usize]
     }
 
-    pub fn get_last_timestamps(&self, before: u32, count: u32) -> Vec<u32> {
+    /// Retrieves the timestamps of a specified number of blocks from a given height in a
+    /// descending order.
+    pub fn get_last_timestamps(&self, from: u32, count: u32) -> Vec<u32> {
         let mut timestamps = Vec::with_capacity(count as usize);
-        for i in (1..count + 1).rev() {
-            let h = self.get_block(before - i);
+        for i in (0..count).rev() {
+            let h = self.get_block(from - i);
             timestamps.push(h.time)
         }
         timestamps