[cleanup] Push account_id_to_shard_uid function in shard_layout (#12622)

`account_id_to_shard_uid` belongs to ShardLayout and not as a stand
alone function.

chain/chain/src/resharding/resharding_v2.rs

@@ -7,7 +7,7 @@ use crate::Chain;
 use near_chain_configs::{MutableConfigValue, ReshardingConfig, ReshardingHandle};
 use near_chain_primitives::error::Error;
 use near_primitives::hash::CryptoHash;
-use near_primitives::shard_layout::{account_id_to_shard_uid, ShardLayout};
+use near_primitives::shard_layout::ShardLayout;
 use near_primitives::types::chunk_extra::ChunkExtra;
 use near_primitives::types::{AccountId, ShardId, StateRoot};
 use near_store::flat::FlatStorageError;
@@ -82,7 +82,7 @@ fn get_checked_account_id_to_shard_uid_fn(
 ) -> impl Fn(&AccountId) -> ShardUId {
     let split_shard_ids: HashSet<_> = new_shards.into_iter().collect();
     move |account_id: &AccountId| {
-        let new_shard_uid = account_id_to_shard_uid(account_id, &next_epoch_shard_layout);
+        let new_shard_uid = next_epoch_shard_layout.account_id_to_shard_uid(account_id);
         // check that all accounts in the shard are mapped the shards that this shard will split
         // to according to shard layout
         assert!(

chain/chunks/src/client.rs

@@ -5,7 +5,7 @@ use itertools::Itertools;
 
 use near_pool::types::TransactionGroupIterator;
 use near_pool::{InsertTransactionResult, PoolIteratorWrapper, TransactionPool};
-use near_primitives::shard_layout::{account_id_to_shard_uid, ShardLayout, ShardUId};
+use near_primitives::shard_layout::{ShardLayout, ShardUId};
 use near_primitives::{
     epoch_info::RngSeed,
     sharding::{EncodedShardChunk, PartialEncodedChunk, ShardChunk, ShardChunkHeader},
@@ -146,7 +146,7 @@ impl ShardedTransactionPool {
 
         for tx in transactions {
             let signer_id = tx.transaction.signer_id();
-            let new_shard_uid = account_id_to_shard_uid(&signer_id, new_shard_layout);
+            let new_shard_uid = new_shard_layout.account_id_to_shard_uid(&signer_id);
             self.insert_transaction(new_shard_uid, tx);
         }
     }
@@ -161,7 +161,7 @@ mod tests {
     use near_primitives::{
         epoch_info::RngSeed,
         hash::CryptoHash,
-        shard_layout::{account_id_to_shard_uid, ShardLayout},
+        shard_layout::ShardLayout,
         transaction::SignedTransaction,
         types::{AccountId, ShardId},
     };
@@ -267,7 +267,7 @@ mod tests {
                     while let Some(tx) = group.next() {
                         total += 1;
                         let account_id = tx.transaction.signer_id();
-                        let tx_shard_uid = account_id_to_shard_uid(account_id, &new_shard_layout);
+                        let tx_shard_uid = new_shard_layout.account_id_to_shard_uid(account_id);
                         tracing::debug!("checking {account_id:?}:{tx_shard_uid} in {shard_uid}");
                         assert_eq!(shard_uid, tx_shard_uid);
                     }

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

@@ -20,7 +20,7 @@ use near_primitives::congestion_info::CongestionInfo;
 use near_primitives::epoch_block_info::BlockInfoV3;
 use near_primitives::epoch_manager::EpochConfig;
 use near_primitives::hash::hash;
-use near_primitives::shard_layout::{account_id_to_shard_uid, ShardLayout};
+use near_primitives::shard_layout::ShardLayout;
 use near_primitives::sharding::{ShardChunkHeader, ShardChunkHeaderV3};
 use near_primitives::stateless_validation::chunk_endorsements_bitmap::ChunkEndorsementsBitmap;
 use near_primitives::stateless_validation::partial_witness::PartialEncodedStateWitness;
@@ -3799,7 +3799,7 @@ fn test_get_shard_uids_pending_resharding_single() {
     let shard_layout_0 = ShardLayout::multi_shard_custom(vec![a.clone()], version);
     let shard_layout_1 = ShardLayout::derive_shard_layout(&shard_layout_0, b);
 
-    let s1 = account_id_to_shard_uid(&a, &shard_layout_0);
+    let s1 = shard_layout_0.account_id_to_shard_uid(&a);
 
     let shard_uids = test_get_shard_uids_pending_resharding_base(&[shard_layout_0, shard_layout_1]);
     assert_eq!(shard_uids, vec![s1].into_iter().collect::<HashSet<_>>());
@@ -3822,8 +3822,8 @@ fn test_get_shard_uids_pending_resharding_double_different() {
     let shard_layout_1 = ShardLayout::derive_shard_layout(&shard_layout_0, a.clone());
     let shard_layout_2 = ShardLayout::derive_shard_layout(&shard_layout_0, c);
 
-    let s0 = account_id_to_shard_uid(&a, &shard_layout_0);
-    let s1 = account_id_to_shard_uid(&b, &shard_layout_0);
+    let s0 = shard_layout_0.account_id_to_shard_uid(&a);
+    let s1 = shard_layout_0.account_id_to_shard_uid(&b);
 
     let shard_uids = test_get_shard_uids_pending_resharding_base(&[
         shard_layout_0,
@@ -3850,7 +3850,7 @@ fn test_get_shard_uids_pending_resharding_double_same() {
     let shard_layout_1 = ShardLayout::derive_shard_layout(&shard_layout_0, b);
     let shard_layout_2 = ShardLayout::derive_shard_layout(&shard_layout_0, c);
 
-    let s1 = account_id_to_shard_uid(&a, &shard_layout_0);
+    let s1 = shard_layout_0.account_id_to_shard_uid(&a);
 
     let shard_uids = test_get_shard_uids_pending_resharding_base(&[
         shard_layout_0,

core/primitives/src/shard_layout.rs

@@ -582,7 +582,7 @@ impl ShardLayout {
         )
     }
 
-    /// Maps an account to the shard that it belongs to given a shard layout
+    /// Maps an account to the shard_id that it belongs to in this shard_layout
     /// For V0, maps according to hash of account id
     /// For V1 and V2, accounts are divided to ranges, each range of account is mapped to a shard.
     pub fn account_id_to_shard_id(&self, account_id: &AccountId) -> ShardId {
@@ -598,8 +598,15 @@ impl ShardLayout {
         }
     }
 
+    /// Maps an account to the shard_uid that it belongs to in this shard_layout
+    #[inline]
+    pub fn account_id_to_shard_uid(&self, account_id: &AccountId) -> ShardUId {
+        ShardUId::from_shard_id_and_layout(self.account_id_to_shard_id(account_id), self)
+    }
+
     /// Given a parent shard id, return the shard uids for the shards in the current shard layout that
     /// are split from this parent shard. If this shard layout has no parent shard layout, return None
+    #[inline]
     pub fn get_children_shards_uids(&self, parent_shard_id: ShardId) -> Option<Vec<ShardUId>> {
         self.get_children_shards_ids(parent_shard_id).map(|shards| {
             shards.into_iter().map(|id| ShardUId::from_shard_id_and_layout(id, self)).collect()
@@ -852,14 +859,6 @@ fn validate_and_derive_shard_parent_map_v2(
     shards_parent_map
 }
 
-/// Maps an account to the shard that it belongs to given a shard_layout
-pub fn account_id_to_shard_uid(account_id: &AccountId, shard_layout: &ShardLayout) -> ShardUId {
-    ShardUId::from_shard_id_and_layout(
-        shard_layout.account_id_to_shard_id(account_id),
-        shard_layout,
-    )
-}
-
 /// `ShardUId` is a unique representation for shards from different shard layouts.
 ///
 /// Comparing to `ShardId`, which is just an ordinal number ranging from 0 to NUM_SHARDS-1,

core/store/src/config.rs

@@ -4,7 +4,7 @@ use crate::trie::{
 use crate::DBCol;
 use near_primitives::chains::MAINNET;
 use near_primitives::epoch_manager::EpochConfigStore;
-use near_primitives::shard_layout::{account_id_to_shard_uid, ShardLayout, ShardUId};
+use near_primitives::shard_layout::{ShardLayout, ShardUId};
 use near_primitives::types::AccountId;
 use near_primitives::version::{ProtocolFeature, PROTOCOL_VERSION};
 use near_time::Duration;
@@ -212,7 +212,7 @@ impl StoreConfig {
             let account_id = AccountId::from_str(account_id)
                 .expect("the hardcoded account id should guarantee to be valid");
             for shard_layout in shard_layouts.iter() {
-                let shard_uid = account_id_to_shard_uid(&account_id, &shard_layout);
+                let shard_uid = shard_layout.account_id_to_shard_uid(&account_id);
                 per_shard_max_bytes.insert(shard_uid, *bytes);
             }
         }

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

@@ -5,7 +5,7 @@ use near_chain_configs::test_genesis::{TestGenesisBuilder, ValidatorsSpec};
 use near_chain_configs::DEFAULT_GC_NUM_EPOCHS_TO_KEEP;
 use near_o11y::testonly::init_test_logger;
 use near_primitives::epoch_manager::EpochConfigStore;
-use near_primitives::shard_layout::{account_id_to_shard_uid, ShardLayout};
+use near_primitives::shard_layout::ShardLayout;
 use near_primitives::types::{AccountId, BlockHeightDelta, Gas, ShardId, ShardIndex};
 use near_primitives::version::{ProtocolFeature, PROTOCOL_VERSION};
 use rand::seq::SliceRandom;
@@ -534,7 +534,7 @@ fn test_resharding_v3_base(params: TestReshardingParameters) {
     base_epoch_config.shard_layout = base_shard_layout.clone();
 
     let new_boundary_account = "account6".parse().unwrap();
-    let parent_shard_uid = account_id_to_shard_uid(&new_boundary_account, &base_shard_layout);
+    let parent_shard_uid = base_shard_layout.account_id_to_shard_uid(&new_boundary_account);
     let mut epoch_config = base_epoch_config.clone();
     epoch_config.shard_layout =
         ShardLayout::derive_shard_layout(&base_shard_layout, new_boundary_account);

integration-tests/src/tests/client/resharding_v2.rs

@@ -12,7 +12,6 @@ use near_primitives::block::{Block, Tip};
 use near_primitives::epoch_manager::{AllEpochConfig, AllEpochConfigTestOverrides, EpochConfig};
 use near_primitives::hash::CryptoHash;
 use near_primitives::serialize::to_base64;
-use near_primitives::shard_layout::account_id_to_shard_uid;
 use near_primitives::stateless_validation::ChunkProductionKey;
 use near_primitives::transaction::{
     Action, DeployContractAction, FunctionCallAction, SignedTransaction,
@@ -444,7 +443,7 @@ impl TestReshardingEnv {
             let id = &tx.get_hash();
 
             let signer_account_id = tx.transaction.signer_id();
-            let shard_uid = account_id_to_shard_uid(signer_account_id, &shard_layout);
+            let shard_uid = shard_layout.account_id_to_shard_uid(signer_account_id);
 
             tracing::trace!(target: "test", tx=?id, ?signer_account_id, ?shard_uid, "checking tx");
 
@@ -522,7 +521,7 @@ fn check_account(env: &TestEnv, account_id: &AccountId, block: &Block) {
     let prev_hash = block.header().prev_hash();
     let shard_layout =
         env.clients[0].epoch_manager.get_shard_layout_from_prev_block(prev_hash).unwrap();
-    let shard_uid = account_id_to_shard_uid(account_id, &shard_layout);
+    let shard_uid = shard_layout.account_id_to_shard_uid(account_id);
     let shard_id = shard_uid.shard_id();
     let shard_index = shard_layout.get_shard_index(shard_id).unwrap();
     for (i, me) in env.validators.iter().enumerate() {