address PR comments

core/primitives/src/shard_layout.rs

@@ -94,6 +94,10 @@ type ShardsSplitMapV2 = BTreeMap<ShardId, Vec<ShardId>>;
 /// A mapping from the child shard to the parent shard.
 type ShardsParentMapV2 = BTreeMap<ShardId, ShardId>;
 
+pub fn shard_uids_to_ids(shard_uids: &[ShardUId]) -> Vec<ShardId> {
+    shard_uids.iter().map(|shard_uid| shard_uid.shard_id()).collect_vec()
+}
+
 fn new_shard_ids_vec(shard_ids: Vec<u64>) -> Vec<ShardId> {
     shard_ids.into_iter().map(Into::into).collect()
 }

integration-tests/src/test_loop/tests/resharding_v3.rs

@@ -5,7 +5,7 @@ use near_async::time::Duration;
 use near_chain_configs::test_genesis::{TestGenesisBuilder, ValidatorsSpec};
 use near_o11y::testonly::init_test_logger;
 use near_primitives::epoch_manager::EpochConfigStore;
-use near_primitives::shard_layout::ShardLayout;
+use near_primitives::shard_layout::{shard_uids_to_ids, ShardLayout};
 use near_primitives::types::{AccountId, BlockHeightDelta, ShardId, ShardIndex};
 use near_primitives::version::{ProtocolFeature, PROTOCOL_VERSION};
 use std::cell::Cell;
@@ -29,8 +29,7 @@ use crate::test_loop::utils::setups::{
     derive_new_epoch_config_from_boundary, two_upgrades_voting_schedule,
 };
 use crate::test_loop::utils::sharding::{
-    get_shards_needs_for_next_epoch, get_tracked_shards, print_and_assert_shard_accounts,
-    shard_uids_to_ids,
+    get_shards_will_care_about, get_tracked_shards, print_and_assert_shard_accounts,
 };
 use crate::test_loop::utils::transactions::{
     check_txs, create_account, deploy_contract, get_smallest_height_head,
@@ -539,14 +538,14 @@ fn test_resharding_v3_base(params: TestReshardingParameters) {
             tip.epoch_id.0,
         );
         for (client_index, client) in clients.iter().enumerate() {
-            let tracked_shards =
-                shard_uids_to_ids(&get_tracked_shards(client, &tip.last_block_hash));
+            let tracked_shards = get_tracked_shards(client, &tip.last_block_hash);
+            let tracked_shards = shard_uids_to_ids(&tracked_shards);
             // That's not accurate in case of tracked shard schedule: it won't return parent shard before resharding boundary, if we track child after resharding.
-            let shards_needs_for_next_epoch =
-                shard_uids_to_ids(&get_shards_needs_for_next_epoch(client, &tip.last_block_hash));
+            let shards_will_care_about = &get_shards_will_care_about(client, &tip.last_block_hash);
+            let shards_will_care_about = shard_uids_to_ids(shards_will_care_about);
             let signer = client.validator_signer.get().unwrap();
             let account_id = signer.validator_id().as_str();
-            println!("client_{client_index}: id={account_id:?} tracks={tracked_shards:?}\tneeds_for_next_epoch={shards_needs_for_next_epoch:?}");
+            println!("client_{client_index}: id={account_id:?} tracks={tracked_shards:?}\twill_care_about={shards_will_care_about:?}");
         }
 
         // Check that all chunks are included.

integration-tests/src/test_loop/utils/setups.rs

@@ -73,6 +73,8 @@ pub fn derive_new_epoch_config_from_boundary(
     epoch_config
 }
 
+/// Two protocol upgrades would happen as soon as possible,
+/// usually in two consecutive epochs, unless upgrade voting decides differently.
 pub fn two_upgrades_voting_schedule(
     target_protocol_version: ProtocolVersion,
 ) -> ProtocolUpgradeVotingSchedule {

integration-tests/src/test_loop/utils/sharding.rs

@@ -1,4 +1,3 @@
-use itertools::Itertools;
 use near_chain::types::Tip;
 use near_client::Client;
 use near_epoch_manager::EpochManagerAdapter;
@@ -148,17 +147,16 @@ pub fn get_tracked_shards(client: &Client, block_hash: &CryptoHash) -> Vec<Shard
     get_tracked_shards_from_prev_block(client, block_header.prev_hash())
 }
 
-pub fn get_shards_needs_for_next_epoch(client: &Client, block_hash: &CryptoHash) -> Vec<ShardUId> {
+pub fn get_shards_will_care_about(client: &Client, block_hash: &CryptoHash) -> Vec<ShardUId> {
     let signer = client.validator_signer.get();
     let account_id = signer.as_ref().map(|s| s.validator_id());
-    let prev_block_hash = *client.chain.get_block_header(block_hash).unwrap().prev_hash();
-    let shard_layout =
-        client.epoch_manager.get_shard_layout_from_prev_block(&prev_block_hash).unwrap();
+    let block_header = client.chain.get_block_header(block_hash).unwrap();
+    let shard_layout = client.epoch_manager.get_shard_layout(&block_header.epoch_id()).unwrap();
     let mut shards_needs_for_next_epoch = vec![];
     for shard_uid in shard_layout.shard_uids() {
         if client.shard_tracker.will_care_about_shard(
             account_id,
-            &prev_block_hash,
+            &block_header.prev_hash(),
             shard_uid.shard_id(),
             true,
         ) {
@@ -167,7 +165,3 @@ pub fn get_shards_needs_for_next_epoch(client: &Client, block_hash: &CryptoHash)
     }
     shards_needs_for_next_epoch
 }
-
-pub fn shard_uids_to_ids(shard_uids: &[ShardUId]) -> Vec<ShardId> {
-    shard_uids.iter().map(|shard_uid| shard_uid.shard_id()).collect_vec()
-}