Convert a few collections to sets

Author: Bownairo

rs/protobuf/generator/src/lib.rs

@@ -243,7 +243,7 @@ fn build_registry_proto(def: &Path, out: &Path) {
 
     config.type_attribute(
         ".registry.hostos_version",
-        "#[derive(serde::Serialize, serde::Deserialize)]",
+        "#[derive(serde::Serialize, serde::Deserialize, Ord, PartialOrd, Eq)]",
     );
     config.type_attribute(
         ".registry.node_rewards.v2",

rs/protobuf/src/gen/registry/registry.hostos_version.v1.rs

@@ -5,7 +5,9 @@
 /// accepted by token holders. Nodes can then be upgraded to any of those
 /// versions. hostos_version_id commonly matches release_package_sha256_hex,
 /// and is used in the key for this record.
-#[derive(serde::Serialize, serde::Deserialize, Clone, PartialEq, ::prost::Message)]
+#[derive(
+    serde::Serialize, serde::Deserialize, Ord, PartialOrd, Eq, Clone, PartialEq, ::prost::Message,
+)]
 pub struct HostosVersionRecord {
     /// The URLs against which a HTTP GET request will return a release package
     /// that corresponds to this version.

rs/registry/canister/src/invariants/common.rs

@@ -1,5 +1,5 @@
 use std::{
-    collections::BTreeMap,
+    collections::{BTreeMap, BTreeSet},
     convert::TryFrom,
     error,
     fmt::{Display, Formatter, Result as FmtResult},
@@ -62,15 +62,15 @@ pub(crate) fn get_all_node_records(snapshot: &RegistrySnapshot) -> BTreeMap<Node
 /// Returns all replica version records in the snapshot.
 pub(crate) fn get_all_replica_version_records(
     snapshot: &RegistrySnapshot,
-) -> Vec<(String, ReplicaVersionRecord)> {
-    let mut replica_versions = Vec::new();
+) -> BTreeMap<String, ReplicaVersionRecord> {
+    let mut replica_versions = BTreeMap::new();
     for (k, v) in snapshot {
         if let Some(key) = str::from_utf8(k)
             .unwrap()
             .strip_prefix(REPLICA_VERSION_KEY_PREFIX)
         {
             let record = ReplicaVersionRecord::decode(v.as_slice()).unwrap();
-            replica_versions.push((key.to_owned(), record));
+            replica_versions.insert(key.to_owned(), record);
         }
     }
 
@@ -114,12 +114,12 @@ pub(crate) fn get_all_chain_key_signing_subnet_list_records(
 // Retrieve all HostOS version records
 pub(crate) fn get_all_hostos_version_records(
     snapshot: &RegistrySnapshot,
-) -> Vec<HostosVersionRecord> {
-    let mut result = Vec::new();
+) -> BTreeSet<HostosVersionRecord> {
+    let mut result = BTreeSet::new();
     for (k, v) in snapshot {
         if k.starts_with(HOSTOS_VERSION_KEY_PREFIX.as_bytes()) {
             let hostos_version_record = HostosVersionRecord::decode(v.as_slice()).unwrap();
-            result.push(hostos_version_record);
+            result.insert(hostos_version_record);
         }
     }
 
@@ -147,8 +147,8 @@ pub(crate) fn get_api_boundary_node_records_from_snapshot(
 /// Returns an all api boundary node ids record from the registry snapshot.
 pub(crate) fn get_api_boundary_node_ids_from_snapshot(
     snapshot: &RegistrySnapshot,
-) -> Result<Vec<NodeId>, InvariantCheckError> {
-    let api_bn_ids: Result<Vec<NodeId>, InvariantCheckError> = snapshot
+) -> Result<BTreeSet<NodeId>, InvariantCheckError> {
+    snapshot
         .keys()
         .cloned()
         .map(|key| {
@@ -157,15 +157,13 @@ pub(crate) fn get_api_boundary_node_ids_from_snapshot(
                 source: None,
             })
         })
-        .collect::<Result<Vec<String>, InvariantCheckError>>()
+        .collect::<Result<BTreeSet<String>, InvariantCheckError>>()
         .map(|keys| {
             keys.into_iter()
                 .filter_map(|key_str| get_api_boundary_node_record_node_id(&key_str))
                 .map(NodeId::from)
                 .collect()
-        });
-
-    api_bn_ids
+        })
 }
 
 /// Returns node record from the snapshot corresponding to a key.
@@ -185,7 +183,7 @@ pub(crate) fn get_node_record_from_snapshot(
         .transpose()
 }
 
-pub(crate) fn get_subnet_ids_from_snapshot(snapshot: &RegistrySnapshot) -> Vec<SubnetId> {
+pub(crate) fn get_subnet_ids_from_snapshot(snapshot: &RegistrySnapshot) -> BTreeSet<SubnetId> {
     get_value_from_snapshot::<SubnetListRecord>(snapshot, make_subnet_list_record_key())
         .map(|r| {
             r.subnets

rs/registry/canister/src/invariants/replica_version.rs

@@ -42,16 +42,9 @@ pub(crate) fn check_replica_version_invariants(
     }
     versions_in_use.append(&mut get_all_api_boundary_node_versions(snapshot));
 
-    let elected_version_ids = get_all_replica_version_records(snapshot)
-        .into_iter()
-        .map(|(k, _)| k);
-
-    let num_elected = elected_version_ids.len();
-    let elected_set = BTreeSet::from_iter(elected_version_ids);
-    assert!(
-        elected_set.len() == num_elected,
-        "A version was elected multiple times."
-    );
+    let elected_set: BTreeSet<_> = get_all_replica_version_records(snapshot)
+        .into_keys()
+        .collect();
     assert!(
         elected_set.is_superset(&versions_in_use),
         "Using a version that isn't elected. Elected versions: {elected_set:?}, in use: {versions_in_use:?}."