From 07d3f947ec3725640bf07bdd0de34f910aaca349 Mon Sep 17 00:00:00 2001
From: Jack Grigg <jack@z.cash>
Date: Wed, 12 Sep 2018 10:07:44 +0100
Subject: [PATCH] Extract a helper method for finding the next epoch

---
 src/consensus/upgrades.cpp  | 20 ++++++++++++++------
 src/consensus/upgrades.h    |  6 ++++++
 src/gtest/test_upgrades.cpp | 29 +++++++++++++++++++++++++++++
 3 files changed, 49 insertions(+), 6 deletions(-)

diff --git a/src/consensus/upgrades.cpp b/src/consensus/upgrades.cpp
index 95bde7925..d346f90b1 100644
--- a/src/consensus/upgrades.cpp
+++ b/src/consensus/upgrades.cpp
@@ -114,20 +114,28 @@ bool IsActivationHeightForAnyUpgrade(
     return false;
 }
 
-boost::optional<int> NextActivationHeight(
-    int nHeight,
-    const Consensus::Params& params)
-{
+boost::optional<int> NextEpoch(int nHeight, const Consensus::Params& params) {
     if (nHeight < 0) {
         return boost::none;
     }
 
-    // Don't count Sprout as an activation height
+    // Sprout is never pending
     for (auto idx = Consensus::BASE_SPROUT + 1; idx < Consensus::MAX_NETWORK_UPGRADES; idx++) {
         if (NetworkUpgradeState(nHeight, params, Consensus::UpgradeIndex(idx)) == UPGRADE_PENDING) {
-            return params.vUpgrades[idx].nActivationHeight;
+            return idx;
         }
     }
 
     return boost::none;
 }
+
+boost::optional<int> NextActivationHeight(
+    int nHeight,
+    const Consensus::Params& params)
+{
+    auto idx = NextEpoch(nHeight, params);
+    if (idx) {
+        return params.vUpgrades[idx.get()].nActivationHeight;
+    }
+    return boost::none;
+}
diff --git a/src/consensus/upgrades.h b/src/consensus/upgrades.h
index 620dc94c4..913a316cb 100644
--- a/src/consensus/upgrades.h
+++ b/src/consensus/upgrades.h
@@ -79,6 +79,12 @@ bool IsActivationHeightForAnyUpgrade(
     int nHeight,
     const Consensus::Params& params);
 
+/**
+ * Returns the index of the next upgrade after the given block height, or
+ * boost::none if there are no more known upgrades.
+ */
+boost::optional<int> NextEpoch(int nHeight, const Consensus::Params& params);
+
 /**
  * Returns the activation height for the next upgrade after the given block height,
  * or boost::none if there are no more known upgrades.
diff --git a/src/gtest/test_upgrades.cpp b/src/gtest/test_upgrades.cpp
index 1066a28d0..54fa48832 100644
--- a/src/gtest/test_upgrades.cpp
+++ b/src/gtest/test_upgrades.cpp
@@ -143,6 +143,35 @@ TEST_F(UpgradesTest, IsActivationHeightForAnyUpgrade) {
     EXPECT_FALSE(IsActivationHeightForAnyUpgrade(1000000, params));
 }
 
+TEST_F(UpgradesTest, NextEpoch) {
+    SelectParams(CBaseChainParams::REGTEST);
+    const Consensus::Params& params = Params().GetConsensus();
+
+    // Consensus::NetworkUpgrade::NO_ACTIVATION_HEIGHT
+    EXPECT_EQ(NextEpoch(-1, params), boost::none);
+    EXPECT_EQ(NextEpoch(0, params), boost::none);
+    EXPECT_EQ(NextEpoch(1, params), boost::none);
+    EXPECT_EQ(NextEpoch(1000000, params), boost::none);
+
+    UpdateNetworkUpgradeParameters(Consensus::UPGRADE_TESTDUMMY, Consensus::NetworkUpgrade::ALWAYS_ACTIVE);
+
+    EXPECT_EQ(NextEpoch(-1, params), boost::none);
+    EXPECT_EQ(NextEpoch(0, params), boost::none);
+    EXPECT_EQ(NextEpoch(1, params), boost::none);
+    EXPECT_EQ(NextEpoch(1000000, params), boost::none);
+
+    int nActivationHeight = 100;
+    UpdateNetworkUpgradeParameters(Consensus::UPGRADE_TESTDUMMY, nActivationHeight);
+
+    EXPECT_EQ(NextEpoch(-1, params), boost::none);
+    EXPECT_EQ(NextEpoch(0, params), static_cast<int>(Consensus::UPGRADE_TESTDUMMY));
+    EXPECT_EQ(NextEpoch(1, params), static_cast<int>(Consensus::UPGRADE_TESTDUMMY));
+    EXPECT_EQ(NextEpoch(nActivationHeight - 1, params), static_cast<int>(Consensus::UPGRADE_TESTDUMMY));
+    EXPECT_EQ(NextEpoch(nActivationHeight, params), boost::none);
+    EXPECT_EQ(NextEpoch(nActivationHeight + 1, params), boost::none);
+    EXPECT_EQ(NextEpoch(1000000, params), boost::none);
+}
+
 TEST_F(UpgradesTest, NextActivationHeight) {
     SelectParams(CBaseChainParams::REGTEST);
     const Consensus::Params& params = Params().GetConsensus();