refactor: add EpochList type

Wrapper struct to make the accessing of the individual epochs easier and
standardized.

Author: obycode

stacks-common/src/types/mod.rs

@@ -1,5 +1,6 @@
 use std::cmp::Ordering;
 use std::fmt;
+use std::ops::{Deref, DerefMut, Index, IndexMut};
 
 #[cfg(feature = "canonical")]
 pub mod sqlite;
@@ -235,21 +236,6 @@ impl StacksEpochId {
             StacksEpochId::Epoch30 => cur_reward_cycle > first_epoch30_reward_cycle,
         }
     }
-
-    /// Return the index for this epoch in the list of epochs
-    pub fn index(&self) -> usize {
-        match self {
-            StacksEpochId::Epoch10 => 0,
-            StacksEpochId::Epoch20 => 1,
-            StacksEpochId::Epoch2_05 => 2,
-            StacksEpochId::Epoch21 => 3,
-            StacksEpochId::Epoch22 => 4,
-            StacksEpochId::Epoch23 => 5,
-            StacksEpochId::Epoch24 => 6,
-            StacksEpochId::Epoch25 => 7,
-            StacksEpochId::Epoch30 => 8,
-        }
-    }
 }
 
 impl std::fmt::Display for StacksEpochId {
@@ -475,3 +461,58 @@ impl<L: PartialEq + Eq> Ord for StacksEpoch<L> {
         self.epoch_id.cmp(&other.epoch_id)
     }
 }
+
+/// A wrapper for holding a list of Epochs, indexable by StacksEpochId
+#[derive(Clone, Debug, Default, Deserialize, PartialEq, Eq)]
+pub struct EpochList<L: Clone>(Vec<StacksEpoch<L>>);
+
+impl<L: Clone> EpochList<L> {
+    pub fn new(epochs: &[StacksEpoch<L>]) -> EpochList<L> {
+        EpochList(epochs.to_vec())
+    }
+
+    pub fn get(&self, index: StacksEpochId) -> Option<&StacksEpoch<L>> {
+        self.0.get(StacksEpoch::find_epoch_by_id(&self.0, index)?)
+    }
+
+    pub fn get_mut(&mut self, index: StacksEpochId) -> Option<&mut StacksEpoch<L>> {
+        let index = StacksEpoch::find_epoch_by_id(&self.0, index)?;
+        self.0.get_mut(index)
+    }
+
+    /// Truncates the list after the given epoch id
+    pub fn truncate_after(&mut self, epoch_id: StacksEpochId) {
+        if let Some(index) = StacksEpoch::find_epoch_by_id(&self.0, epoch_id) {
+            self.0.truncate(index + 1);
+        }
+    }
+}
+
+impl<L: Clone> Index<StacksEpochId> for EpochList<L> {
+    type Output = StacksEpoch<L>;
+    fn index(&self, index: StacksEpochId) -> &StacksEpoch<L> {
+        self.get(index)
+            .expect("Invalid StacksEpochId: could not find corresponding epoch")
+    }
+}
+
+impl<L: Clone> IndexMut<StacksEpochId> for EpochList<L> {
+    fn index_mut(&mut self, index: StacksEpochId) -> &mut StacksEpoch<L> {
+        self.get_mut(index)
+            .expect("Invalid StacksEpochId: could not find corresponding epoch")
+    }
+}
+
+impl<L: Clone> Deref for EpochList<L> {
+    type Target = [StacksEpoch<L>];
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+impl<L: Clone> DerefMut for EpochList<L> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.0
+    }
+}

stackslib/src/burnchains/bitcoin/indexer.rs

@@ -45,7 +45,7 @@ use crate::burnchains::{
     Burnchain, BurnchainBlockHeader, Error as burnchain_error, MagicBytes, BLOCKSTACK_MAGIC_MAINNET,
 };
 use crate::core::{
-    StacksEpoch, StacksEpochExtension, STACKS_EPOCHS_MAINNET, STACKS_EPOCHS_REGTEST,
+    EpochList, StacksEpoch, StacksEpochExtension, STACKS_EPOCHS_MAINNET, STACKS_EPOCHS_REGTEST,
     STACKS_EPOCHS_TESTNET,
 };
 use crate::util_lib::db::Error as DBError;
