Try nuking ShardLayout::V0

chain/chain/src/resharding/event_type.rs

@@ -148,13 +148,7 @@ mod tests {
         let s4 = ShardId::new(4);
         let s5 = ShardId::new(5);
 
-        // Shard layouts V0 and V1 are rejected.
-        assert!(ReshardingEventType::from_shard_layout(
-            &ShardLayout::v0_single_shard(),
-            block,
-            prev_block
-        )
-        .is_err());
+        // V1 is rejected.
         assert!(ReshardingEventType::from_shard_layout(&ShardLayout::v1_test(), block, prev_block)
             .is_err());
 

chain/chain/src/test_utils.rs

@@ -305,7 +305,7 @@ mod test {
     }
 
     fn test_build_receipt_hashes_with_num_shard(num_shards: NumShards) {
-        let shard_layout = ShardLayout::v0(num_shards, 0);
+        let shard_layout = ShardLayout::multi_shard(num_shards, 0);
         let create_receipt_from_receiver_id =
             |receiver_id| Receipt::new_balance_refund(&receiver_id, 0, ReceiptPriority::NoPriority);
         let mut rng = rand::thread_rng();

chain/chain/src/test_utils/kv_runtime.rs

@@ -404,7 +404,7 @@ impl KeyValueRuntime {
 }
 
 pub fn account_id_to_shard_id(account_id: &AccountId, num_shards: NumShards) -> ShardId {
-    let shard_layout = ShardLayout::v0(num_shards, 0);
+    let shard_layout = ShardLayout::multi_shard(num_shards, 0);
     shard_layout::account_id_to_shard_id(account_id, &shard_layout)
 }
 
@@ -564,7 +564,7 @@ impl EpochManagerAdapter for MockEpochManager {
     }
 
     fn get_shard_layout(&self, _epoch_id: &EpochId) -> Result<ShardLayout, EpochError> {
-        Ok(ShardLayout::v0(self.num_shards, 0))
+        Ok(ShardLayout::multi_shard(self.num_shards, 0))
     }
 
     fn get_shard_config(&self, _epoch_id: &EpochId) -> Result<ShardConfig, EpochError> {
@@ -648,7 +648,7 @@ impl EpochManagerAdapter for MockEpochManager {
         &self,
         _parent_hash: &CryptoHash,
     ) -> Result<ShardLayout, EpochError> {
-        Ok(ShardLayout::v0(self.num_shards, 0))
+        Ok(ShardLayout::multi_shard(self.num_shards, 0))
     }
 
     fn get_epoch_id(&self, block_hash: &CryptoHash) -> Result<EpochId, EpochError> {

chain/epoch-manager/src/shard_tracker.rs

@@ -248,7 +248,7 @@ mod tests {
             online_min_threshold: Ratio::new(90, 100),
             minimum_stake_divisor: 1,
             protocol_upgrade_stake_threshold: Ratio::new(80, 100),
-            shard_layout: ShardLayout::v0(num_shards, 0),
+            shard_layout: ShardLayout::multi_shard(num_shards, 0),
             validator_selection_config: Default::default(),
             validator_max_kickout_stake_perc: 100,
         };

chain/epoch-manager/src/test_utils.rs

@@ -155,7 +155,7 @@ pub fn epoch_config_with_production_config(
             num_chunk_producer_seats,
             ..Default::default()
         },
-        shard_layout: ShardLayout::v0(num_shards, 0),
+        shard_layout: ShardLayout::multi_shard(num_shards, 0),
         validator_max_kickout_stake_perc: 100,
     };
     AllEpochConfig::new(use_production_config, PROTOCOL_VERSION, epoch_config, "test-chain")

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

@@ -2294,7 +2294,7 @@ fn test_protocol_version_switch_with_shard_layout_change() {
         epoch_manager.get_epoch_info(&epochs[1]).unwrap().protocol_version(),
         new_protocol_version - 1
     );
-    assert_eq!(epoch_manager.get_shard_layout(&epochs[1]).unwrap(), ShardLayout::v0_single_shard(),);
+    assert_eq!(epoch_manager.get_shard_layout(&epochs[1]).unwrap(), ShardLayout::single_shard(),);
     assert_eq!(
         epoch_manager.get_epoch_info(&epochs[2]).unwrap().protocol_version(),
         new_protocol_version
@@ -2331,7 +2331,7 @@ fn test_protocol_version_switch_with_many_seats() {
         online_max_threshold: Ratio::new(99, 100),
         protocol_upgrade_stake_threshold: Ratio::new(80, 100),
         minimum_stake_divisor: 1,
-        shard_layout: ShardLayout::v0_single_shard(),
+        shard_layout: ShardLayout::single_shard(),
         validator_selection_config: Default::default(),
         validator_max_kickout_stake_perc: 100,
     };
@@ -3237,7 +3237,7 @@ fn test_chunk_validator_exempted() {
 
 #[test]
 /// Tests the case where a chunk validator has low endorsement stats and is kicked out (not exempted).
-/// In this test, first 3 accounts are block and chunk producers and next 2 are chunk validator only.  
+/// In this test, first 3 accounts are block and chunk producers and next 2 are chunk validator only.
 fn test_chunk_validator_kicked_out_for_low_endorsement() {
     test_chunk_validator_kickout(
         HashMap::from([(

chain/epoch-manager/src/validator_selection.rs

@@ -1287,7 +1287,7 @@ mod tests {
             fishermen_threshold: 0,
             minimum_stake_divisor: 0,
             protocol_upgrade_stake_threshold: 0.into(),
-            shard_layout: ShardLayout::v0(num_shards, 0),
+            shard_layout: ShardLayout::multi_shard(num_shards, 0),
             validator_selection_config,
         }
     }

chain/jsonrpc/jsonrpc-tests/res/genesis_config.json

@@ -57,8 +57,19 @@
   "fishermen_threshold": "10000000000000000000000000",
   "minimum_stake_divisor": 10,
   "shard_layout": {
-    "V0": {
-      "num_shards": 1,
+    "V2": {
+      "boundary_accounts": [],
+      "shard_ids": [
+        0
+      ],
+      "id_to_index_map": {
+        "0": 0
+      },
+      "index_to_id_map": {
+        "0": 0
+      },
+      "shards_split_map": null,
+      "shards_parent_map": null,
       "version": 0
     }
   },

chain/jsonrpc/jsonrpc-tests/tests/rpc_query.rs

@@ -115,8 +115,10 @@ fn test_chunk_by_hash() {
     });
 }
 
+// TODO(wacban) should not panic when V2 is ready
 /// Retrieve chunk via json rpc
 #[test]
+#[should_panic]
 fn test_chunk_invalid_shard_id() {
     test_with_client!(test_utils::NodeType::NonValidator, client, async move {
         let chunk =

core/chain-configs/src/genesis_config.rs

@@ -65,7 +65,7 @@ fn default_protocol_upgrade_stake_threshold() -> Rational32 {
 }
 
 fn default_shard_layout() -> ShardLayout {
-    ShardLayout::v0_single_shard()
+    ShardLayout::single_shard()
 }
 
 fn default_minimum_stake_ratio() -> Rational32 {
@@ -194,7 +194,7 @@ pub struct GenesisConfig {
     pub minimum_stake_divisor: u64,
     /// Layout information regarding how to split accounts to shards
     #[serde(default = "default_shard_layout")]
-    #[default(ShardLayout::v0_single_shard())]
+    #[default(default_shard_layout())]
     pub shard_layout: ShardLayout,
     #[serde(default = "default_num_chunk_only_producer_seats")]
     #[default(300)]

core/chain-configs/src/test_genesis.rs

@@ -75,8 +75,7 @@ impl TestGenesisBuilder {
 
     pub fn epoch_config_mut(&mut self) -> &mut EpochConfig {
         if self.epoch_config.is_none() {
-            let mut epoch_config =
-                Genesis::test_epoch_config(1, ShardLayout::v0_single_shard(), 100);
+            let mut epoch_config = Genesis::test_epoch_config(1, ShardLayout::single_shard(), 100);
             epoch_config.block_producer_kickout_threshold = 0;
             epoch_config.chunk_producer_kickout_threshold = 0;
             epoch_config.chunk_validator_only_kickout_threshold = 0;
@@ -121,7 +120,7 @@ impl TestGenesisBuilder {
     }
 
     pub fn shard_layout_single(&mut self) -> &mut Self {
-        self.epoch_config_mut().shard_layout = ShardLayout::v0_single_shard();
+        self.epoch_config_mut().shard_layout = ShardLayout::single_shard();
         self
     }
 

core/chain-configs/src/test_utils.rs

@@ -138,7 +138,7 @@ impl Genesis {
             accounts,
             num_validator_seats,
             vec![num_validator_seats],
-            ShardLayout::v0_single_shard(),
+            ShardLayout::single_shard(),
         )
     }
 
@@ -154,7 +154,7 @@ impl Genesis {
             accounts,
             num_validator_seats,
             num_validator_seats_per_shard,
-            ShardLayout::v0(num_shards, 0),
+            ShardLayout::multi_shard(num_shards, 0),
         )
     }
 
@@ -169,7 +169,7 @@ impl Genesis {
             accounts,
             num_validator_seats,
             num_validator_seats_per_shard,
-            ShardLayout::v0(num_shards, 1),
+            ShardLayout::multi_shard(num_shards, 1),
         )
     }
 }

core/primitives/src/shard_layout.rs

@@ -164,17 +164,7 @@ impl ShardLayoutV1 {
 }
 
 /// Making the shard ids non-contiguous.
-#[derive(
-    BorshSerialize,
-    BorshDeserialize,
-    serde::Serialize,
-    serde::Deserialize,
-    Clone,
-    Debug,
-    PartialEq,
-    Eq,
-    ProtocolSchema,
-)]
+#[derive(BorshSerialize, BorshDeserialize, Clone, Debug, PartialEq, Eq, ProtocolSchema)]
 pub struct ShardLayoutV2 {
     /// The boundary accounts are the accounts on boundaries between shards.
     /// Each shard contains a range of accounts from one boundary account to
@@ -212,6 +202,114 @@ pub struct ShardLayoutV2 {
     version: ShardVersion,
 }
 
+/// Counterpart to `ShardLayoutV2` composed of maps with string keys to aid
+/// serde serialization.
+#[derive(serde::Serialize, serde::Deserialize)]
+struct SerdeShardLayoutV2 {
+    boundary_accounts: Vec<AccountId>,
+    shard_ids: Vec<ShardId>,
+    id_to_index_map: BTreeMap<String, ShardIndex>,
+    index_to_id_map: BTreeMap<String, ShardId>,
+    shards_split_map: Option<BTreeMap<String, Vec<ShardId>>>,
+    shards_parent_map: Option<BTreeMap<String, ShardId>>,
+    version: ShardVersion,
+}
+
+impl From<&ShardLayoutV2> for SerdeShardLayoutV2 {
+    fn from(layout: &ShardLayoutV2) -> Self {
+        let id_to_index_map =
+            layout.id_to_index_map.iter().map(|(k, v)| (k.to_string(), *v)).collect();
+
+        let index_to_id_map =
+            layout.index_to_id_map.iter().map(|(k, v)| (k.to_string(), *v)).collect();
+
+        let shards_split_map = layout
+            .shards_split_map
+            .as_ref()
+            .map(|map| map.iter().map(|(k, v)| (k.to_string(), v.clone())).collect());
+
+        let shards_parent_map = layout
+            .shards_parent_map
+            .as_ref()
+            .map(|map| map.iter().map(|(k, v)| (k.to_string(), *v)).collect());
+
+        Self {
+            boundary_accounts: layout.boundary_accounts.clone(),
+            shard_ids: layout.shard_ids.clone(),
+            id_to_index_map,
+            index_to_id_map,
+            shards_split_map,
+            shards_parent_map,
+            version: layout.version,
+        }
+    }
+}
+
+impl TryFrom<SerdeShardLayoutV2> for ShardLayoutV2 {
+    type Error = Box<dyn std::error::Error + Send + Sync>;
+
+    fn try_from(layout: SerdeShardLayoutV2) -> Result<Self, Self::Error> {
+        let id_to_index_map = layout
+            .id_to_index_map
+            .into_iter()
+            .map(|(k, v)| Ok((k.parse::<u64>()?.into(), v)))
+            .collect::<Result<_, Self::Error>>()?;
+
+        let index_to_id_map = layout
+            .index_to_id_map
+            .into_iter()
+            .map(|(k, v)| Ok((k.parse()?, v)))
+            .collect::<Result<_, Self::Error>>()?;
+
+        let shards_split_map = layout
+            .shards_split_map
+            .map(|map| {
+                map.into_iter()
+                    .map(|(k, v)| Ok((k.parse::<u64>()?.into(), v)))
+                    .collect::<Result<_, Self::Error>>()
+            })
+            .transpose()?;
+
+        let shards_parent_map = layout
+            .shards_parent_map
+            .map(|map| {
+                map.into_iter()
+                    .map(|(k, v)| Ok((k.parse::<u64>()?.into(), v)))
+                    .collect::<Result<_, Self::Error>>()
+            })
+            .transpose()?;
+
+        Ok(Self {
+            boundary_accounts: layout.boundary_accounts,
+            shard_ids: layout.shard_ids,
+            id_to_index_map,
+            index_to_id_map,
+            shards_split_map,
+            shards_parent_map,
+            version: layout.version,
+        })
+    }
+}
+
+impl serde::Serialize for ShardLayoutV2 {
+    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
+    where
+        S: serde::Serializer,
+    {
+        SerdeShardLayoutV2::from(self).serialize(serializer)
+    }
+}
+
+impl<'de> serde::Deserialize<'de> for ShardLayoutV2 {
+    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
+    where
+        D: serde::Deserializer<'de>,
+    {
+        let serde_layout = SerdeShardLayoutV2::deserialize(deserializer)?;
+        ShardLayoutV2::try_from(serde_layout).map_err(serde::de::Error::custom)
+    }
+}
+
 impl ShardLayoutV2 {
     pub fn account_id_to_shard_id(&self, account_id: &AccountId) -> ShardId {
         // TODO(resharding) - This could be optimized.
@@ -241,9 +339,34 @@ pub enum ShardLayoutError {
 }
 
 impl ShardLayout {
-    /* Some constructors */
-    pub fn v0_single_shard() -> Self {
-        Self::v0(1, 0)
+    /// Handy constructor for a single-shard layout, mostly for test purposes
+    pub fn single_shard() -> Self {
+        Self::multi_shard(1, 0)
+    }
+
+    /// Can be used to construct a multi-shard layout, mostly for test purposes
+    pub fn multi_shard(num_shards: NumShards, version: ShardVersion) -> Self {
+        assert!(num_shards > 0, "at least 1 shard is required");
+
+        let boundary_accounts = (0..num_shards - 1)
+            .map(|i| format!("shard{}.test.near", i).parse().unwrap())
+            .collect::<Vec<AccountId>>();
+        let shard_ids = (0..num_shards).map(ShardId::new).collect::<Vec<ShardId>>();
+        let (id_to_index_map, index_to_id_map) = shard_ids
+            .iter()
+            .enumerate()
+            .map(|(i, &shard_id)| ((shard_id, i), (i, shard_id)))
+            .unzip();
+
+        Self::V2(ShardLayoutV2 {
+            boundary_accounts,
+            shard_ids,
+            id_to_index_map,
+            index_to_id_map,
+            shards_split_map: None,
+            shards_parent_map: None,
+            version,
+        })
     }
 
     /// Return a V0 Shardlayout

integration-tests/src/tests/client/features/stateless_validation.rs

@@ -354,8 +354,10 @@ fn test_protocol_upgrade_81() {
     assert_eq!(config.chunk_producer_kickout_threshold, 90);
 }
 
+// TODO(wacban) should not panic when V2 is ready
 /// Test that Client rejects ChunkStateWitnesses with invalid shard_id
 #[test]
+#[should_panic(expected = "no entry found for key")]
 fn test_chunk_state_witness_bad_shard_id() {
     init_integration_logger();