feat: Move key family helpers to top level

rs/registry/canister/src/common/key_family.rs

@@ -0,0 +1,81 @@
+use std::default::Default;
+
+use crate::{registry::Registry, storage::with_chunks};
+use ic_registry_canister_chunkify::decode_high_capacity_registry_value;
+use ic_registry_transport::pb::v1::HighCapacityRegistryValue;
+
+/// Similar to `get_key_family` on the `RegistryClient`, return a list of
+/// tuples, (ID, value).  This strips the prefix from the key and returns the
+/// value as a decoded struct.
+///
+/// This function must return keys in order of their bytes, which should
+/// also be the same order as the string representations.
+pub(crate) fn get_key_family<T: prost::Message + Default>(
+    registry: &Registry,
+    prefix: &str,
+) -> Vec<(String, T)> {
+    get_key_family_iter(registry, prefix).collect()
+}
+
+/// This function must return keys in order of their bytes, which should
+/// also be the same order as the string representations.
+pub(crate) fn get_key_family_iter<'a, T: prost::Message + Default>(
+    registry: &'a Registry,
+    prefix: &'a str,
+) -> impl Iterator<Item = (String, T)> + 'a {
+    get_key_family_iter_at_version(registry, prefix, registry.latest_version())
+}
+
+/// This function must return keys in order of their bytes, which should
+/// also be the same order as the string representations.
+pub(crate) fn get_key_family_iter_at_version<'a, T: prost::Message + Default>(
+    registry: &'a Registry,
+    prefix: &'a str,
+    version: u64,
+) -> impl Iterator<Item = (String, T)> + 'a {
+    get_key_family_raw_iter_at_version(registry, prefix, version).filter_map(|(id, value)| {
+        let latest_value: Option<T> =
+            with_chunks(|chunks| decode_high_capacity_registry_value::<T, _>(value, chunks));
+
+        let latest_value = latest_value?;
+
+        Some((id, latest_value))
+    })
+}
+
+/// This function must return keys in order of their bytes, which should
+/// also be the same order as the string representations.
+pub(crate) fn get_key_family_raw_iter_at_version<'a>(
+    registry: &'a Registry,
+    prefix: &'a str,
+    version: u64,
+) -> impl Iterator<Item = (String, &'a HighCapacityRegistryValue)> + 'a {
+    let prefix_bytes = prefix.as_bytes();
+    let start = prefix_bytes.to_vec();
+
+    // Note, using the 'store' which is a BTreeMap is what guarantees the order of keys.
+    registry
+        .store
+        .range(start..)
+        .take_while(|(k, _)| k.starts_with(prefix_bytes))
+        .filter_map(move |(key, values)| {
+            let latest_value: &HighCapacityRegistryValue =
+                values.iter().rev().find(|value| value.version <= version)?;
+
+            if !latest_value.is_present() {
+                return None; // Deleted or otherwise empty value.
+            }
+
+            let id = key
+                .strip_prefix(prefix_bytes)
+                .and_then(|v| std::str::from_utf8(v).ok())
+                .unwrap()
+                .to_string();
+
+            Some((id, latest_value))
+        })
+}
+
+#[path = "key_family_tests.rs"]
+#[cfg(test)]
+mod tests;

rs/registry/canister/src/mutations/node_management/common_tests.rs → rs/registry/canister/src/common/key_family_tests.rs

@@ -1,5 +1,10 @@
 use super::*;
-use ic_registry_transport::{pb::v1::LargeValueChunkKeys, upsert};
+
+use ic_registry_transport::{
+    pb::v1::{LargeValueChunkKeys, RegistryMutation},
+    upsert,
+};
+use prost::Message;
 
 fn new_upsert(key: &str, value: &str) -> RegistryMutation {
     let value = LargeValueChunkKeys {

rs/registry/canister/src/common/mod.rs

@@ -1,4 +1,5 @@
 pub const LOG_PREFIX: &str = "[Registry] ";
 
+pub mod key_family;
 #[cfg(test)]
 pub mod test_helpers;

rs/registry/canister/src/get_node_providers_monthly_xdr_rewards.rs

@@ -1,4 +1,4 @@
-use crate::mutations::node_management::common::get_key_family_iter_at_version;
+use crate::common::key_family::get_key_family_iter_at_version;
 use crate::{pb::v1::NodeProvidersMonthlyXdrRewards, registry::Registry};
 use ic_protobuf::registry::{
     dc::v1::DataCenterRecord, node_operator::v1::NodeOperatorRecord,

rs/registry/canister/src/mutations/api_boundary_node.rs

@@ -1,7 +1,8 @@
 use std::str::FromStr;
 
 use crate::{
-    common::LOG_PREFIX, mutations::node_management::common::get_key_family_iter, registry::Registry,
+    common::{LOG_PREFIX, key_family::get_key_family_iter},
+    registry::Registry,
 };
 
 use ic_base_types::NodeId;

rs/registry/canister/src/mutations/do_remove_node_operators.rs

@@ -1,5 +1,6 @@
 use crate::{
-    common::LOG_PREFIX, mutations::node_management::common::get_key_family_iter, registry::Registry,
+    common::{LOG_PREFIX, key_family::get_key_family_iter},
+    registry::Registry,
 };
 
 #[cfg(target_arch = "wasm32")]

rs/registry/canister/src/mutations/node_management/common.rs

@@ -1,21 +1,21 @@
-use std::{default::Default, str::FromStr};
+use std::str::FromStr;
 
-use crate::{common::LOG_PREFIX, registry::Registry, storage::with_chunks};
+use crate::common::key_family::get_key_family;
+use crate::{common::LOG_PREFIX, registry::Registry};
 use ic_base_types::{NodeId, PrincipalId, SubnetId};
 use ic_crypto_node_key_validation::ValidNodePublicKeys;
 use ic_protobuf::registry::node::v1::NodeRewardType;
 use ic_protobuf::registry::{
     node::v1::NodeRecord, node_operator::v1::NodeOperatorRecord, subnet::v1::SubnetListRecord,
 };
-use ic_registry_canister_chunkify::decode_high_capacity_registry_value;
 use ic_registry_keys::{
     FirewallRulesScope, NODE_RECORD_KEY_PREFIX, make_crypto_node_key, make_crypto_tls_cert_key,
     make_firewall_rules_record_key, make_node_operator_record_key, make_node_record_key,
     make_subnet_list_record_key,
 };
 use ic_registry_transport::{
     delete, insert,
-    pb::v1::{HighCapacityRegistryValue, RegistryMutation, RegistryValue},
+    pb::v1::{RegistryMutation, RegistryValue},
     update,
 };
 use ic_types::crypto::KeyPurpose;
@@ -287,79 +287,3 @@ pub fn node_exists_with_ipv4(registry: &Registry, ipv4_addr: &str) -> bool {
         })
         .is_some()
 }
-
-/// Similar to `get_key_family` on the `RegistryClient`, return a list of
-/// tuples, (ID, value).  This strips the prefix from the key and returns the
-/// value as a decoded struct.
-///
-/// This function must return keys in order of their bytes, which should
-/// also be the same order as the string representations.
-pub(crate) fn get_key_family<T: prost::Message + Default>(
-    registry: &Registry,
-    prefix: &str,
-) -> Vec<(String, T)> {
-    get_key_family_iter(registry, prefix).collect()
-}
-
-/// This function must return keys in order of their bytes, which should
-/// also be the same order as the string representations.
-pub(crate) fn get_key_family_iter<'a, T: prost::Message + Default>(
-    registry: &'a Registry,
-    prefix: &'a str,
-) -> impl Iterator<Item = (String, T)> + 'a {
-    get_key_family_iter_at_version(registry, prefix, registry.latest_version())
-}
-
-/// This function must return keys in order of their bytes, which should
-/// also be the same order as the string representations.
-pub(crate) fn get_key_family_iter_at_version<'a, T: prost::Message + Default>(
-    registry: &'a Registry,
-    prefix: &'a str,
-    version: u64,
-) -> impl Iterator<Item = (String, T)> + 'a {
-    get_key_family_raw_iter_at_version(registry, prefix, version).filter_map(|(id, value)| {
-        let latest_value: Option<T> =
-            with_chunks(|chunks| decode_high_capacity_registry_value::<T, _>(value, chunks));
-
-        let latest_value = latest_value?;
-
-        Some((id, latest_value))
-    })
-}
-
-/// This function must return keys in order of their bytes, which should
-/// also be the same order as the string representations.
-pub(crate) fn get_key_family_raw_iter_at_version<'a>(
-    registry: &'a Registry,
-    prefix: &'a str,
-    version: u64,
-) -> impl Iterator<Item = (String, &'a HighCapacityRegistryValue)> + 'a {
-    let prefix_bytes = prefix.as_bytes();
-    let start = prefix_bytes.to_vec();
-
-    // Note, using the 'store' which is a BTreeMap is what guarantees the order of keys.
-    registry
-        .store
-        .range(start..)
-        .take_while(|(k, _)| k.starts_with(prefix_bytes))
-        .filter_map(move |(key, values)| {
-            let latest_value: &HighCapacityRegistryValue =
-                values.iter().rev().find(|value| value.version <= version)?;
-
-            if !latest_value.is_present() {
-                return None; // Deleted or otherwise empty value.
-            }
-
-            let id = key
-                .strip_prefix(prefix_bytes)
-                .and_then(|v| std::str::from_utf8(v).ok())
-                .unwrap()
-                .to_string();
-
-            Some((id, latest_value))
-        })
-}
-
-#[path = "common_tests.rs"]
-#[cfg(test)]
-mod tests;

rs/registry/canister/src/mutations/routing_table.rs

@@ -1,7 +1,7 @@
 use crate::{
-    common::LOG_PREFIX,
-    mutations::node_management::common::{
-        get_key_family_iter_at_version, get_key_family_raw_iter_at_version,
+    common::{
+        LOG_PREFIX,
+        key_family::{get_key_family_iter_at_version, get_key_family_raw_iter_at_version},
     },
     pb::v1::SubnetForCanister,
     registry::Registry,
@@ -459,7 +459,7 @@ mod tests {
     use crate::common::test_helpers::invariant_compliant_registry;
 
     use super::*;
-    use crate::mutations::node_management::common::get_key_family;
+    use crate::common::key_family::get_key_family;
     use assert_matches::assert_matches;
     use ic_base_types::CanisterId;
     use ic_registry_keys::CANISTER_RANGES_PREFIX;