more random things

chain/chain/src/runtime/mod.rs

@@ -29,8 +29,8 @@ use near_primitives::state_part::PartId;
 use near_primitives::transaction::SignedTransaction;
 use near_primitives::trie_key::TrieKey;
 use near_primitives::types::{
-    new_shard_id_tmp, AccountId, Balance, BlockHeight, EpochHeight, EpochId, EpochInfoProvider,
-    Gas, MerkleHash, ShardId, StateChangeCause, StateRoot, StateRootNode,
+    shard_id_as_u32, AccountId, Balance, BlockHeight, EpochHeight, EpochId, EpochInfoProvider, Gas,
+    MerkleHash, ShardId, StateChangeCause, StateRoot, StateRootNode,
 };
 use near_primitives::version::{ProtocolFeature, ProtocolVersion};
 use near_primitives::views::{
@@ -223,7 +223,7 @@ impl NightshadeRuntime {
             epoch_manager.get_epoch_id_from_prev_block(prev_hash).map_err(Error::from)?;
         let shard_version =
             epoch_manager.get_shard_layout(&epoch_id).map_err(Error::from)?.version();
-        Ok(ShardUId { version: shard_version, shard_id: new_shard_id_tmp(shard_id) as u32 })
+        Ok(ShardUId { version: shard_version, shard_id: shard_id_as_u32(shard_id) })
     }
 
     fn get_shard_uid_from_epoch_id(
@@ -234,7 +234,7 @@ impl NightshadeRuntime {
         let epoch_manager = self.epoch_manager.read();
         let shard_version =
             epoch_manager.get_shard_layout(epoch_id).map_err(Error::from)?.version();
-        Ok(ShardUId { version: shard_version, shard_id: new_shard_id_tmp(shard_id) as u32 })
+        Ok(ShardUId { version: shard_version, shard_id: shard_id_as_u32(shard_id) })
     }
 
     fn account_id_to_shard_uid(

chain/chain/src/runtime/tests.rs

@@ -30,7 +30,8 @@ use near_primitives::block::Tip;
 use near_primitives::challenge::{ChallengesResult, PartialState, SlashedValidator};
 use near_primitives::transaction::{Action, DeleteAccountAction, StakeAction, TransferAction};
 use near_primitives::types::{
-    BlockHeightDelta, Nonce, ValidatorId, ValidatorInfoIdentifier, ValidatorKickoutReason,
+    new_shard_id_tmp, BlockHeightDelta, Nonce, ValidatorId, ValidatorInfoIdentifier,
+    ValidatorKickoutReason,
 };
 use near_primitives::validator_signer::ValidatorSigner;
 use near_primitives::views::{

chain/chunks/src/client.rs

@@ -6,6 +6,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::types::shard_id_as_u16;
 use near_primitives::{
     epoch_info::RngSeed,
     sharding::{EncodedShardChunk, PartialEncodedChunk, ShardChunk, ShardChunkHeader},
@@ -74,6 +75,7 @@ impl ShardedTransactionPool {
     /// For better security we want the seed to different in each shard.
     /// For testing purposes we want it to be the reproducible and derived from the `self.rng_seed` and `shard_id`
     fn random_seed(base_seed: &RngSeed, shard_id: ShardId) -> RngSeed {
+        let shard_id = shard_id_as_u16(shard_id);
         let mut res = *base_seed;
         res[0] = shard_id as u8;
         res[1] = (shard_id / 256) as u8;

chain/network/src/peer_manager/tests/snapshot_hosts.rs

@@ -18,6 +18,8 @@ use near_o11y::WithSpanContextExt;
 use near_primitives::hash::CryptoHash;
 use near_primitives::network::PeerId;
 use near_primitives::types::new_shard_id_tmp;
+use near_primitives::types::shard_id_as_u64;
+use near_primitives::types::shard_id_max;
 use near_primitives::types::EpochHeight;
 use near_primitives::types::ShardId;
 use peer_manager::testonly::FDS_PER_PEER;
@@ -371,12 +373,16 @@ async fn large_shard_id_in_cache() {
     let peer1 = pm.start_inbound(chain.clone(), peer1_config.clone()).await.handshake(clock).await;
 
     tracing::info!(target:"test", "Send a SnapshotHostInfo message with very large shard ids.");
-    let max_shard_id: ShardId = ShardId::MAX;
+    let max_shard_id = shard_id_max();
+    let max_shard_id_minus_one = shard_id_as_u64(max_shard_id) - 1;
+    let max_shard_id_minus_one = new_shard_id_tmp(max_shard_id_minus_one);
     let big_shard_info = Arc::new(SnapshotHostInfo::new(
         peer1_config.node_id(),
         CryptoHash::hash_borsh(1234_u64),
         1234,
-        vec![0, 1232232, max_shard_id - 1, max_shard_id].into_iter().map(Into::into).collect(),
+        vec![new_shard_id_tmp(0), new_shard_id_tmp(1232232), max_shard_id_minus_one, max_shard_id]
+            .into_iter()
+            .collect(),
         &peer1_config.node_key,
     ));
 
@@ -447,7 +453,7 @@ async fn too_many_shards_truncate() {
     assert_eq!(info.shards.len(), MAX_SHARDS_PER_SNAPSHOT_HOST_INFO);
     for &shard_id in &info.shards {
         // Shard ids are taken from the original vector
-        assert!(shard_id < 2 * MAX_SHARDS_PER_SNAPSHOT_HOST_INFO as u64);
+        assert!(shard_id_as_u64(shard_id) < 2 * MAX_SHARDS_PER_SNAPSHOT_HOST_INFO as u64);
     }
     // The shard_ids are sorted and unique (no two elements are equal, hence the < condition instead of <=)
     assert!(info.shards.windows(2).all(|twoelems| twoelems[0] < twoelems[1]));

core/primitives-core/Cargo.toml

@@ -33,14 +33,14 @@ insta.workspace = true
 expect-test.workspace = true
 
 [features]
-default = ["new_shard_id"] # DO NOT COMMIT THIS
-# default = []
+default = []
 protocol_feature_fix_staking_threshold = []
 protocol_feature_fix_contract_loading_cost = []
 protocol_feature_reject_blocks_with_outdated_protocol_version = []
 protocol_feature_nonrefundable_transfer_nep491 = []
 
 # TODO(wacban) remove after the transition is done
+# default = ["new_shard_id"] # DO NOT COMMIT THIS
 new_shard_id = []
 
 nightly = [

core/primitives-core/src/types.rs

@@ -76,6 +76,14 @@ pub fn new_shard_id_vec_tmp(vec: &[u64]) -> Vec<ShardId> {
     vec.iter().copied().map(new_shard_id_tmp).collect()
 }
 
+pub fn shard_id_max() -> ShardId {
+    #[cfg(not(feature = "new_shard_id"))]
+    return ShardId::MAX;
+
+    #[cfg(feature = "new_shard_id")]
+    return ShardId::max();
+}
+
 // TODO(wacban) This is a temporary solution to aid the transition to having
 // ShardId as a newtype. It should be replaced / removed / inlined once the
 // transition is complete.

core/primitives/src/trie_key.rs

@@ -853,7 +853,7 @@ mod tests {
         let shard_layout = ShardLayout::for_protocol_version(PROTOCOL_VERSION);
         let max_id = shard_layout.shard_ids().max().unwrap();
         assert!(
-            max_id <= u16::MAX as u64,
+            shard_id_as_u64(max_id) <= u16::MAX as u64,
             "buffered receipt trie key optimization broken, must fit in a u16"
         );
     }

core/store/src/trie/state_parts.rs

@@ -30,7 +30,7 @@ use near_primitives::hash::{hash, CryptoHash};
 use near_primitives::state::FlatStateValue;
 use near_primitives::state_part::PartId;
 use near_primitives::state_record::is_contract_code_key;
-use near_primitives::types::{ShardId, StateRoot};
+use near_primitives::types::{shard_id_max, ShardId, StateRoot};
 use near_vm_runner::ContractCode;
 use std::collections::{HashMap, HashSet};
 use std::sync::Arc;
@@ -129,7 +129,7 @@ impl Trie {
         part_id: PartId,
     ) -> Result<(PartialState, Vec<u8>, Vec<u8>), StorageError> {
         let shard_id: ShardId = self.flat_storage_chunk_view.as_ref().map_or(
-            ShardId::MAX, // Fake value for metrics.
+            shard_id_max(), // Fake value for metrics.
             |chunk_view| chunk_view.shard_uid().shard_id(),
         );
         let _span = tracing::debug_span!(
@@ -183,7 +183,7 @@ impl Trie {
         state_trie: &Trie,
     ) -> Result<PartialState, StorageError> {
         let shard_id: ShardId = self.flat_storage_chunk_view.as_ref().map_or(
-            ShardId::MAX, // Fake value for metrics.
+            shard_id_max(), // Fake value for metrics.
             |chunk_view| chunk_view.shard_uid().shard_id(),
         );
         let _span = tracing::debug_span!(

core/store/src/trie/trie_storage.rs

@@ -10,7 +10,7 @@ use near_o11y::metrics::prometheus::core::{GenericCounter, GenericGauge};
 use near_primitives::challenge::PartialState;
 use near_primitives::hash::CryptoHash;
 use near_primitives::shard_layout::ShardUId;
-use near_primitives::types::ShardId;
+use near_primitives::types::{shard_id_as_u64, ShardId};
 use std::collections::{HashMap, HashSet, VecDeque};
 use std::num::NonZeroUsize;
 use std::sync::{Arc, Mutex};
@@ -103,7 +103,7 @@ impl TrieCacheInner {
         assert!(total_size_limit > 0);
         // `itoa` is much faster for printing shard_id to a string than trivial alternatives.
         let mut buffer = itoa::Buffer::new();
-        let shard_id_str = buffer.format(shard_id);
+        let shard_id_str = buffer.format(shard_id_as_u64(shard_id));
 
         let metrics_labels: [&str; 2] = [&shard_id_str, if is_view { "1" } else { "0" }];
         let metrics = TrieCacheMetrics {

runtime/runtime/src/congestion_control.rs

@@ -9,7 +9,7 @@ use near_primitives::receipt::{
     Receipt, ReceiptEnum, ReceiptOrStateStoredReceipt, StateStoredReceipt,
     StateStoredReceiptMetadata,
 };
-use near_primitives::types::{new_shard_id_tmp, EpochInfoProvider, Gas, ShardId};
+use near_primitives::types::{shard_id_as_u16, EpochInfoProvider, Gas, ShardId};
 use near_primitives::version::ProtocolFeature;
 use near_store::trie::receipts_column_helper::{
     DelayedReceiptQueue, ReceiptIterator, ShardsOutgoingReceiptBuffer, TrieQueue,
@@ -460,7 +460,7 @@ pub fn bootstrap_congestion_info(
         // It is also irrelevant, since the bootstrapped value is only used at
         // the start of applying a chunk on this shard. Other shards will only
         // see and act on the first congestion info after that.
-        allowed_shard: new_shard_id_tmp(shard_id) as u16,
+        allowed_shard: shard_id_as_u16(shard_id),
     }))
 }