added is_double_vote and is_surround_vote (lighthouse-150)

eth2/types/src/slashable_vote_data.rs

@@ -1,4 +1,5 @@
 use super::AttestationData;
+use crate::spec::ChainSpec;
 use crate::test_utils::TestRandom;
 use bls::AggregateSignature;
 use rand::RngCore;
@@ -13,6 +14,21 @@ pub struct SlashableVoteData {
     pub aggregate_signature: AggregateSignature,
 }
 
+impl SlashableVoteData {
+    pub fn is_double_vote(&self, other: &SlashableVoteData, spec: &ChainSpec) -> bool {
+        self.data.slot.epoch(spec.epoch_length) == other.data.slot.epoch(spec.epoch_length)
+    }
+
+    pub fn is_surround_vote(&self, other: &SlashableVoteData, spec: &ChainSpec) -> bool {
+        let source_epoch_1 = self.data.justified_epoch;
+        let source_epoch_2 = other.data.justified_epoch;
+        let target_epoch_1 = self.data.slot.epoch(spec.epoch_length);
+        let target_epoch_2 = other.data.slot.epoch(spec.epoch_length);
+
+        (source_epoch_1 < source_epoch_2) && (target_epoch_2 < target_epoch_1)
+    }
+}
+
 impl Encodable for SlashableVoteData {
     fn ssz_append(&self, s: &mut SszStream) {
         s.append_vec(&self.custody_bit_0_indices);
@@ -66,9 +82,71 @@ impl<T: RngCore> TestRandom<T> for SlashableVoteData {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::slot_epoch_height::{Epoch, Slot};
+    use crate::spec::ChainSpec;
     use crate::test_utils::{SeedableRng, TestRandom, XorShiftRng};
     use ssz::ssz_encode;
 
+    #[test]
+    pub fn test_is_double_vote_true() {
+        let spec = ChainSpec::foundation();
+        let slashable_vote_first = create_slashable_vote_data(1, 1, &spec);
+        let slashable_vote_second = create_slashable_vote_data(1, 1, &spec);
+
+        assert_eq!(
+            slashable_vote_first.is_double_vote(&slashable_vote_second, &spec),
+            true
+        )
+    }
+
+    #[test]
+    pub fn test_is_double_vote_false() {
+        let spec = ChainSpec::foundation();
+        let slashable_vote_first = create_slashable_vote_data(1, 1, &spec);
+        let slashable_vote_second = create_slashable_vote_data(2, 1, &spec);
+
+        assert_eq!(
+            slashable_vote_first.is_double_vote(&slashable_vote_second, &spec),
+            false
+        );
+    }
+
+    #[test]
+    pub fn test_is_surround_vote_true() {
+        let spec = ChainSpec::foundation();
+        let slashable_vote_first = create_slashable_vote_data(2, 1, &spec);
+        let slashable_vote_second = create_slashable_vote_data(1, 2, &spec);
+
+        assert_eq!(
+            slashable_vote_first.is_surround_vote(&slashable_vote_second, &spec),
+            true
+        );
+    }
+
+    #[test]
+    pub fn test_is_surround_vote_false_source_epoch_fails() {
+        let spec = ChainSpec::foundation();
+        let slashable_vote_first = create_slashable_vote_data(2, 2, &spec);
+        let slashable_vote_second = create_slashable_vote_data(1, 1, &spec);
+
+        assert_eq!(
+            slashable_vote_first.is_surround_vote(&slashable_vote_second, &spec),
+            false
+        );
+    }
+
+    #[test]
+    pub fn test_is_surround_vote_false_target_epoch_fails() {
+        let spec = ChainSpec::foundation();
+        let slashable_vote_first = create_slashable_vote_data(1, 1, &spec);
+        let slashable_vote_second = create_slashable_vote_data(2, 2, &spec);
+
+        assert_eq!(
+            slashable_vote_first.is_surround_vote(&slashable_vote_second, &spec),
+            false
+        );
+    }
+
     #[test]
     pub fn test_ssz_round_trip() {
         let mut rng = XorShiftRng::from_seed([42; 16]);
@@ -91,4 +169,17 @@ mod tests {
         // TODO: Add further tests
         // https://github.com/sigp/lighthouse/issues/170
     }
+
+    fn create_slashable_vote_data(
+        slot_factor: u64,
+        justified_epoch: u64,
+        spec: &ChainSpec,
+    ) -> SlashableVoteData {
+        let mut rng = XorShiftRng::from_seed([42; 16]);
+        let mut slashable_vote = SlashableVoteData::random_for_test(&mut rng);
+
+        slashable_vote.data.slot = Slot::new(slot_factor * spec.epoch_length);
+        slashable_vote.data.justified_epoch = Epoch::new(justified_epoch);
+        slashable_vote
+    }
 }