@@ -91,12 +91,12 @@ impl TryFrom<u32> for BitcoinNetworkType {
 /// Get the default epochs definitions for the given BitcoinNetworkType.
 /// Should *not* be used except by the BitcoinIndexer when no epochs vector
 /// was specified.
-pub fn get_bitcoin_stacks_epochs(network_id: BitcoinNetworkType) -> Vec<StacksEpoch> {
-    match network_id {
-        BitcoinNetworkType::Mainnet => STACKS_EPOCHS_MAINNET.to_vec(),
-        BitcoinNetworkType::Testnet => STACKS_EPOCHS_TESTNET.to_vec(),
-        BitcoinNetworkType::Regtest => STACKS_EPOCHS_REGTEST.to_vec(),
-    }
+pub fn get_bitcoin_stacks_epochs(network_id: BitcoinNetworkType) -> EpochList {
+    EpochList::new(match network_id {
+        BitcoinNetworkType::Mainnet => &*STACKS_EPOCHS_MAINNET,
+        BitcoinNetworkType::Testnet => &*STACKS_EPOCHS_TESTNET,
+        BitcoinNetworkType::Regtest => &*STACKS_EPOCHS_REGTEST,
+    })
 }
 
 #[derive(Debug, Clone, PartialEq)]
@@ -112,7 +112,7 @@ pub struct BitcoinIndexerConfig {
     pub spv_headers_path: String,
     pub first_block: u64,
     pub magic_bytes: MagicBytes,
-    pub epochs: Option<Vec<StacksEpoch>>,
+    pub epochs: Option<EpochList>,
 }
 
 #[derive(Debug)]
@@ -1041,7 +1041,7 @@ impl BurnchainIndexer for BitcoinIndexer {
     /// 2) Use hard-coded static values, otherwise.
     ///
     /// It is an error (panic) to set custom epochs if running on `Mainnet`.
-    fn get_stacks_epochs(&self) -> Vec<StacksEpoch> {
+    fn get_stacks_epochs(&self) -> EpochList {
         StacksEpoch::get_epochs(self.runtime.network_id, self.config.epochs.as_ref())
     }
 

stackslib/src/burnchains/burnchain.rs

@@ -29,6 +29,7 @@ use stacks_common::util::hash::to_hex;
 use stacks_common::util::vrf::VRFPublicKey;
 use stacks_common::util::{get_epoch_time_ms, get_epoch_time_secs, log, sleep_ms};
 
+use super::EpochList;
 use crate::burnchains::affirmation::update_pox_affirmation_maps;
 use crate::burnchains::bitcoin::address::{
     to_c32_version_byte, BitcoinAddress, LegacyBitcoinAddressType,
@@ -718,7 +719,7 @@ impl Burnchain {
         readwrite: bool,
         first_block_header_hash: BurnchainHeaderHash,
         first_block_header_timestamp: u64,
-        epochs: Vec<StacksEpoch>,
+        epochs: EpochList,
     ) -> Result<(SortitionDB, BurnchainDB), burnchain_error> {
         Burnchain::setup_chainstate_dirs(&self.working_dir)?;
 

stackslib/src/burnchains/indexer.rs

@@ -62,7 +62,7 @@ pub trait BurnchainIndexer {
     fn get_first_block_height(&self) -> u64;
     fn get_first_block_header_hash(&self) -> Result<BurnchainHeaderHash, burnchain_error>;
     fn get_first_block_header_timestamp(&self) -> Result<u64, burnchain_error>;
-    fn get_stacks_epochs(&self) -> Vec<StacksEpoch>;
+    fn get_stacks_epochs(&self) -> EpochList;
 
     fn get_headers_path(&self) -> String;
     fn get_headers_height(&self) -> Result<u64, burnchain_error>;

stackslib/src/chainstate/burn/db/sortdb.rs

@@ -2969,9 +2969,9 @@ impl SortitionDB {
         db_tx: &Transaction,
         epochs: &[StacksEpoch],
     ) -> Result<(), db_error> {
-        let epochs = StacksEpoch::validate_epochs(epochs);
+        let epochs: &[StacksEpoch] = &StacksEpoch::validate_epochs(epochs);
         let existing_epochs = Self::get_stacks_epochs(db_tx)?;
-        if existing_epochs == epochs {
+        if &existing_epochs == epochs {
             return Ok(());
         }
 
@@ -3482,9 +3482,10 @@ impl SortitionDB {
                         tx.commit()?;
                     } else if version == expected_version {
                         // this transaction is almost never needed
-                        let validated_epochs = StacksEpoch::validate_epochs(epochs);
+                        let validated_epochs: &[StacksEpoch] =
+                            &StacksEpoch::validate_epochs(epochs);
                         let existing_epochs = Self::get_stacks_epochs(self.conn())?;
-                        if existing_epochs == validated_epochs {
+                        if &existing_epochs == validated_epochs {
                             return Ok(());
                         }
 

stackslib/src/cli.rs

@@ -53,7 +53,7 @@ pub struct StacksChainConfig {
     pub first_burn_header_hash: BurnchainHeaderHash,
     pub first_burn_header_timestamp: u64,
     pub pox_constants: PoxConstants,
-    pub epochs: Vec<StacksEpoch>,
+    pub epochs: EpochList,
 }
 
 impl StacksChainConfig {
@@ -65,7 +65,7 @@ impl StacksChainConfig {
                 .unwrap(),
             first_burn_header_timestamp: BITCOIN_MAINNET_FIRST_BLOCK_TIMESTAMP.into(),
             pox_constants: PoxConstants::mainnet_default(),
-            epochs: STACKS_EPOCHS_MAINNET.to_vec(),
+            epochs: EpochList::new(&*STACKS_EPOCHS_MAINNET),
         }
     }
 }