Impl Ord on ForkName for ChainSpec usage (#5531)

* Impl Ord on ForkName for ChainSpec usage * add fork ord consistent test
sigp · Apr 10, 2024 · ced6538 · ced6538
1 parent d527d12
commit ced6538
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 33 deletions.
diff --git a/consensus/types/src/chain_spec.rs b/consensus/types/src/chain_spec.rs
@@ -334,12 +334,12 @@ impl ChainSpec {
     }
 
     pub fn inactivity_penalty_quotient_for_fork(&self, fork_name: ForkName) -> u64 {
-        match fork_name {
-            ForkName::Base => self.inactivity_penalty_quotient,
-            ForkName::Altair => self.inactivity_penalty_quotient_altair,
-            ForkName::Merge => self.inactivity_penalty_quotient_bellatrix,
-            ForkName::Capella => self.inactivity_penalty_quotient_bellatrix,
-            ForkName::Deneb | ForkName::Electra => self.inactivity_penalty_quotient_bellatrix,
+        if fork_name >= ForkName::Merge {
+            self.inactivity_penalty_quotient_bellatrix
+        } else if fork_name >= ForkName::Altair {
+            self.inactivity_penalty_quotient_altair
+        } else {
+            self.inactivity_penalty_quotient
         }
     }
 
@@ -348,13 +348,13 @@ impl ChainSpec {
         &self,
         state: &BeaconState<E>,
     ) -> u64 {
-        match state {
-            BeaconState::Base(_) => self.proportional_slashing_multiplier,
-            BeaconState::Altair(_) => self.proportional_slashing_multiplier_altair,
-            BeaconState::Merge(_) => self.proportional_slashing_multiplier_bellatrix,
-            BeaconState::Capella(_) => self.proportional_slashing_multiplier_bellatrix,
-            BeaconState::Deneb(_) => self.proportional_slashing_multiplier_bellatrix,
-            BeaconState::Electra(_) => self.proportional_slashing_multiplier_bellatrix,
+        let fork_name = state.fork_name_unchecked();
+        if fork_name >= ForkName::Merge {
+            self.proportional_slashing_multiplier_bellatrix
+        } else if fork_name >= ForkName::Altair {
+            self.proportional_slashing_multiplier_altair
+        } else {
+            self.proportional_slashing_multiplier
         }
     }
 
@@ -363,13 +363,13 @@ impl ChainSpec {
         &self,
         state: &BeaconState<E>,
     ) -> u64 {
-        match state {
-            BeaconState::Base(_) => self.min_slashing_penalty_quotient,
-            BeaconState::Altair(_) => self.min_slashing_penalty_quotient_altair,
-            BeaconState::Merge(_) => self.min_slashing_penalty_quotient_bellatrix,
-            BeaconState::Capella(_) => self.min_slashing_penalty_quotient_bellatrix,
-            BeaconState::Deneb(_) => self.min_slashing_penalty_quotient_bellatrix,
-            BeaconState::Electra(_) => self.min_slashing_penalty_quotient_bellatrix,
+        let fork_name = state.fork_name_unchecked();
+        if fork_name >= ForkName::Merge {
+            self.min_slashing_penalty_quotient_bellatrix
+        } else if fork_name >= ForkName::Altair {
+            self.min_slashing_penalty_quotient_altair
+        } else {
+            self.min_slashing_penalty_quotient
         }
     }
 
@@ -531,22 +531,19 @@ impl ChainSpec {
     }
 
     pub fn max_blocks_by_root_request(&self, fork_name: ForkName) -> usize {
-        match fork_name {
-            ForkName::Base | ForkName::Altair | ForkName::Merge | ForkName::Capella => {
-                self.max_blocks_by_root_request
-            }
-            ForkName::Deneb | ForkName::Electra => self.max_blocks_by_root_request_deneb,
+        if fork_name >= ForkName::Deneb {
+            self.max_blocks_by_root_request_deneb
+        } else {
+            self.max_blocks_by_root_request
         }
     }
 
     pub fn max_request_blocks(&self, fork_name: ForkName) -> usize {
-        let max_request_blocks = match fork_name {
-            ForkName::Base | ForkName::Altair | ForkName::Merge | ForkName::Capella => {
-                self.max_request_blocks
-            }
-            ForkName::Deneb | ForkName::Electra => self.max_request_blocks_deneb,
-        };
-        max_request_blocks as usize
+        if fork_name >= ForkName::Deneb {
+            self.max_request_blocks_deneb as usize
+        } else {
+            self.max_request_blocks as usize
+        }
     }
 
     /// Returns a `ChainSpec` compatible with the Ethereum Foundation specification.

diff --git a/consensus/types/src/fork_name.rs b/consensus/types/src/fork_name.rs
@@ -4,7 +4,9 @@ use ssz_derive::{Decode, Encode};
 use std::fmt::{self, Display, Formatter};
 use std::str::FromStr;
 
-#[derive(Debug, Clone, Copy, Decode, Encode, PartialEq, Eq, Hash, Serialize, Deserialize)]
+#[derive(
+    Debug, Clone, Copy, Decode, Encode, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize,
+)]
 #[serde(try_from = "String")]
 #[serde(into = "String")]
 #[ssz(enum_behaviour = "tag")]
@@ -272,4 +274,13 @@ mod test {
         }
         assert_eq!(ForkName::latest(), fork);
     }
+
+    #[test]
+    fn fork_ord_consistent() {
+        for (prev_fork, fork) in ForkName::list_all().into_iter().tuple_windows() {
+            assert_eq!(prev_fork.next_fork(), Some(fork));
+            assert_eq!(fork.previous_fork(), Some(prev_fork));
+            assert!(prev_fork < fork);
+        }
+    }
 }