diff --git a/zebra-state/src/lib.rs b/zebra-state/src/lib.rs index 76d96a91..a470d25a 100644 --- a/zebra-state/src/lib.rs +++ b/zebra-state/src/lib.rs @@ -156,18 +156,15 @@ pub enum Response { /// Get the heights of the blocks for constructing a block_locator list fn block_locator_heights(tip_height: BlockHeight) -> impl Iterator { // Stop at the reorg limit, or the genesis block. - let min_locator_height = tip_height - .0 - .checked_sub(MAX_BLOCK_REORG_HEIGHT.0) - .map(BlockHeight) - .unwrap_or(BlockHeight(0)); + let min_locator_height = tip_height.0.saturating_sub(MAX_BLOCK_REORG_HEIGHT.0); let locators = iter::successors(Some(1u32), |h| h.checked_mul(2)) - .flat_map(move |step| tip_height.0.checked_sub(step)) - .filter(move |&height| height > min_locator_height.0) - .map(BlockHeight); - let locators = iter::once(tip_height) + .flat_map(move |step| tip_height.0.checked_sub(step)); + let locators = iter::once(tip_height.0) .chain(locators) - .chain(iter::once(min_locator_height)); + .take_while(move |&height| height > min_locator_height) + .chain(iter::once(min_locator_height)) + .map(BlockHeight); + let locators: Vec<_> = locators.collect(); tracing::info!( ?tip_height, @@ -175,6 +172,7 @@ fn block_locator_heights(tip_height: BlockHeight) -> impl Iterator assert_eq!(path.file_name(), Some(OsStr::new("testnet"))), } } + + /// Block heights, and the expected minimum block locator height + static BLOCK_LOCATOR_CASES: &[(u32, u32)] = &[ + (0, 0), + (1, 0), + (10, 0), + (98, 0), + (99, 0), + (100, 1), + (101, 2), + (1000, 901), + (10000, 9901), + ]; + + /// Check that the block locator heights are sensible. + #[test] + fn test_block_locator_heights() { + for (height, min_height) in BLOCK_LOCATOR_CASES.iter().cloned() { + let locator = block_locator_heights(BlockHeight(height)).collect::>(); + + assert!(!locator.is_empty(), "locators must not be empty"); + if (height - min_height) > 1 { + assert!( + locator.len() > 2, + "non-trivial locators must have some intermediate heights" + ); + } + + assert_eq!( + locator[0], + BlockHeight(height), + "locators must start with the tip height" + ); + + // Check that the locator is sorted, and that it has no duplicates + // TODO: replace with dedup() and is_sorted_by() when sorting stabilises. + assert!(locator.windows(2).all(|v| match v { + [a, b] => a.0 > b.0, + _ => unreachable!("windows returns exact sized slices"), + })); + + let final_height = locator[locator.len() - 1]; + assert_eq!( + final_height, + BlockHeight(min_height), + "locators must end with the specified final height" + ); + assert!(height - final_height.0 <= MAX_BLOCK_REORG_HEIGHT.0, + format!("locator for {} must not be more than the maximum reorg height {} below the tip, but {} is {} blocks below the tip", + height, + MAX_BLOCK_REORG_HEIGHT.0, + final_height.0, + height - final_height.0)); + } + } }