Use peeking_take_while in BlockReplayer (sigp#4803)

## Issue Addressed

While reviewing sigp#4801 I noticed that our use of `take_while` in the block replayer means that if a state root iterator _with gaps_ is provided, some additonal state roots will be dropped unnecessarily. In practice the impact is small, because once there's _one_ state root miss, the whole tree hash cache needs to be built anyway, and subsequent misses are less costly. However this was still a little inefficient, so I figured it's better to fix it.

## Proposed Changes

Use [`peeking_take_while`](https://docs.rs/itertools/latest/itertools/trait.Itertools.html#method.peeking_take_while) to avoid consuming the next element when checking whether it satisfies the slot predicate.

## Additional Info

There's a gist here that shows the basic dynamics in isolation: https://play.rust-lang.org/?version=stable&mode=debug&edition=2021&gist=40b623cc0febf9ed51705d476ab140c5. Changing the `peeking_take_while` to a `take_while` causes the assert to fail. Similarly I've added a new test `block_replayer_peeking_state_roots` which fails if the same change is applied inside `get_state_root`.

Author: michaelsproul

consensus/state_processing/src/block_replayer.rs

@@ -3,6 +3,8 @@ use crate::{
     BlockProcessingError, BlockSignatureStrategy, ConsensusContext, SlotProcessingError,
     VerifyBlockRoot,
 };
+use itertools::Itertools;
+use std::iter::Peekable;
 use std::marker::PhantomData;
 use types::{BeaconState, BlindedPayload, ChainSpec, EthSpec, Hash256, SignedBeaconBlock, Slot};
 
@@ -25,7 +27,7 @@ pub struct BlockReplayer<
     'a,
     Spec: EthSpec,
     Error = BlockReplayError,
-    StateRootIter = StateRootIterDefault<Error>,
+    StateRootIter: Iterator<Item = Result<(Hash256, Slot), Error>> = StateRootIterDefault<Error>,
 > {
     state: BeaconState<Spec>,
     spec: &'a ChainSpec,
@@ -36,7 +38,7 @@ pub struct BlockReplayer<
     post_block_hook: Option<PostBlockHook<'a, Spec, Error>>,
     pre_slot_hook: Option<PreSlotHook<'a, Spec, Error>>,
     post_slot_hook: Option<PostSlotHook<'a, Spec, Error>>,
-    state_root_iter: Option<StateRootIter>,
+    pub(crate) state_root_iter: Option<Peekable<StateRootIter>>,
     state_root_miss: bool,
     _phantom: PhantomData<Error>,
 }
@@ -138,7 +140,7 @@ where
     /// `self.state.slot` to the `target_slot` supplied to `apply_blocks` (inclusive of both
     /// endpoints).
     pub fn state_root_iter(mut self, iter: StateRootIter) -> Self {
-        self.state_root_iter = Some(iter);
+        self.state_root_iter = Some(iter.peekable());
         self
     }
 
@@ -192,7 +194,7 @@ where
         // If a state root iterator is configured, use it to find the root.
         if let Some(ref mut state_root_iter) = self.state_root_iter {
             let opt_root = state_root_iter
-                .take_while(|res| res.as_ref().map_or(true, |(_, s)| *s <= slot))
+                .peeking_take_while(|res| res.as_ref().map_or(true, |(_, s)| *s <= slot))
                 .find(|res| res.as_ref().map_or(true, |(_, s)| *s == slot))
                 .transpose()?;
 

consensus/state_processing/src/per_block_processing/tests.rs

@@ -5,7 +5,7 @@ use crate::per_block_processing::errors::{
     DepositInvalid, HeaderInvalid, IndexedAttestationInvalid, IntoWithIndex,
     ProposerSlashingInvalid,
 };
-use crate::{per_block_processing, StateProcessingStrategy};
+use crate::{per_block_processing, BlockReplayError, BlockReplayer, StateProcessingStrategy};
 use crate::{
     per_block_processing::{process_operations, verify_exit::verify_exit},
     BlockSignatureStrategy, ConsensusContext, VerifyBlockRoot, VerifySignatures,
@@ -1035,3 +1035,51 @@ async fn fork_spanning_exit() {
     )
     .expect_err("phase0 exit does not verify against bellatrix state");
 }
+
+/// Check that the block replayer does not consume state roots unnecessarily.
+#[tokio::test]
+async fn block_replayer_peeking_state_roots() {
+    let harness = get_harness::<MainnetEthSpec>(EPOCH_OFFSET, VALIDATOR_COUNT).await;
+
+    let target_state = harness.get_current_state();
+    let target_block_root = harness.head_block_root();
+    let target_block = harness
+        .chain
+        .get_blinded_block(&target_block_root)
+        .unwrap()
+        .unwrap();
+
+    let parent_block_root = target_block.parent_root();
+    let parent_block = harness
+        .chain
+        .get_blinded_block(&parent_block_root)
+        .unwrap()
+        .unwrap();
+    let parent_state = harness
+        .chain
+        .get_state(&parent_block.state_root(), Some(parent_block.slot()))
+        .unwrap()
+        .unwrap();
+
+    // Omit the state root for `target_state` but provide a dummy state root at the *next* slot.
+    // If the block replayer is peeking at the state roots rather than consuming them, then the
+    // dummy state should still be there after block replay completes.
+    let dummy_state_root = Hash256::repeat_byte(0xff);
+    let dummy_slot = target_state.slot() + 1;
+    let state_root_iter = vec![Ok::<_, BlockReplayError>((dummy_state_root, dummy_slot))];
+    let block_replayer = BlockReplayer::new(parent_state, &harness.chain.spec)
+        .state_root_iter(state_root_iter.into_iter())
+        .no_signature_verification()
+        .apply_blocks(vec![target_block], None)
+        .unwrap();
+
+    assert_eq!(
+        block_replayer
+            .state_root_iter
+            .unwrap()
+            .next()
+            .unwrap()
+            .unwrap(),
+        (dummy_state_root, dummy_slot)
+    );
+}