diff --git a/chain/chain/src/resharding/resharding_v2.rs b/chain/chain/src/resharding/resharding_v2.rs index e2e56bde68e..be853b33293 100644 --- a/chain/chain/src/resharding/resharding_v2.rs +++ b/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!( diff --git a/chain/chunks/src/client.rs b/chain/chunks/src/client.rs index 45af5126a62..0a334e89cf2 100644 --- a/chain/chunks/src/client.rs +++ b/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); } diff --git a/chain/epoch-manager/src/tests/mod.rs b/chain/epoch-manager/src/tests/mod.rs index 6c11fe7e0ea..a9dfb2f79a2 100644 --- a/chain/epoch-manager/src/tests/mod.rs +++ b/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::>()); @@ -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, diff --git a/core/primitives/src/shard_layout.rs b/core/primitives/src/shard_layout.rs index 03fa6fb0df9..c96f4d70fb2 100644 --- a/core/primitives/src/shard_layout.rs +++ b/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> { 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, diff --git a/core/store/src/config.rs b/core/store/src/config.rs index ff0798bd1ae..6aa06864c03 100644 --- a/core/store/src/config.rs +++ b/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); } } diff --git a/integration-tests/src/test_loop/tests/resharding_v3.rs b/integration-tests/src/test_loop/tests/resharding_v3.rs index 0c8c6dd9744..0bda17d80bf 100644 --- a/integration-tests/src/test_loop/tests/resharding_v3.rs +++ b/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); diff --git a/integration-tests/src/tests/client/resharding_v2.rs b/integration-tests/src/tests/client/resharding_v2.rs index 044ab25be12..c6782846c8d 100644 --- a/integration-tests/src/tests/client/resharding_v2.rs +++ b/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() {