Approval checking unit tests

node/core/approval-voting/src/approval_checking.rs

@@ -17,10 +17,11 @@
 //! Utilities for checking whether a candidate has been approved under a given block.
 
 use polkadot_node_primitives::approval::DelayTranche;
+use polkadot_primitives::v1::ValidatorIndex;
 use bitvec::slice::BitSlice;
 use bitvec::order::Lsb0 as BitOrderLsb0;
 
-use crate::persisted_entries::{ApprovalEntry, CandidateEntry};
+use crate::persisted_entries::{TrancheEntry, ApprovalEntry, CandidateEntry};
 use crate::time::Tick;
 
 /// The required tranches of assignments needed to determine whether a candidate is approved.
@@ -263,6 +264,39 @@ impl State {
 	}
 }
 
+/// Constructs an infinite iterator from an array of `TrancheEntry` values. Any missing tranches
+/// are filled with empty assignments, as they are needed to compute the approved tranches.
+fn filled_tranche_iterator<'a>(
+	tranches: &'a [TrancheEntry],
+) -> impl Iterator<Item = (DelayTranche, &[(ValidatorIndex, Tick)])> {
+	let mut gap_end = None;
+
+	let approval_entries_filled = tranches
+		.iter()
+		.flat_map(move |tranche_entry| {
+			let tranche = tranche_entry.tranche();
+			let assignments = tranche_entry.assignments();
+
+			// The new gap_start immediately follows the prior gap_end, if one exists.
+			// Otherwise, on the first pass, the new gap_start is set to the first
+			// tranche so that the range below will be empty.
+			let gap_start = gap_end.map(|end| end + 1).unwrap_or(tranche);
+			gap_end = Some(tranche);
+
+			(gap_start..tranche).map(|i| (i, &[] as &[_]))
+				.chain(std::iter::once((tranche, assignments)))
+		});
+
+	let pre_end = tranches.first().map(|t| t.tranche());
+	let post_start = tranches.last().map_or(0, |t| t.tranche() + 1);
+
+	let pre = pre_end.into_iter()
+		.flat_map(|pre_end| (0..pre_end).map(|i| (i, &[] as &[_])));
+	let post = (post_start..).map(|i| (i, &[] as &[_]));
+
+	pre.chain(approval_entries_filled).chain(post)
+}
+
 /// Determine the amount of tranches of assignments needed to determine approval of a candidate.
 pub fn tranches_to_approve(
 	approval_entry: &ApprovalEntry,
@@ -288,31 +322,7 @@ pub fn tranches_to_approve(
 	// these empty tranches, so we create an iterator to fill the gaps.
 	//
 	// This iterator has an infinitely long amount of non-empty tranches appended to the end.
-	let tranches_with_gaps_filled = {
-		let mut gap_end = 0;
-
-		let approval_entries_filled = approval_entry.tranches()
-			.iter()
-			.flat_map(move |tranche_entry| {
-				let tranche = tranche_entry.tranche();
-				let assignments = tranche_entry.assignments();
-
-				let gap_start = gap_end + 1;
-				gap_end = tranche;
-
-				(gap_start..tranche).map(|i| (i, &[] as &[_]))
-					.chain(std::iter::once((tranche, assignments)))
-			});
-
-		let pre_end = approval_entry.tranches().first().map(|t| t.tranche());
-		let post_start = approval_entry.tranches().last().map_or(0, |t| t.tranche() + 1);
-
-		let pre = pre_end.into_iter()
-			.flat_map(|pre_end| (0..pre_end).map(|i| (i, &[] as &[_])));
-		let post = (post_start..).map(|i| (i, &[] as &[_]));
-
-		pre.chain(approval_entries_filled).chain(post)
-	};
+	let tranches_with_gaps_filled = filled_tranche_iterator(approval_entry.tranches());
 
 	tranches_with_gaps_filled
 		.scan(Some(initial_state), |state, (tranche, assignments)| {
@@ -384,7 +394,7 @@ pub fn tranches_to_approve(
 mod tests {
 	use super::*;
 
-	use polkadot_primitives::v1::{GroupIndex, ValidatorIndex};
+	use polkadot_primitives::v1::GroupIndex;
 	use bitvec::bitvec;
 	use bitvec::order::Lsb0 as BitOrderLsb0;
 
@@ -938,6 +948,51 @@ mod tests {
 			},
 		);
 	}
+
+	#[test]
+	fn filled_tranche_iterator_yields_sequential_tranches() {
+		const PREFIX: u32 = 10;
+
+		let test_tranches = vec![
+			vec![],					// empty set
+			vec![0],				// zero start
+			vec![0, 3],				// zero start with gap
+			vec![2],				// non-zero start
+			vec![2, 4],				// non-zero start with gap
+			vec![0, 1, 2],			// zero start with run and no gap
+			vec![2, 3, 4, 8],		// non-zero start with run and gap
+			vec![0, 1, 2, 5, 6, 7],	// zero start with runs and gap
+		];
+
+		for test_tranche in test_tranches {
+			let mut approval_entry: ApprovalEntry = approval_db::v1::ApprovalEntry {
+				tranches: Vec::new(),
+				backing_group: GroupIndex(0),
+				our_assignment: None,
+				our_approval_sig: None,
+				assignments: bitvec![BitOrderLsb0, u8; 0; 3],
+				approved: false,
+			}.into();
+
+			// Populate the requested tranches. The assignemnts aren't inspected in
+			// this test.
+			for &t in &test_tranche {
+				approval_entry.import_assignment(t, ValidatorIndex(0), 0)
+			}
+
+			let filled_tranches = filled_tranche_iterator(approval_entry.tranches());
+
+			// Take the first PREFIX entries and map them to their tranche.
+			let tranches: Vec<DelayTranche> = filled_tranches
+				.take(PREFIX as usize)
+				.map(|e| e.0)
+				.collect();
+
+			// We expect this sequence to be sequential.
+			let exp_tranches: Vec<DelayTranche> = (0..PREFIX).collect();
+			assert_eq!(tranches, exp_tranches, "for test tranches: {:?}", test_tranche);
+		}
+	}
 }
 
 #[test]