Reduce ForkName boilerplate in fork-context (#6933)

Noted that there's a bit of fork boiler plate in fork context. If we list a mapping of ForkName -> fork_version in the ForkName enum we can get rid of it :) Not much, but should make the next fork a tiny tit less annoying
sigp · Feb 7, 2025 · 59afe41 · 59afe41
1 parent 9c45a0e
commit 59afe41
Show file tree

Hide file tree

Showing 2 changed files with 47 additions and 63 deletions.
diff --git a/consensus/types/src/fork_context.rs b/consensus/types/src/fork_context.rs
@@ -22,61 +22,22 @@ impl ForkContext {
         genesis_validators_root: Hash256,
         spec: &ChainSpec,
     ) -> Self {
-        let mut fork_to_digest = vec![(
-            ForkName::Base,
-            ChainSpec::compute_fork_digest(spec.genesis_fork_version, genesis_validators_root),
-        )];
-
-        // Only add Altair to list of forks if it's enabled
-        // Note: `altair_fork_epoch == None` implies altair hasn't been activated yet on the config.
-        if spec.altair_fork_epoch.is_some() {
-            fork_to_digest.push((
-                ForkName::Altair,
-                ChainSpec::compute_fork_digest(spec.altair_fork_version, genesis_validators_root),
-            ));
-        }
-
-        // Only add Bellatrix to list of forks if it's enabled
-        // Note: `bellatrix_fork_epoch == None` implies bellatrix hasn't been activated yet on the config.
-        if spec.bellatrix_fork_epoch.is_some() {
-            fork_to_digest.push((
-                ForkName::Bellatrix,
-                ChainSpec::compute_fork_digest(
-                    spec.bellatrix_fork_version,
-                    genesis_validators_root,
-                ),
-            ));
-        }
-
-        if spec.capella_fork_epoch.is_some() {
-            fork_to_digest.push((
-                ForkName::Capella,
-                ChainSpec::compute_fork_digest(spec.capella_fork_version, genesis_validators_root),
-            ));
-        }
-
-        if spec.deneb_fork_epoch.is_some() {
-            fork_to_digest.push((
-                ForkName::Deneb,
-                ChainSpec::compute_fork_digest(spec.deneb_fork_version, genesis_validators_root),
-            ));
-        }
-
-        if spec.electra_fork_epoch.is_some() {
-            fork_to_digest.push((
-                ForkName::Electra,
-                ChainSpec::compute_fork_digest(spec.electra_fork_version, genesis_validators_root),
-            ));
-        }
-
-        if spec.fulu_fork_epoch.is_some() {
-            fork_to_digest.push((
-                ForkName::Fulu,
-                ChainSpec::compute_fork_digest(spec.fulu_fork_version, genesis_validators_root),
-            ));
-        }
-
-        let fork_to_digest: HashMap<ForkName, [u8; 4]> = fork_to_digest.into_iter().collect();
+        let fork_to_digest: HashMap<ForkName, [u8; 4]> = ForkName::list_all()
+            .into_iter()
+            .filter_map(|fork| {
+                if fork.fork_epoch(spec).is_some() {
+                    Some((
+                        fork,
+                        ChainSpec::compute_fork_digest(
+                            ForkName::fork_version(fork, spec),
+                            genesis_validators_root,
+                        ),
+                    ))
+                } else {
+                    None
+                }
+            })
+            .collect();
 
         let digest_to_fork = fork_to_digest
             .clone()

diff --git a/consensus/types/src/fork_name.rs b/consensus/types/src/fork_name.rs
@@ -34,14 +34,37 @@ impl ForkName {
     }
 
     pub fn list_all_fork_epochs(spec: &ChainSpec) -> Vec<(ForkName, Option<Epoch>)> {
-        vec![
-            (ForkName::Altair, spec.altair_fork_epoch),
-            (ForkName::Bellatrix, spec.bellatrix_fork_epoch),
-            (ForkName::Capella, spec.capella_fork_epoch),
-            (ForkName::Deneb, spec.deneb_fork_epoch),
-            (ForkName::Electra, spec.electra_fork_epoch),
-            (ForkName::Fulu, spec.fulu_fork_epoch),
-        ]
+        ForkName::list_all()
+            .into_iter()
+            // Skip Base
+            .skip(1)
+            .map(|fork| (fork, fork.fork_epoch(spec)))
+            .collect()
+    }
+
+    pub fn fork_epoch(self, spec: &ChainSpec) -> Option<Epoch> {
+        match self {
+            Self::Base => Some(Epoch::new(0)),
+            Self::Altair => spec.altair_fork_epoch,
+            Self::Bellatrix => spec.bellatrix_fork_epoch,
+            Self::Capella => spec.capella_fork_epoch,
+            Self::Deneb => spec.deneb_fork_epoch,
+            Self::Electra => spec.electra_fork_epoch,
+            Self::Fulu => spec.fulu_fork_epoch,
+        }
+    }
+
+    /// Returns the fork version of a fork
+    pub fn fork_version(self, spec: &ChainSpec) -> [u8; 4] {
+        match self {
+            Self::Base => spec.genesis_fork_version,
+            Self::Altair => spec.altair_fork_version,
+            Self::Bellatrix => spec.bellatrix_fork_version,
+            Self::Capella => spec.capella_fork_version,
+            Self::Deneb => spec.deneb_fork_version,
+            Self::Electra => spec.electra_fork_version,
+            Self::Fulu => spec.fulu_fork_version,
+        }
     }
 
     pub fn latest() -> ForkName {