Avoid code duplication

node/core/provisioner/src/disputes/mod.rs

@@ -23,7 +23,8 @@ use polkadot_node_primitives::{CandidateVotes, DisputeStatus};
 use polkadot_node_subsystem::{messages::DisputeCoordinatorMessage, overseer, ActivatedLeaf};
 use polkadot_primitives::v2::{
 	supermajority_threshold, CandidateHash, DisputeState, DisputeStatement, DisputeStatementSet,
-	Hash, MultiDisputeStatementSet, SessionIndex, ValidatorIndex, ValidatorSignature,
+	Hash, InvalidDisputeStatementKind, MultiDisputeStatementSet, SessionIndex,
+	ValidDisputeStatementKind, ValidatorIndex, ValidatorSignature,
 };
 use std::collections::HashMap;
 
@@ -99,8 +100,8 @@ where
 				}
 				let onchain_state =
 					onchain_state.expect("It's guaranteed that the Option contains a value");
-				votes.valid.retain(|v| should_keep_vote(true, v, onchain_state));
-				votes.invalid.retain(|v| should_keep_vote(false, v, onchain_state));
+				votes.valid.retain(|v| should_keep_vote(v, onchain_state));
+				votes.invalid.retain(|v| should_keep_vote(v, onchain_state));
 				(session_index, candidate_hash, votes)
 			})
 			.collect::<Vec<_>>();
@@ -211,12 +212,29 @@ fn partition_recent_disputes(
 	partitioned
 }
 
-fn should_keep_vote<T>(
-	local_vote: bool,
+// Helper trait to obtain the value of vote for `InvalidDisputeStatementKind` and `ValidDisputeStatementKind`.
+// The alternative was to pass a bool to `fn should_keep_vote` explicitly but it's pointless as the value is already 'encoded' in the type.
+trait VoteType {
+	fn vote_value() -> bool;
+}
+
+impl VoteType for InvalidDisputeStatementKind {
+	fn vote_value() -> bool {
+		false
+	}
+}
+
+impl VoteType for ValidDisputeStatementKind {
+	fn vote_value() -> bool {
+		true
+	}
+}
+
+fn should_keep_vote<T: VoteType>(
 	v: &(T, ValidatorIndex, ValidatorSignature),
 	onchain_state: &DisputeState,
 ) -> bool {
-	// TODO: what to do if both bits are set to 1?
+	let local_vote = T::vote_value();
 	let (_, validator_index, _) = v;
 	let in_validators_for = *onchain_state
 		.validators_for
@@ -231,15 +249,14 @@ fn should_keep_vote<T>(
 
 	if in_validators_for && in_validators_against {
 		// The validator has double voted and runtime knows about this. Ignore this vote.
-		return false;
+		return false
 	}
 
-	if local_vote	&& in_validators_against ||
-		 !local_vote && in_validators_for {
-			// local vote differs from the onchain vote
-			// we need this vote to punish the offending validator
-			return true;
-		 }
+	if local_vote && in_validators_against || !local_vote && in_validators_for {
+		// local vote differs from the onchain vote
+		// we need this vote to punish the offending validator
+		return true
+	}
 
 	// The vote is valid. Return true if it is not seen onchain.
 	!(in_validators_for || in_validators_against)

node/core/provisioner/src/disputes/tests.rs

@@ -34,20 +34,20 @@ fn should_keep_vote_behaves() {
 	let local_invalid_unknown =
 		(InvalidDisputeStatementKind::Explicit, ValidatorIndex(3), test_helpers::dummy_signature());
 
-	assert_eq!(should_keep_vote(true, &local_valid_known, &onchain_state), false);
-	assert_eq!(should_keep_vote(true, &local_valid_unknown, &onchain_state), true);
-	assert_eq!(should_keep_vote(false, &local_invalid_known, &onchain_state), false);
-	assert_eq!(should_keep_vote(false, &local_invalid_unknown, &onchain_state), true);
+	assert_eq!(should_keep_vote(&local_valid_known, &onchain_state), false);
+	assert_eq!(should_keep_vote(&local_valid_unknown, &onchain_state), true);
+	assert_eq!(should_keep_vote(&local_invalid_known, &onchain_state), false);
+	assert_eq!(should_keep_vote(&local_invalid_unknown, &onchain_state), true);
 
 	//double voting - onchain knows
 	let local_double_vote_onchain_knows =
 		(InvalidDisputeStatementKind::Explicit, ValidatorIndex(4), test_helpers::dummy_signature());
-	assert_eq!(should_keep_vote(false, &local_double_vote_onchain_knows, &onchain_state), false);
+	assert_eq!(should_keep_vote(&local_double_vote_onchain_knows, &onchain_state), false);
 
 	//double voting - onchain doesn't know
 	let local_double_vote_onchain_doesnt_knows =
 		(InvalidDisputeStatementKind::Explicit, ValidatorIndex(0), test_helpers::dummy_signature());
-	assert_eq!(should_keep_vote(false, &local_double_vote_onchain_doesnt_knows, &onchain_state), true);
+	assert_eq!(should_keep_vote(&local_double_vote_onchain_doesnt_knows, &onchain_state), true);
 
 	// empty onchain state
 	let empty_onchain_state = DisputeState {
@@ -56,7 +56,10 @@ fn should_keep_vote_behaves() {
 		start: 1,
 		concluded_at: None,
 	};
-	assert_eq!(should_keep_vote(false, &local_double_vote_onchain_doesnt_knows, &empty_onchain_state), true);
+	assert_eq!(
+		should_keep_vote(&local_double_vote_onchain_doesnt_knows, &empty_onchain_state),
+		true
+	);
 }
 
 #[test]