Merge old and new lists of child keys when swapping hotkey

pallets/subtensor/src/lib.rs

@@ -208,6 +208,12 @@ pub mod pallet {
         T::InitialMinDelegateTake::get()
     }
 
+    #[pallet::type_value]
+    /// Default limit for number of childkeys on one hotkey
+    pub fn DefaultChildKeyLimit<T: Config>() -> u16 {
+        5
+    }
+
     #[pallet::type_value]
     /// Default minimum childkey take.
     pub fn DefaultMinChildKeyTake<T: Config>() -> u16 {

pallets/subtensor/src/staking/set_children.rs

@@ -1,4 +1,5 @@
 use super::*;
+use frame_support::traits::Get;
 
 impl<T: Config> Pallet<T> {
     /// ---- The implementation for the extrinsic do_set_child_singular: Sets a single child.
@@ -96,7 +97,10 @@ impl<T: Config> Pallet<T> {
         );
 
         // --- 4.1. Ensure that the number of children does not exceed 5.
-        ensure!(children.len() <= 5, Error::<T>::TooManyChildren);
+        ensure!(
+            children.len() <= DefaultChildKeyLimit::<T>::get() as usize,
+            Error::<T>::TooManyChildren
+        );
 
         // --- 5. Ensure that each child is not the hotkey.
         for (_, child_i) in &children {

pallets/subtensor/src/swap/swap_hotkey.rs

@@ -1,6 +1,7 @@
 use super::*;
 use frame_support::weights::Weight;
 use sp_core::Get;
+use substrate_fixed::types::I96F32;
 
 impl<T: Config> Pallet<T> {
     /// Swaps the hotkey of a coldkey account.
@@ -82,7 +83,7 @@ impl<T: Config> Pallet<T> {
         Self::burn_tokens(actual_burn_amount);
 
         // 14. Perform the hotkey swap
-        let _ = Self::perform_hotkey_swap(old_hotkey, new_hotkey, &coldkey, &mut weight);
+        Self::perform_hotkey_swap(old_hotkey, new_hotkey, &coldkey, &mut weight)?;
 
         // 15. Update the last transaction block for the coldkey
         Self::set_last_tx_block(&coldkey, block);
@@ -321,11 +322,80 @@ impl<T: Config> Pallet<T> {
         // ChildKeys( parent, netuid ) --> Vec<(proportion,child)> -- the child keys of the parent.
         for netuid in Self::get_all_subnet_netuids() {
             // Get the children of the old hotkey for this subnet
-            let my_children: Vec<(u64, T::AccountId)> = ChildKeys::<T>::get(old_hotkey, netuid);
+            let old_hotkey_children: Vec<(u64, T::AccountId)> =
+                ChildKeys::<T>::get(old_hotkey, netuid);
+            // Read old children of the new hotkey
+            let new_hotkey_children: Vec<(u64, T::AccountId)> =
+                ChildKeys::<T>::get(new_hotkey, netuid);
+            // Merged childkey list (initialize with old hotkey children)
+            let mut merged_childkey_list: Vec<(u64, T::AccountId)> =
+                ChildKeys::<T>::get(old_hotkey, netuid);
+
+            // Calculate total of proportions between two lists being merged (old and new)
+            let proportion_total: I96F32 = I96F32::from_num(
+                old_hotkey_children
+                    .iter()
+                    .map(|(proportion, _)| *proportion as u128)
+                    .sum::<u128>(),
+            )
+            .saturating_add(I96F32::from_num(
+                new_hotkey_children
+                    .iter()
+                    .map(|(proportion, _)| *proportion as u128)
+                    .sum::<u128>(),
+            ));
+
+            // Merge old and new child lists. Do not use HashSet because we only have 5 keys per list,
+            // and iteration is more efficient.
+            new_hotkey_children.iter().for_each(|new_hotkey_child| {
+                if !merged_childkey_list
+                    .iter()
+                    .any(|merged_hotkey_child| merged_hotkey_child.1 == new_hotkey_child.1)
+                {
+                    merged_childkey_list.push(new_hotkey_child.clone());
+                }
+            });
+
+            if merged_childkey_list.len() > DefaultChildKeyLimit::<T>::get() as usize {
+                return Err(Error::<T>::TooManyChildren.into());
+            }
+
+            // Update proportions in the merged list
+            // Re-normalize child proportions if total of proportions exceeds 1.0
+            // (because two child lists were merged and we can't go over 1.0 total proportion)
+            let proportion_normalization_coeff: I96F32 = if proportion_total > u64::MAX {
+                I96F32::from_num(u64::MAX).saturating_div(proportion_total)
+            } else {
+                I96F32::from_num(1)
+            };
+            merged_childkey_list.iter_mut().for_each(|child| {
+                let proportion_from_old_hotkey = old_hotkey_children
+                    .iter()
+                    .find(|&old_hotkey_child| old_hotkey_child.1 == child.1)
+                    .map(|old_hotkey_child| old_hotkey_child.0)
+                    .unwrap_or(0_u64);
+                let proportion_from_new_hotkey = new_hotkey_children
+                    .iter()
+                    .find(|&new_hotkey_child| new_hotkey_child.1 == child.1)
+                    .map(|new_hotkey_child| new_hotkey_child.0)
+                    .unwrap_or(0_u64);
+
+                let mut normalized_proportion: I96F32 =
+                    I96F32::from_num(proportion_from_old_hotkey)
+                        .saturating_mul(proportion_normalization_coeff);
+                normalized_proportion = normalized_proportion.saturating_add(
+                    I96F32::from_num(proportion_from_new_hotkey)
+                        .saturating_mul(proportion_normalization_coeff),
+                );
+
+                *child = (normalized_proportion.to_num::<u64>(), child.1.clone());
+            });
+
             // Remove the old hotkey's child entries
             ChildKeys::<T>::remove(old_hotkey, netuid);
-            // Insert the same child entries for the new hotkey
-            ChildKeys::<T>::insert(new_hotkey, netuid, my_children);
+
+            // Insert the new hotkey's child entries
+            ChildKeys::<T>::insert(new_hotkey, netuid, merged_childkey_list);
         }
 
         // 13. Swap ParentKeys.

pallets/subtensor/tests/swap_hotkey.rs

@@ -985,6 +985,210 @@ fn test_swap_child_keys() {
     });
 }
 
+#[test]
+fn test_swap_existing_child_keys() {
+    new_test_ext(1).execute_with(|| {
+        let old_hotkey = U256::from(1);
+        let new_hotkey = U256::from(2);
+        let coldkey = U256::from(3);
+        let netuid = 0u16;
+        let children1 = vec![(100u64, U256::from(4)), (200u64, U256::from(5))];
+        let children2 = vec![(100u64, U256::from(6)), (200u64, U256::from(7))];
+        let mut weight = Weight::zero();
+
+        // Initialize ChildKeys for old_hotkey
+        add_network(netuid, 1, 0);
+        ChildKeys::<Test>::insert(old_hotkey, netuid, children1.clone());
+        ChildKeys::<Test>::insert(new_hotkey, netuid, children2.clone());
+
+        // Perform the swap
+        SubtensorModule::perform_hotkey_swap(&old_hotkey, &new_hotkey, &coldkey, &mut weight);
+
+        // Verify the swap
+        assert_eq!(
+            ChildKeys::<Test>::get(new_hotkey, netuid),
+            [
+                (100u64, U256::from(4)),
+                (200u64, U256::from(5)),
+                (100u64, U256::from(6)),
+                (200u64, U256::from(7)),
+            ]
+        );
+        assert!(ChildKeys::<Test>::get(old_hotkey, netuid).is_empty());
+    });
+}
+
+#[test]
+fn test_swap_existing_child_keys_too_many_children() {
+    new_test_ext(1).execute_with(|| {
+        let old_hotkey = U256::from(1);
+        let new_hotkey = U256::from(2);
+        let coldkey = U256::from(3);
+        let netuid = 0u16;
+        let children1 = vec![
+            (100u64, U256::from(4)),
+            (200u64, U256::from(5)),
+            (200u64, U256::from(6)),
+        ];
+        let children2 = vec![
+            (100u64, U256::from(7)),
+            (200u64, U256::from(8)),
+            (200u64, U256::from(9)),
+        ];
+        let mut weight = Weight::zero();
+
+        // Initialize ChildKeys for old_hotkey
+        add_network(netuid, 1, 0);
+        ChildKeys::<Test>::insert(old_hotkey, netuid, children1.clone());
+        ChildKeys::<Test>::insert(new_hotkey, netuid, children2.clone());
+
+        // Perform the swap
+        assert_err!(
+            SubtensorModule::perform_hotkey_swap(&old_hotkey, &new_hotkey, &coldkey, &mut weight),
+            Error::<Test>::TooManyChildren,
+        );
+
+        // Verify no change
+        assert_eq!(ChildKeys::<Test>::get(old_hotkey, netuid), children1);
+        assert_eq!(ChildKeys::<Test>::get(new_hotkey, netuid), children2);
+    });
+}
+
+#[test]
+fn test_swap_existing_child_keys_too_many_children_extrinsic() {
+    new_test_ext(1).execute_with(|| {
+        let old_hotkey = U256::from(1);
+        let new_hotkey = U256::from(2);
+        let coldkey = U256::from(3);
+        let netuid = 1u16;
+        let swap_cost = 1_000_000_000u64 * 2;
+        let children1 = vec![
+            (100u64, U256::from(4)),
+            (200u64, U256::from(5)),
+            (200u64, U256::from(6)),
+        ];
+        let children2 = vec![
+            (100u64, U256::from(7)),
+            (200u64, U256::from(8)),
+            (200u64, U256::from(9)),
+        ];
+
+        // Initialize ChildKeys for old_hotkey
+        add_network(netuid, 1, 0);
+        register_ok_neuron(netuid, old_hotkey, coldkey, 0);
+        SubtensorModule::add_balance_to_coldkey_account(&coldkey, swap_cost);
+        ChildKeys::<Test>::insert(old_hotkey, netuid, children1.clone());
+        ChildKeys::<Test>::insert(new_hotkey, netuid, children2.clone());
+
+        // Perform the swap using high level function
+        assert_err!(
+            SubtensorModule::do_swap_hotkey(
+                <<Test as Config>::RuntimeOrigin>::signed(coldkey),
+                &old_hotkey,
+                &new_hotkey,
+            ),
+            Error::<Test>::TooManyChildren
+        );
+
+        // Verify no change
+        assert_eq!(ChildKeys::<Test>::get(old_hotkey, netuid), children1);
+        assert_eq!(ChildKeys::<Test>::get(new_hotkey, netuid), children2);
+    });
+}
+
+#[test]
+fn test_swap_existing_child_keys_proper_merge() {
+    new_test_ext(1).execute_with(|| {
+        let old_hotkey = U256::from(1);
+        let new_hotkey = U256::from(2);
+        let coldkey = U256::from(3);
+        let netuid = 0u16;
+        let children1 = vec![
+            (100u64, U256::from(4)),
+            (200u64, U256::from(5)),
+            (300u64, U256::from(6)),
+        ];
+        let children2 = vec![
+            (100u64, U256::from(6)),
+            (200u64, U256::from(7)),
+            (300u64, U256::from(8)),
+        ];
+        let mut weight = Weight::zero();
+
+        // Initialize ChildKeys for old_hotkey
+        add_network(netuid, 1, 0);
+        ChildKeys::<Test>::insert(old_hotkey, netuid, children1.clone());
+        ChildKeys::<Test>::insert(new_hotkey, netuid, children2.clone());
+
+        // Perform the swap
+        assert_ok!(SubtensorModule::perform_hotkey_swap(
+            &old_hotkey,
+            &new_hotkey,
+            &coldkey,
+            &mut weight
+        ),);
+
+        // Verify the swap
+        assert_eq!(
+            ChildKeys::<Test>::get(new_hotkey, netuid),
+            [
+                (100u64, U256::from(4)),
+                (200u64, U256::from(5)),
+                (400u64, U256::from(6)),
+                (200u64, U256::from(7)),
+                (300u64, U256::from(8)),
+            ]
+        );
+        assert!(ChildKeys::<Test>::get(old_hotkey, netuid).is_empty());
+    });
+}
+
+#[test]
+fn test_swap_existing_child_keys_proportion_overflow() {
+    new_test_ext(1).execute_with(|| {
+        let old_hotkey = U256::from(1);
+        let new_hotkey = U256::from(2);
+        let coldkey = U256::from(3);
+        let netuid = 0u16;
+        let children1 = vec![
+            (u64::MAX / 5, U256::from(4)),
+            (u64::MAX / 5 * 4, U256::from(5)),
+        ];
+        let children2 = vec![
+            (u64::MAX / 5 * 4, U256::from(5)),
+            (u64::MAX / 5, U256::from(6)),
+        ];
+        let mut weight = Weight::zero();
+
+        // Initialize ChildKeys for old_hotkey
+        add_network(netuid, 1, 0);
+        ChildKeys::<Test>::insert(old_hotkey, netuid, children1.clone());
+        ChildKeys::<Test>::insert(new_hotkey, netuid, children2.clone());
+
+        // Perform the swap
+        assert_ok!(SubtensorModule::perform_hotkey_swap(
+            &old_hotkey,
+            &new_hotkey,
+            &coldkey,
+            &mut weight
+        ),);
+
+        // Verify the swap
+        // key 4: max / 5 renormalized = max / 10
+        // key 5: 8 max / 5 renormalized = 4 max / 5
+        // key 6: max / 5 renormalized = max / 10
+        assert_eq!(
+            ChildKeys::<Test>::get(new_hotkey, netuid),
+            [
+                (u64::MAX / 10, U256::from(4)),
+                (u64::MAX / 5 * 4, U256::from(5)),
+                (u64::MAX / 10, U256::from(6)),
+            ]
+        );
+        assert!(ChildKeys::<Test>::get(old_hotkey, netuid).is_empty());
+    });
+}
+
 // SKIP_WASM_BUILD=1 RUST_LOG=debug cargo test --test swap_hotkey -- test_swap_parent_keys --exact --nocapture
 #[test]
 fn test_swap_parent_keys() {