From 9eafad467c115dd3164eb63d48807f3431b30ad8 Mon Sep 17 00:00:00 2001
From: Ashwin Sekar <ashwin@solana.com>
Date: Fri, 21 Oct 2022 13:17:40 -0700
Subject: [PATCH] Add convenience methods to VoteInstruction to distinguish
 vote types (#28526)

* Add convenience methods to VoteInstruction to distinguish vote types

* use matches! macro instead
---
 core/src/latest_unprocessed_votes.rs |  9 +++----
 programs/vote/src/vote_state/mod.rs  |  2 +-
 runtime/src/vote_parser.rs           | 10 +------
 sdk/program/src/vote/instruction.rs  | 40 ++++++++++++++++++++++++++++
 sdk/program/src/vote/state/mod.rs    |  4 +++
 5 files changed, 50 insertions(+), 15 deletions(-)

diff --git a/core/src/latest_unprocessed_votes.rs b/core/src/latest_unprocessed_votes.rs
index 2e4fe33bc..4fe4096c3 100644
--- a/core/src/latest_unprocessed_votes.rs
+++ b/core/src/latest_unprocessed_votes.rs
@@ -53,16 +53,15 @@ impl LatestValidatorVotePacket {
             .ok_or(DeserializedPacketError::VoteTransactionError)?;
 
         match limited_deserialize::<VoteInstruction>(&instruction.data) {
-            Ok(VoteInstruction::UpdateVoteState(vote_state_update))
-            | Ok(VoteInstruction::UpdateVoteStateSwitch(vote_state_update, _))
-            | Ok(VoteInstruction::CompactUpdateVoteState(vote_state_update))
-            | Ok(VoteInstruction::CompactUpdateVoteStateSwitch(vote_state_update, _)) => {
+            Ok(vote_state_update_instruction)
+                if vote_state_update_instruction.is_single_vote_state_update() =>
+            {
                 let &pubkey = message
                     .message
                     .static_account_keys()
                     .get(0)
                     .ok_or(DeserializedPacketError::VoteTransactionError)?;
-                let slot = vote_state_update.last_voted_slot().unwrap_or(0);
+                let slot = vote_state_update_instruction.last_voted_slot().unwrap_or(0);
 
                 Ok(Self {
                     vote: Some(vote),
diff --git a/programs/vote/src/vote_state/mod.rs b/programs/vote/src/vote_state/mod.rs
index c939cd591..d3b02b4f7 100644
--- a/programs/vote/src/vote_state/mod.rs
+++ b/programs/vote/src/vote_state/mod.rs
@@ -105,7 +105,7 @@ impl VoteTransaction {
 
     pub fn last_voted_slot(&self) -> Option<Slot> {
         match self {
-            VoteTransaction::Vote(vote) => vote.slots.last().copied(),
+            VoteTransaction::Vote(vote) => vote.last_voted_slot(),
             VoteTransaction::VoteStateUpdate(vote_state_update)
             | VoteTransaction::CompactVoteStateUpdate(vote_state_update) => {
                 vote_state_update.last_voted_slot()
diff --git a/runtime/src/vote_parser.rs b/runtime/src/vote_parser.rs
index 9505345ed..814fb47a5 100644
--- a/runtime/src/vote_parser.rs
+++ b/runtime/src/vote_parser.rs
@@ -23,15 +23,7 @@ pub(crate) fn is_simple_vote_transaction(transaction: &SanitizedTransaction) ->
         if program_pubkey == &solana_vote_program::id() {
             if let Ok(vote_instruction) = limited_deserialize::<VoteInstruction>(&instruction.data)
             {
-                return matches!(
-                    vote_instruction,
-                    VoteInstruction::Vote(_)
-                        | VoteInstruction::VoteSwitch(_, _)
-                        | VoteInstruction::UpdateVoteState(_)
-                        | VoteInstruction::UpdateVoteStateSwitch(_, _)
-                        | VoteInstruction::CompactUpdateVoteState(_)
-                        | VoteInstruction::CompactUpdateVoteStateSwitch(..)
-                );
+                return vote_instruction.is_simple_vote();
             }
         }
     }
diff --git a/sdk/program/src/vote/instruction.rs b/sdk/program/src/vote/instruction.rs
index 8cd6b1d6e..f82b05949 100644
--- a/sdk/program/src/vote/instruction.rs
+++ b/sdk/program/src/vote/instruction.rs
@@ -2,6 +2,7 @@
 
 use {
     crate::{
+        clock::Slot,
         hash::Hash,
         instruction::{AccountMeta, Instruction},
         pubkey::Pubkey,
@@ -147,6 +148,45 @@ pub enum VoteInstruction {
     ),
 }
 
+impl VoteInstruction {
+    pub fn is_simple_vote(&self) -> bool {
+        matches!(
+            self,
+            Self::Vote(_)
+                | Self::VoteSwitch(_, _)
+                | Self::UpdateVoteState(_)
+                | Self::UpdateVoteStateSwitch(_, _)
+                | Self::CompactUpdateVoteState(_)
+                | Self::CompactUpdateVoteStateSwitch(_, _),
+        )
+    }
+
+    pub fn is_single_vote_state_update(&self) -> bool {
+        matches!(
+            self,
+            Self::UpdateVoteState(_)
+                | Self::UpdateVoteStateSwitch(_, _)
+                | Self::CompactUpdateVoteState(_)
+                | Self::CompactUpdateVoteStateSwitch(_, _),
+        )
+    }
+
+    /// Only to be used on vote instructions (guard with is_simple_vote),  panics otherwise
+    pub fn last_voted_slot(&self) -> Option<Slot> {
+        assert!(self.is_simple_vote());
+        match self {
+            Self::Vote(v) | Self::VoteSwitch(v, _) => v.last_voted_slot(),
+            Self::UpdateVoteState(vote_state_update)
+            | Self::UpdateVoteStateSwitch(vote_state_update, _)
+            | Self::CompactUpdateVoteState(vote_state_update)
+            | Self::CompactUpdateVoteStateSwitch(vote_state_update, _) => {
+                vote_state_update.last_voted_slot()
+            }
+            _ => panic!("Tried to get slot on non simple vote instruction"),
+        }
+    }
+}
+
 fn initialize_account(vote_pubkey: &Pubkey, vote_init: &VoteInit) -> Instruction {
     let account_metas = vec![
         AccountMeta::new(*vote_pubkey, false),
diff --git a/sdk/program/src/vote/state/mod.rs b/sdk/program/src/vote/state/mod.rs
index 49114ea39..b5279d00e 100644
--- a/sdk/program/src/vote/state/mod.rs
+++ b/sdk/program/src/vote/state/mod.rs
@@ -50,6 +50,10 @@ impl Vote {
             timestamp: None,
         }
     }
+
+    pub fn last_voted_slot(&self) -> Option<Slot> {
+        self.slots.last().copied()
+    }
 }
 
 #[derive(Serialize, Default, Deserialize, Debug, PartialEq, Eq, Copy, Clone, AbiExample)]