From 5983f166c94dd5ab172e38bf12a3330a3ed9004c Mon Sep 17 00:00:00 2001
From: Calvin Kim <calvin@kcalvinalvin.info>
Date: Sun, 2 Feb 2025 15:11:59 +0900
Subject: [PATCH] refactor: simplify GetAncestor

The if statement in GetAncestor was quite hard to make sense of.
Simplifying it improves readability and the added test ensures that the
performance remains the same.
---
 src/chain.cpp                 | 16 +++-----
 src/test/blockchain_tests.cpp | 71 +++++++++++++++++++++++++++++++++++
 2 files changed, 77 insertions(+), 10 deletions(-)

diff --git a/src/chain.cpp b/src/chain.cpp
index 82007a8a1e7..97b74656d01 100644
--- a/src/chain.cpp
+++ b/src/chain.cpp
@@ -96,24 +96,20 @@ const CBlockIndex* CBlockIndex::GetAncestor(int height) const
         return nullptr;
     }
 
+    // Traverse back until we find the desired pindex.
     const CBlockIndex* pindexWalk = this;
-    int heightWalk = nHeight;
-    while (heightWalk > height) {
-        int heightSkip = GetSkipHeight(heightWalk);
-        int heightSkipPrev = GetSkipHeight(heightWalk - 1);
+    while (pindexWalk != nullptr && pindexWalk->nHeight != height) {
+        // Prefer the ancestor if there is one we can take.
         if (pindexWalk->pskip != nullptr &&
-            (heightSkip == height ||
-             (heightSkip > height && !(heightSkipPrev < heightSkip - 2 &&
-                                       heightSkipPrev >= height)))) {
-            // Only follow pskip if pprev->pskip isn't better than pskip->pprev.
+                GetSkipHeight(pindexWalk->nHeight) >= height) {
             pindexWalk = pindexWalk->pskip;
-            heightWalk = heightSkip;
         } else {
+            // We couldn't take the ancestor skip so traverse back to the parent.
             assert(pindexWalk->pprev);
             pindexWalk = pindexWalk->pprev;
-            heightWalk--;
         }
     }
+
     return pindexWalk;
 }
 
diff --git a/src/test/blockchain_tests.cpp b/src/test/blockchain_tests.cpp
index 4ecc15041cc..56964645cfb 100644
--- a/src/test/blockchain_tests.cpp
+++ b/src/test/blockchain_tests.cpp
@@ -117,4 +117,75 @@ BOOST_AUTO_TEST_CASE(num_chain_tx_max)
     BOOST_CHECK_EQUAL(block_index.m_chain_tx_count, std::numeric_limits<uint64_t>::max());
 }
 
+/** Turn the lowest '1' bit in the binary representation of a number into a '0'. */
+int static inline InvertLowestOne(int n) { return n & (n - 1); }
+
+/** Compute what height to jump back to with the CBlockIndex::pskip pointer. */
+int static inline GetSkipHeight(int height) {
+    if (height < 2)
+        return 0;
+
+    // Determine which height to jump back to. Any number strictly lower than height is acceptable,
+    // but the following expression seems to perform well in simulations (max 110 steps to go back
+    // up to 2**18 blocks).
+    return (height & 1) ? InvertLowestOne(InvertLowestOne(height - 1)) + 1 : InvertLowestOne(height);
+}
+
+int static inline CountStepsOld(int start_height, int target_height) {
+    if (target_height > start_height || target_height < 0) {
+        return 0;
+    }
+
+    int count = 0;
+    int heightWalk = start_height;
+    while (heightWalk > target_height) {
+        int heightSkip = GetSkipHeight(heightWalk);
+        int heightSkipPrev = GetSkipHeight(heightWalk - 1);
+        if ((heightSkip == target_height ||
+             (heightSkip > target_height && !(heightSkipPrev < heightSkip - 2 &&
+                                       heightSkipPrev >= target_height)))) {
+            // Only follow pskip if pprev->pskip isn't better than pskip->pprev.
+            heightWalk = heightSkip;
+        } else {
+            heightWalk--;
+        }
+
+        count++;
+    }
+    return count;
+}
+
+int static inline CountStepsNew(int start_height, int target_height) {
+    if (target_height > start_height || target_height < 0) {
+        return 0;
+    }
+
+    // Traverse back until we find the desired pindex.
+    int count = 0;
+    while (start_height != target_height) {
+        auto skip_height = GetSkipHeight(start_height);
+        // Prefer the ancestor if there is one we can take.
+        if (skip_height >= target_height) {
+            start_height = skip_height;
+        } else {
+            // We couldn't take the ancestor skip so traverse back to the parent.
+            start_height--;
+        }
+        count++;
+    }
+
+    return count;
+}
+
+BOOST_AUTO_TEST_CASE(test_test)
+{
+    auto max_height = 1 << 25;
+    for (auto i = 0; i < max_height; i++) {
+        auto x = CountStepsOld(max_height, i);
+        auto y = CountStepsNew(max_height, i);
+        BOOST_CHECK_EQUAL(x, y);
+    }
+}
+
+
 BOOST_AUTO_TEST_SUITE_END()