refactor: replace EpochInfoProvider::account_id_to_shard_id with shar… (

#12754) [issue 12542](#12542) --------- Co-authored-by: slava <[email protected]>
near · Jan 17, 2025 · b6f82d6 · b6f82d6
1 parent 43fc799
commit b6f82d6
Show file tree

Hide file tree

Showing 5 changed files with 7 additions and 33 deletions.
diff --git a/chain/chain/src/runtime/tests.rs b/chain/chain/src/runtime/tests.rs
@@ -403,13 +403,8 @@ impl TestEnv {
     }
 
     pub fn view_account(&self, account_id: &AccountId) -> AccountView {
-        let shard_id = EpochInfoProvider::account_id_to_shard_id(
-            &*self.epoch_manager,
-            account_id,
-            &self.head.epoch_id,
-        )
-        .unwrap();
         let shard_layout = self.epoch_manager.get_shard_layout(&self.head.epoch_id).unwrap();
+        let shard_id = shard_layout.account_id_to_shard_id(account_id);
         let shard_index = shard_layout.get_shard_index(shard_id).unwrap();
         let shard_uid = self.epoch_manager.shard_id_to_uid(shard_id, &self.head.epoch_id).unwrap();
         self.runtime

diff --git a/chain/epoch-manager/src/lib.rs b/chain/epoch-manager/src/lib.rs
@@ -126,14 +126,6 @@ impl EpochInfoProvider for EpochManagerHandle {
         epoch_manager.config.chain_id().into()
     }
 
-    fn account_id_to_shard_id(
-        &self,
-        account_id: &AccountId,
-        epoch_id: &EpochId,
-    ) -> Result<ShardId, EpochError> {
-        EpochManagerAdapter::account_id_to_shard_id(self, account_id, epoch_id)
-    }
-
     fn shard_layout(&self, epoch_id: &EpochId) -> Result<ShardLayout, EpochError> {
         let epoch_manager = self.read();
         epoch_manager.get_shard_layout(epoch_id)

diff --git a/core/primitives/src/test_utils.rs b/core/primitives/src/test_utils.rs
@@ -21,7 +21,7 @@ use crate::version::PROTOCOL_VERSION;
 use crate::views::{ExecutionStatusView, FinalExecutionOutcomeView, FinalExecutionStatus};
 use near_crypto::vrf::Value;
 use near_crypto::{EmptySigner, PublicKey, SecretKey, Signature, Signer};
-use near_primitives_core::types::{BlockHeight, MerkleHash, ProtocolVersion, ShardId};
+use near_primitives_core::types::{BlockHeight, MerkleHash, ProtocolVersion};
 use std::collections::HashMap;
 use std::sync::Arc;
 
@@ -1037,14 +1037,6 @@ impl EpochInfoProvider for MockEpochInfoProvider {
         "localnet".into()
     }
 
-    fn account_id_to_shard_id(
-        &self,
-        account_id: &AccountId,
-        _epoch_id: &EpochId,
-    ) -> Result<ShardId, EpochError> {
-        Ok(self.shard_layout.account_id_to_shard_id(account_id))
-    }
-
     fn shard_layout(&self, _epoch_id: &EpochId) -> Result<ShardLayout, EpochError> {
         Ok(self.shard_layout.clone())
     }

diff --git a/core/primitives/src/types.rs b/core/primitives/src/types.rs
@@ -1177,13 +1177,6 @@ pub trait EpochInfoProvider: Send + Sync {
     /// Get the chain_id of the chain this epoch belongs to
     fn chain_id(&self) -> String;
 
-    /// Which shard the account belongs to in the given epoch.
-    fn account_id_to_shard_id(
-        &self,
-        account_id: &AccountId,
-        epoch_id: &EpochId,
-    ) -> Result<ShardId, EpochError>;
-
     fn shard_layout(&self, epoch_id: &EpochId) -> Result<ShardLayout, EpochError>;
 }
 

diff --git a/runtime/runtime/src/congestion_control.rs b/runtime/runtime/src/congestion_control.rs
@@ -362,7 +362,8 @@ impl ReceiptSinkV2 {
         epoch_info_provider: &dyn EpochInfoProvider,
     ) -> Result<(), RuntimeError> {
         let shard = epoch_info_provider
-            .account_id_to_shard_id(receipt.receiver_id(), &apply_state.epoch_id)?;
+            .shard_layout(&apply_state.epoch_id)?
+            .account_id_to_shard_id(receipt.receiver_id());
 
         let size = compute_receipt_size(&receipt)?;
         let gas = compute_receipt_congestion_gas(&receipt, &apply_state.config)?;
@@ -792,8 +793,9 @@ impl<'a> DelayedReceiptQueueWrapper<'a> {
         let receiver_id = receipt.get_receipt().receiver_id();
         let receipt_shard_id = self
             .epoch_info_provider
-            .account_id_to_shard_id(receiver_id, &self.epoch_id)
-            .expect("account_id_to_shard_id should never fail");
+            .shard_layout(&self.epoch_id)
+            .unwrap()
+            .account_id_to_shard_id(receiver_id);
         receipt_shard_id == self.shard_id
     }