Improve readability of vote lockout processing (#16987)

* Improve readability of vote lockout processing

* clippy

* simplify comment

* feedback

core/src/consensus.rs

@@ -264,7 +264,7 @@ impl Tower {
             };
             for vote in &vote_state.votes {
                 lockout_intervals
-                    .entry(vote.expiration_slot())
+                    .entry(vote.last_locked_out_slot())
                     .or_insert_with(Vec::new)
                     .push((vote.slot, key));
             }
@@ -534,20 +534,22 @@ impl Tower {
             return true;
         }
 
+        // Check if a slot is locked out by simulating adding a vote for that
+        // slot to the current lockouts to pop any expired votes. If any of the
+        // remaining voted slots are on a different fork from the checked slot,
+        // it's still locked out.
         let mut lockouts = self.lockouts.clone();
         lockouts.process_slot_vote_unchecked(slot);
         for vote in &lockouts.votes {
-            if vote.slot == slot {
+            if slot != vote.slot && !ancestors[&slot].contains(&vote.slot) {
-                continue;
-            }
-            if !ancestors[&slot].contains(&vote.slot) {
                 return true;
             }
         }
+
         if let Some(root_slot) = lockouts.root_slot {
+            if slot != root_slot {
                 // This case should never happen because bank forks purges all
                 // non-descendants of the root every time root is set
-            if slot != root_slot {
                 assert!(
                     ancestors[&slot].contains(&root_slot),
                     "ancestors: {:?}, slot: {} root: {}",
@@ -728,8 +730,9 @@ impl Tower {
                         .unwrap()
                         .fork_stats
                         .lockout_intervals;
-                    // Find any locked out intervals in this bank with endpoint >= last_vote,
+                    // Find any locked out intervals for vote accounts in this bank with
-                    // implies they are locked out at last_vote
+                    // `lockout_interval_end` >= `last_vote`, which implies they are locked out at
+                    // `last_vote` on another fork.
                     for (_lockout_interval_end, intervals_keyed_by_end) in lockout_intervals.range((Included(last_voted_slot), Unbounded)) {
                         for (lockout_interval_start, vote_account_pubkey) in intervals_keyed_by_end {
                             if locked_out_vote_accounts.contains(vote_account_pubkey) {

programs/vote/src/vote_state/mod.rs

@@ -84,13 +84,15 @@ impl Lockout {
         (INITIAL_LOCKOUT as u64).pow(self.confirmation_count)
     }
 
-    // The slot height at which this vote expires (cannot vote for any slot
+    // The last slot at which a vote is still locked out. Validators should not
-    // less than this)
+    // vote on a slot in another fork which is less than or equal to this slot
-    pub fn expiration_slot(&self) -> Slot {
+    // to avoid having their stake slashed.
+    pub fn last_locked_out_slot(&self) -> Slot {
         self.slot + self.lockout()
     }
-    pub fn is_expired(&self, slot: Slot) -> bool {
+
-        self.expiration_slot() < slot
+    pub fn is_locked_out_at_slot(&self, slot: Slot) -> bool {
+        self.last_locked_out_slot() >= slot
     }
 }
 
@@ -354,22 +356,24 @@ impl VoteState {
         }
         self.check_slots_are_valid(vote, slot_hashes)?;
 
-        vote.slots.iter().for_each(|s| self.process_slot(*s, epoch));
+        vote.slots
+            .iter()
+            .for_each(|s| self.process_next_vote_slot(*s, epoch));
         Ok(())
     }
 
-    pub fn process_slot(&mut self, slot: Slot, epoch: Epoch) {
+    pub fn process_next_vote_slot(&mut self, next_vote_slot: Slot, epoch: Epoch) {
         // Ignore votes for slots earlier than we already have votes for
         if self
             .last_voted_slot()
-            .map_or(false, |last_voted_slot| slot <= last_voted_slot)
+            .map_or(false, |last_voted_slot| next_vote_slot <= last_voted_slot)
         {
             return;
         }
 
-        let vote = Lockout::new(slot);
+        let vote = Lockout::new(next_vote_slot);
 
-        self.pop_expired_votes(slot);
+        self.pop_expired_votes(next_vote_slot);
 
         // Once the stack is full, pop the oldest lockout and distribute rewards
         if self.votes.len() == MAX_LOCKOUT_HISTORY {
@@ -540,9 +544,13 @@ impl VoteState {
         Ok(pubkey)
     }
 
-    fn pop_expired_votes(&mut self, slot: Slot) {
+    // Pop all recent votes that are not locked out at the next vote slot.  This
-        loop {
+    // allows validators to switch forks once their votes for another fork have
-            if self.last_lockout().map_or(false, |v| v.is_expired(slot)) {
+    // expired. This also allows validators continue voting on recent blocks in
+    // the same fork without increasing lockouts.
+    fn pop_expired_votes(&mut self, next_vote_slot: Slot) {
+        while let Some(vote) = self.last_lockout() {
+            if !vote.is_locked_out_at_slot(next_vote_slot) {
                 self.votes.pop_back();
             } else {
                 break;
@@ -1350,11 +1358,12 @@ mod tests {
         // second vote
         let top_vote = vote_state.votes.front().unwrap().slot;
         vote_state
-            .process_slot_vote_unchecked(vote_state.last_lockout().unwrap().expiration_slot());
+            .process_slot_vote_unchecked(vote_state.last_lockout().unwrap().last_locked_out_slot());
         assert_eq!(Some(top_vote), vote_state.root_slot);
 
         // Expire everything except the first vote
-        vote_state.process_slot_vote_unchecked(vote_state.votes.front().unwrap().expiration_slot());
+        vote_state
+            .process_slot_vote_unchecked(vote_state.votes.front().unwrap().last_locked_out_slot());
         // First vote and new vote are both stored for a total of 2 votes
         assert_eq!(vote_state.votes.len(), 2);
     }