From 46e895ee5ab579df05d5a1a689504e7b0886ee60 Mon Sep 17 00:00:00 2001
From: RunasSudo <runassudo@yingtongli.me>
Date: Mon, 9 Aug 2021 00:17:14 +1000
Subject: [PATCH] Correct handling of exhausted votes during random sample
 surplus distribution

---
 src/stv/sample.rs | 33 ++++++++++++++++++---------------
 1 file changed, 18 insertions(+), 15 deletions(-)

diff --git a/src/stv/sample.rs b/src/stv/sample.rs
index c831639..56217cf 100644
--- a/src/stv/sample.rs
+++ b/src/stv/sample.rs
@@ -181,8 +181,6 @@ where
 			count_card.votes.assign(state.quota.as_ref().unwrap());
 			checksum -= surplus;
 			
-			count_card.finalised = true; // Mark surpluses as done
-			
 			// Update loss by fraction
 			state.loss_fraction.transfer(&-checksum);
 		}
@@ -209,6 +207,7 @@ where
 						let surplus = &state.candidates[elected_candidate].votes - state.quota.as_ref().unwrap();
 						state.exhausted.transfer(&surplus);
 						state.candidates.get_mut(elected_candidate).unwrap().transfer(&-surplus);
+						break;
 					}
 				}
 			}
@@ -236,31 +235,35 @@ where
 			
 			while &state.candidates[elected_candidate].votes > state.quota.as_ref().unwrap() {
 				// Transfer one vote to next available preference
-				let vote = numbered_votes.remove(&index).unwrap();
-				transfer_ballot(state, opts, elected_candidate, vote, opts.transferable_only)?;
-				if state.num_elected == state.election.seats {
-					return Ok(());
-				}
-				
-				index += skip_value;
-				if index >= total_ballots {
-					iteration += 1;
-					index = iteration + skip_value - 1;
-					
-					if iteration >= skip_value {
+				match numbered_votes.remove(&index) {
+					Some(vote) => {
+						transfer_ballot(state, opts, elected_candidate, vote, opts.transferable_only)?;
+						if state.num_elected == state.election.seats {
+							return Ok(());
+						}
+					}
+					None => {
 						// We have run out of ballot papers
 						// Remaining ballot papers exhaust
-						
 						let surplus = &state.candidates[elected_candidate].votes - state.quota.as_ref().unwrap();
 						state.exhausted.transfer(&surplus);
 						state.candidates.get_mut(elected_candidate).unwrap().transfer(&-surplus);
 						break;
 					}
 				}
+				
+				index += skip_value;
+				if index >= total_ballots {
+					iteration += 1;
+					index = iteration + skip_value - 1;
+				}
 			}
 		}
 	}
 	
+	let count_card = state.candidates.get_mut(elected_candidate).unwrap();
+	count_card.finalised = true; // Mark surpluses as done
+	
 	return Ok(());
 }