fix test_chunk_state_witness_bad_shard_id test

chain/client/src/view_client.rs

@@ -280,7 +280,7 @@ impl ViewClientActor {
             let cps: Vec<AccountId> = shard_ids
                 .iter()
                 .map(|&shard_id| {
-                    let cp = epoch_info.sample_chunk_producer(block_height, shard_id);
+                    let cp = epoch_info.sample_chunk_producer(block_height, shard_id).unwrap();
                     let cp = epoch_info.get_validator(cp).account_id().clone();
                     cp
                 })

chain/epoch-manager/src/lib.rs

@@ -1029,7 +1029,7 @@ impl EpochManager {
         shard_id: ShardId,
     ) -> Result<ValidatorStake, EpochError> {
         let epoch_info = self.get_epoch_info(epoch_id)?;
-        let validator_id = Self::chunk_producer_from_info(&epoch_info, height, shard_id);
+        let validator_id = Self::chunk_producer_from_info(&epoch_info, height, shard_id)?;
         Ok(epoch_info.get_validator(validator_id))
     }
 
@@ -1556,8 +1556,12 @@ impl EpochManager {
         epoch_info: &EpochInfo,
         height: BlockHeight,
         shard_id: ShardId,
-    ) -> ValidatorId {
-        epoch_info.sample_chunk_producer(height, shard_id)
+    ) -> Result<ValidatorId, EpochError> {
+        epoch_info.sample_chunk_producer(height, shard_id).ok_or_else(|| {
+            EpochError::ChunkProducerSelectionError(format!(
+                "Invalid shard {shard_id} for height {height}"
+            ))
+        })
     }
 
     /// Returns true, if given current block info, next block supposed to be in the next epoch.

chain/epoch-manager/src/tests/mod.rs

@@ -1131,7 +1131,7 @@ fn test_expected_chunks_prev_block_not_produced() {
         let prev_block_info = epoch_manager.get_block_info(&prev_block).unwrap();
         let prev_height = prev_block_info.height();
         let expected_chunk_producer =
-            EpochManager::chunk_producer_from_info(&epoch_info, prev_height + 1, 0);
+            EpochManager::chunk_producer_from_info(&epoch_info, prev_height + 1, 0).unwrap();
         // test1 does not produce blocks during first epoch
         if block_producer == 0 && epoch_id == initial_epoch_id {
             expected += 1;
@@ -1473,7 +1473,8 @@ fn test_chunk_validator_kickout() {
                         &epoch_info,
                         height,
                         shard_id as u64,
-                    );
+                    )
+                    .unwrap();
                     // test1 skips chunks
                     if chunk_producer == 0 {
                         expected += 1;

chain/epoch-manager/src/types.rs

@@ -138,7 +138,8 @@ impl EpochInfoAggregator {
                 epoch_info,
                 prev_block_height + 1,
                 i as ShardId,
-            );
+            )
+            .unwrap();
             let tracker = self.shard_tracker.entry(i as ShardId).or_insert_with(HashMap::new);
             tracker
                 .entry(chunk_validator_id)

chain/epoch-manager/src/validator_selection.rs

@@ -606,7 +606,7 @@ mod tests {
                 let cp = epoch_info.sample_chunk_producer(h, shard_id);
                 // Don't read too much into this. The reason the ValidatorId always
                 // equals the ShardId is because the validators are assigned to shards in order.
-                assert_eq!(cp, shard_id)
+                assert_eq!(cp, Some(shard_id))
             }
         }
 
@@ -632,7 +632,7 @@ mod tests {
         for shard_id in 0..num_shards {
             let mut counts: [i32; 2] = [0, 0];
             for h in 0..100_000 {
-                let cp = epoch_info.sample_chunk_producer(h, shard_id);
+                let cp = epoch_info.sample_chunk_producer(h, shard_id).unwrap();
                 // if ValidatorId is in the second half then it is the lower
                 // stake validator (because they are sorted by decreasing stake).
                 let index = if cp >= num_shards { 1 } else { 0 };

core/primitives/src/epoch_manager.rs

@@ -1081,33 +1081,37 @@ pub mod epoch_info {
             }
         }
 
-        pub fn sample_chunk_producer(&self, height: BlockHeight, shard_id: ShardId) -> ValidatorId {
+        pub fn sample_chunk_producer(
+            &self,
+            height: BlockHeight,
+            shard_id: ShardId,
+        ) -> Option<ValidatorId> {
             match &self {
                 Self::V1(v1) => {
                     let cp_settlement = &v1.chunk_producers_settlement;
-                    let shard_cps = &cp_settlement[shard_id as usize];
-                    shard_cps[(height as u64 % (shard_cps.len() as u64)) as usize]
+                    let shard_cps = cp_settlement.get(shard_id as usize)?;
+                    shard_cps.get((height as u64 % (shard_cps.len() as u64)) as usize).copied()
                 }
                 Self::V2(v2) => {
                     let cp_settlement = &v2.chunk_producers_settlement;
-                    let shard_cps = &cp_settlement[shard_id as usize];
-                    shard_cps[(height as u64 % (shard_cps.len() as u64)) as usize]
+                    let shard_cps = cp_settlement.get(shard_id as usize)?;
+                    shard_cps.get((height as u64 % (shard_cps.len() as u64)) as usize).copied()
                 }
                 Self::V3(v3) => {
                     let protocol_version = self.protocol_version();
                     let seed =
                         Self::chunk_produce_seed(protocol_version, &v3.rng_seed, height, shard_id);
                     let shard_id = shard_id as usize;
-                    let sample = v3.chunk_producers_sampler[shard_id].sample(seed);
-                    v3.chunk_producers_settlement[shard_id][sample]
+                    let sample = v3.chunk_producers_sampler.get(shard_id)?.sample(seed);
+                    v3.chunk_producers_settlement.get(shard_id)?.get(sample).copied()
                 }
                 Self::V4(v4) => {
                     let protocol_version = self.protocol_version();
                     let seed =
                         Self::chunk_produce_seed(protocol_version, &v4.rng_seed, height, shard_id);
                     let shard_id = shard_id as usize;
-                    let sample = v4.chunk_producers_sampler[shard_id].sample(seed);
-                    v4.chunk_producers_settlement[shard_id][sample]
+                    let sample = v4.chunk_producers_sampler.get(shard_id)?.sample(seed);
+                    v4.chunk_producers_settlement.get(shard_id)?.get(sample).copied()
                 }
             }
         }

core/primitives/src/errors.rs

@@ -865,6 +865,8 @@ pub enum EpochError {
     },
     /// Error selecting validators for a chunk.
     ChunkValidatorSelectionError(String),
+    /// Error selecting chunk producer for a shard.
+    ChunkProducerSelectionError(String),
 }
 
 impl std::error::Error for EpochError {}
@@ -892,6 +894,9 @@ impl Display for EpochError {
             EpochError::ChunkValidatorSelectionError(err) => {
                 write!(f, "Error selecting validators for a chunk: {}", err)
             }
+            EpochError::ChunkProducerSelectionError(err) => {
+                write!(f, "Error selecting chunk producer: {}", err)
+            }
         }
     }
 }
@@ -915,6 +920,9 @@ impl Debug for EpochError {
             EpochError::ChunkValidatorSelectionError(err) => {
                 write!(f, "ChunkValidatorSelectionError({})", err)
             }
+            EpochError::ChunkProducerSelectionError(err) => {
+                write!(f, "ChunkProducerSelectionError({})", err)
+            }
         }
     }
 }

tools/state-viewer/src/epoch_info.rs

@@ -96,7 +96,7 @@ fn display_block_and_chunk_producers(
         let cps: Vec<String> = shard_ids
             .iter()
             .map(|&shard_id| {
-                let cp = epoch_info.sample_chunk_producer(block_height, shard_id);
+                let cp = epoch_info.sample_chunk_producer(block_height, shard_id).unwrap();
                 let cp = epoch_info.get_validator(cp).account_id().clone();
                 cp.as_str().to_string()
             })
@@ -280,7 +280,8 @@ fn display_validator_info(
                     .iter()
                     .map(|&shard_id| (block_height, shard_id))
                     .filter(|&(block_height, shard_id)| {
-                        epoch_info.sample_chunk_producer(block_height, shard_id) == *validator_id
+                        epoch_info.sample_chunk_producer(block_height, shard_id)
+                            == Some(*validator_id)
                     })
                     .collect::<Vec<(BlockHeight, ShardId)>>()
             })