fix: correct epoch counting

applications/minotari_node/src/commands/command/list_validator_nodes.rs

@@ -62,19 +62,16 @@ impl CommandContext {
     /// Function to process the list-connections command
     pub async fn list_validator_nodes(&mut self, args: Args) -> Result<(), Error> {
         let metadata = self.blockchain_db.get_chain_metadata().await?;
-        let constants = self
-            .consensus_rules
-            .consensus_constants(metadata.height_of_longest_chain());
         let height = args
             .epoch
-            .map(|epoch| constants.epoch_to_block_height(epoch))
+            .map(|epoch| self.consensus_rules.epoch_to_block_height(epoch))
             .unwrap_or_else(|| metadata.height_of_longest_chain());
         let vns = self.blockchain_db.fetch_active_validator_nodes(height).await?;
 
         println!();
         println!(
             "Registered validator nodes for epoch {}",
-            constants.block_height_to_epoch(height).as_u64()
+            self.consensus_rules.block_height_to_epoch(height).as_u64()
         );
         println!("----------------------------------");
         if vns.is_empty() {

base_layer/core/src/chain_storage/lmdb_db/lmdb_db.rs

@@ -1400,24 +1400,21 @@ impl LMDBDatabase {
     ) -> Result<(), ChainStorageError> {
         let store = self.validator_node_store(txn);
         let constants = self.get_consensus_constants(header.height);
-        let current_epoch = constants.block_height_to_epoch(header.height);
-
-        let prev_shard_key = store.get_shard_key(
-            current_epoch
-                .as_u64()
-                .saturating_sub(constants.validator_node_validity_period_epochs().as_u64()) *
-                constants.epoch_length(),
-            current_epoch.as_u64() * constants.epoch_length(),
-            vn_reg.public_key(),
-        )?;
+        let current_epoch = self.block_height_to_epoch(header.height);
+        // TODO: What if the validity period has changed?
+        let start_height =
+            self.epoch_to_block_height(current_epoch.saturating_sub(constants.validator_node_validity_period_epochs()));
+        let end_height = self.epoch_to_block_height(current_epoch);
+
+        let prev_shard_key = store.get_shard_key(start_height, end_height, vn_reg.public_key())?;
         let shard_key = vn_reg.derive_shard_key(
             prev_shard_key,
             current_epoch,
             constants.validator_node_registration_shuffle_interval(),
             &header.prev_hash,
         );
 
-        let next_epoch = constants.block_height_to_epoch(header.height) + VnEpoch(1);
+        let next_epoch = current_epoch + VnEpoch(1);
         let validator_node = ValidatorNodeEntry {
             shard_key,
             start_epoch: next_epoch,
@@ -1687,6 +1684,14 @@ impl LMDBDatabase {
     fn get_consensus_constants(&self, height: u64) -> &ConsensusConstants {
         self.consensus_manager.consensus_constants(height)
     }
+
+    fn block_height_to_epoch(&self, height: u64) -> VnEpoch {
+        self.consensus_manager.block_height_to_epoch(height)
+    }
+
+    fn epoch_to_block_height(&self, epoch: VnEpoch) -> u64 {
+        self.consensus_manager.epoch_to_block_height(epoch)
+    }
 }
 
 pub fn create_recovery_lmdb_database<P: AsRef<Path>>(path: P) -> Result<(), ChainStorageError> {
@@ -2450,7 +2455,7 @@ impl BlockchainBackend for LMDBDatabase {
         let constants = self.consensus_manager.consensus_constants(height);
 
         // Get the current epoch for the height
-        let end_epoch = constants.block_height_to_epoch(height);
+        let end_epoch = self.consensus_manager.block_height_to_epoch(height);
         // Subtract the registration validaty period to get the start epoch
         let start_epoch = end_epoch.saturating_sub(constants.validator_node_validity_period_epochs());
         // Convert these back to height as validators regs are indexed by height
@@ -2466,7 +2471,7 @@ impl BlockchainBackend for LMDBDatabase {
         let constants = self.get_consensus_constants(height);
 
         // Get the epoch height boundaries for our query
-        let current_epoch = constants.block_height_to_epoch(height);
+        let current_epoch = self.consensus_manager.block_height_to_epoch(height);
         let start_epoch = current_epoch.saturating_sub(constants.validator_node_validity_period_epochs());
         let start_height = start_epoch.as_u64() * constants.epoch_length();
         let end_height = current_epoch.as_u64() * constants.epoch_length();

base_layer/core/src/consensus/consensus_constants.rs

@@ -330,16 +330,6 @@ impl ConsensusConstants {
         self.vn_registration_lock_height
     }
 
-    /// Returns the current epoch from the given height
-    pub fn block_height_to_epoch(&self, height: u64) -> VnEpoch {
-        VnEpoch(height / self.vn_epoch_length)
-    }
-
-    /// Returns the block height of the start of the given epoch
-    pub fn epoch_to_block_height(&self, epoch: VnEpoch) -> u64 {
-        epoch.as_u64() * self.vn_epoch_length
-    }
-
     pub fn epoch_length(&self) -> u64 {
         self.vn_epoch_length
     }
@@ -848,6 +838,16 @@ impl ConsensusConstantsBuilder {
         self
     }
 
+    pub fn with_effective_height(mut self, height: u64) -> Self {
+        self.consensus.effective_from_height = height;
+        self
+    }
+
+    pub fn with_vn_epoch_length(mut self, length: u64) -> Self {
+        self.consensus.vn_epoch_length = length;
+        self
+    }
+
     pub fn build(self) -> ConsensusConstants {
         self.consensus
     }

base_layer/core/src/consensus/consensus_manager.rs

@@ -23,6 +23,7 @@
 use std::sync::Arc;
 
 use tari_common::configuration::Network;
+use tari_common_types::epoch::VnEpoch;
 use thiserror::Error;
 
 #[cfg(feature = "base_node")]
@@ -112,6 +113,51 @@ impl ConsensusManager {
         constants
     }
 
+    /// Returns the current epoch from the given height
+    pub fn block_height_to_epoch(&self, height: u64) -> VnEpoch {
+        let mut epoch = 0;
+        let mut leftover_height = 0;
+        let mut active_effective_height = 0;
+        let mut active_epoch_length = self.inner.consensus_constants[0].epoch_length();
+        for c in &self.inner.consensus_constants[1..] {
+            if c.effective_from_height() > height {
+                break;
+            }
+            epoch += (c.effective_from_height() - active_effective_height + leftover_height) / active_epoch_length;
+            leftover_height = std::cmp::min(
+                c.epoch_length(),
+                (c.effective_from_height() - active_effective_height + leftover_height) % active_epoch_length,
+            );
+            active_effective_height = c.effective_from_height();
+            active_epoch_length = c.epoch_length();
+        }
+        epoch += (height - active_effective_height + leftover_height) / active_epoch_length;
+        VnEpoch(epoch)
+    }
+
+    /// Returns the block height of the start of the given epoch
+    pub fn epoch_to_block_height(&self, epoch: VnEpoch) -> u64 {
+        let mut cur_epoch = 0;
+        let mut leftover_height = 0;
+        let mut active_effective_height = 0;
+        let mut active_epoch_length = self.inner.consensus_constants[0].epoch_length();
+        for c in &self.inner.consensus_constants[1..] {
+            if cur_epoch + (c.effective_from_height() - active_effective_height + leftover_height) / active_epoch_length >
+                epoch.as_u64()
+            {
+                break;
+            }
+            cur_epoch += (c.effective_from_height() - active_effective_height + leftover_height) / active_epoch_length;
+            leftover_height = std::cmp::min(
+                c.epoch_length(),
+                (c.effective_from_height() - active_effective_height + leftover_height) % active_epoch_length,
+            );
+            active_effective_height = c.effective_from_height();
+            active_epoch_length = c.epoch_length();
+        }
+        (epoch.as_u64() - cur_epoch) * active_epoch_length + active_effective_height - leftover_height
+    }
+
     /// Create a new TargetDifficulty for the given proof of work using constants that are effective from the given
     /// height
     #[cfg(feature = "base_node")]
@@ -271,3 +317,64 @@ pub enum ConsensusBuilderError {
     #[error("Cannot set a genesis block with a network other than LocalNet")]
     CannotSetGenesisBlock,
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::consensus::ConsensusConstantsBuilder;
+
+    #[test]
+    fn test_epoch_to_height_and_back() {
+        let manager = ConsensusManager::builder(Network::LocalNet)
+            .add_consensus_constants(
+                ConsensusConstantsBuilder::new(Network::LocalNet)
+                    .with_effective_height(0)
+                    .with_vn_epoch_length(15)
+                    .build(),
+            )
+            .add_consensus_constants(
+                ConsensusConstantsBuilder::new(Network::LocalNet)
+                    .with_effective_height(100)
+                    .with_vn_epoch_length(6)
+                    .build(),
+            )
+            .add_consensus_constants(
+                ConsensusConstantsBuilder::new(Network::LocalNet)
+                    .with_effective_height(200)
+                    .with_vn_epoch_length(8)
+                    .build(),
+            )
+            .add_consensus_constants(
+                ConsensusConstantsBuilder::new(Network::LocalNet)
+                    .with_effective_height(300)
+                    .with_vn_epoch_length(13)
+                    .build(),
+            )
+            .add_consensus_constants(
+                ConsensusConstantsBuilder::new(Network::LocalNet)
+                    .with_effective_height(400)
+                    .with_vn_epoch_length(17)
+                    .build(),
+            )
+            .add_consensus_constants(
+                ConsensusConstantsBuilder::new(Network::LocalNet)
+                    .with_effective_height(500)
+                    .with_vn_epoch_length(7)
+                    .build(),
+            )
+            .build()
+            .unwrap();
+        assert_eq!(manager.block_height_to_epoch(99), VnEpoch(6)); // The next epoch should change at 105
+        assert_eq!(manager.block_height_to_epoch(100), VnEpoch(7)); // But with the new length the epoch should change right away
+        assert_eq!(manager.block_height_to_epoch(199), VnEpoch(23)); // The next epoch should change at 202
+        assert_eq!(manager.block_height_to_epoch(202), VnEpoch(23)); // But we have new length with size +2 so the epoch change will happen at 204
+        assert_eq!(manager.block_height_to_epoch(204), VnEpoch(24));
+        // Now test couple more back and forth
+        for epoch in 0..=100 {
+            assert_eq!(
+                manager.block_height_to_epoch(manager.epoch_to_block_height(VnEpoch(epoch))),
+                VnEpoch(epoch)
+            );
+        }
+    }
+}