feat: get_main_chain_length

canister/src/blocktree.rs

@@ -3,6 +3,7 @@ use ic_btc_interface::Network;
 use ic_btc_types::{Block, BlockHash};
 use std::fmt;
 mod serde;
+use std::cmp::Ordering;
 use std::ops::{Add, Sub};
 
 /// Represents a non-empty block chain as:
@@ -381,6 +382,38 @@ impl BlockTree {
         }
         blocks
     }
+
+    /// Returns the length of the main chain based on the heaviest cumulative difficulty.
+    pub fn get_main_chain_length(&self, network: Network) -> usize {
+        let (_, depth) = self.find_main_chain(network);
+        depth.get() as usize
+    }
+
+    /// Recursively finds the heaviest chain and its common ancestor height.
+    fn find_main_chain(&self, network: Network) -> (DifficultyBasedDepth, Depth) {
+        let mut max_difficulty_based_depth = DifficultyBasedDepth::new(0);
+        let mut max_depth = Depth::new(0);
+
+        for child in &self.children {
+            let (difficulty_based_depth, depth) = child.find_main_chain(network);
+            match max_difficulty_based_depth.cmp(&difficulty_based_depth) {
+                Ordering::Less => {
+                    max_difficulty_based_depth = difficulty_based_depth;
+                    max_depth = depth;
+                }
+                Ordering::Equal => {
+                    max_difficulty_based_depth = difficulty_based_depth;
+                    max_depth = Depth::new(0);
+                }
+                Ordering::Greater => {}
+            }
+        }
+
+        (
+            max_difficulty_based_depth + DifficultyBasedDepth::new(self.root.difficulty(network)),
+            max_depth + Depth::new(1)
+        )
+    }
 }
 
 /// An error thrown when trying to add a block that isn't a successor

canister/src/unstable_blocks.rs

@@ -209,6 +209,11 @@ impl UnstableBlocks {
             .collect::<Vec<_>>()
             .into_iter()
     }
+
+    /// Returns the length of the main chain based on the heaviest cumulative difficulty.
+    pub fn get_main_chain_length(&self) -> usize {
+        self.tree.get_main_chain_length(self.network)
+    }
 }
 
 /// Returns a reference to the `anchor` block iff ∃ a child `C` of `anchor` that is stable.
@@ -318,17 +323,19 @@ pub fn get_main_chain(blocks: &UnstableBlocks) -> BlockChain {
 /// Returns the length of the "main chain".
 /// See `get_main_chain` for what defines a main chain.
 pub fn get_main_chain_length(blocks: &UnstableBlocks) -> usize {
-    let blocks_by_height = blocks.blocks_with_depths_by_heights();
+    // let blocks_by_height = blocks.blocks_with_depths_by_heights();
 
-    // Traverse the heights in reverse order. The highest height with a single block corresponds to
-    // the tip of the main chain.
-    for height in (0..blocks_by_height.len()).rev() {
-        if blocks_by_height[height].len() == 1 {
-            return height + 1;
-        }
-    }
+    // // Traverse the heights in reverse order. The highest height with a single block corresponds to
+    // // the tip of the main chain.
+    // for height in (0..blocks_by_height.len()).rev() {
+    //     if blocks_by_height[height].len() == 1 {
+    //         return height + 1;
+    //     }
+    // }
+
+    // unreachable!("There must be at least one height with exactly one block.");
 
-    unreachable!("There must be at least one height with exactly one block.");
+    blocks.get_main_chain_length()
 }
 
 pub fn get_block_hashes(blocks: &UnstableBlocks) -> Vec<BlockHash> {