3 files changed, 38 insertions, 8 deletions
diff --git a/src/cryptonote_basic/hardfork.cpp b/src/cryptonote_basic/hardfork.cpp
index 87a394918..447d79aee 100644
--- a/src/cryptonote_basic/hardfork.cpp
+++ b/src/cryptonote_basic/hardfork.cpp
@@ -222,7 +222,6 @@ bool HardFork::reorganize_from_block_height(uint64_t height)
   if (height >= db.height())
     return false;
 
-  db.set_batch_transactions(true);
   bool stop_batch = db.batch_start();
 
   versions.clear();
@@ -306,6 +305,29 @@ bool HardFork::rescan_from_chain_height(uint64_t height)
   return rescan_from_block_height(height - 1);
 }
 
+void HardFork::on_block_popped(uint64_t nblocks)
+{
+  CHECK_AND_ASSERT_THROW_MES(nblocks > 0, "nblocks must be greater than 0");
+
+  CRITICAL_REGION_LOCAL(lock);
+
+  const uint64_t new_chain_height = db.height();
+  const uint64_t old_chain_height = new_chain_height + nblocks;
+  uint8_t version;
+  uint64_t height;
+  for (height = old_chain_height - 1; height >= new_chain_height; --height)
+  {
+    versions.pop_back();
+    version = db.get_hard_fork_version(height);
+    versions.push_front(version);
+  }
+
+  // does not take voting into account
+  for (current_fork_index = heights.size() - 1; current_fork_index > 0; --current_fork_index)
+    if (height >= heights[current_fork_index].height)
+      break;
+}
+
 int HardFork::get_voted_fork_index(uint64_t height) const
 {
   CRITICAL_REGION_LOCAL(lock);
diff --git a/src/cryptonote_basic/hardfork.h b/src/cryptonote_basic/hardfork.h
index a63a66976..a3fc25dfa 100644
--- a/src/cryptonote_basic/hardfork.h
+++ b/src/cryptonote_basic/hardfork.h
@@ -150,6 +150,16 @@ namespace cryptonote
     bool reorganize_from_chain_height(uint64_t height);
 
     /**
+     * @brief called when one or more blocks are popped from the blockchain
+     *
+     * The current fork will be updated by looking up the db,
+     * which is much cheaper than recomputing everything
+     *
+     * @param new_chain_height the height of the chain after popping
+     */
+    void on_block_popped(uint64_t new_chain_height);
+
+    /**
      * @brief returns current state at the given time
      *
      * Based on the approximate time of the last known hard fork,
diff --git a/src/cryptonote_core/blockchain.cpp b/src/cryptonote_core/blockchain.cpp
index 19dec069c..d8163721e 100644
--- a/src/cryptonote_core/blockchain.cpp
+++ b/src/cryptonote_core/blockchain.cpp
@@ -641,6 +641,9 @@ block Blockchain::pop_block_from_blockchain()
     throw;
   }
 
+  // make sure the hard fork object updates its current version
+  m_hardfork->on_block_popped(1);
+
   // return transactions from popped block to the tx_pool
   for (transaction& tx : popped_txs)
   {
@@ -651,12 +654,7 @@ block Blockchain::pop_block_from_blockchain()
       // FIXME: HardFork
       // Besides the below, popping a block should also remove the last entry
       // in hf_versions.
-      //
-      // FIXME: HardFork
-      // This is not quite correct, as we really want to add the txes
-      // to the pool based on the version determined after all blocks
-      // are popped.
-      uint8_t version = get_current_hard_fork_version();
+      uint8_t version = get_ideal_hard_fork_version(m_db->height());
 
       // We assume that if they were in a block, the transactions are already
       // known to the network as a whole. However, if we had mined that block,
@@ -1201,7 +1199,7 @@ bool Blockchain::validate_miner_transaction(const block& b, size_t cumulative_bl
     return false;
   }
   // From hard fork 2, we allow a miner to claim less block reward than is allowed, in case a miner wants less dust
-  if (m_hardfork->get_current_version() < 2)
+  if (version < 2)
   {
     if(base_reward + fee != money_in_use)
     {