diff --git a/pallets/subtensor/src/swap/swap_hotkey.rs b/pallets/subtensor/src/swap/swap_hotkey.rs index efa6c53f8..b16c180c4 100644 --- a/pallets/subtensor/src/swap/swap_hotkey.rs +++ b/pallets/subtensor/src/swap/swap_hotkey.rs @@ -337,7 +337,20 @@ impl Pallet { // Remove the old hotkey's child entries ChildKeys::::remove(old_hotkey, netuid); // Insert the same child entries for the new hotkey - ChildKeys::::insert(new_hotkey, netuid, my_children); + ChildKeys::::insert(new_hotkey, netuid, my_children.clone()); + for (_, child_key_i) in my_children { + // For each child, update their parent list + let mut child_parents: Vec<(u64, T::AccountId)> = + ParentKeys::::get(child_key_i.clone(), netuid); + for parent in child_parents.iter_mut() { + // If the parent is the old hotkey, replace it with the new hotkey + if parent.1 == *old_hotkey { + parent.1 = new_hotkey.clone(); + } + } + // Update the child's parent list + ParentKeys::::insert(child_key_i, netuid, child_parents); + } } // 13. Swap ParentKeys. diff --git a/pallets/subtensor/tests/swap_hotkey.rs b/pallets/subtensor/tests/swap_hotkey.rs index bf5ecb301..53524c72c 100644 --- a/pallets/subtensor/tests/swap_hotkey.rs +++ b/pallets/subtensor/tests/swap_hotkey.rs @@ -1148,3 +1148,100 @@ fn test_swap_complex_parent_child_structure() { ); }); } + +#[test] +fn test_swap_parent_hotkey_childkey_maps() { + new_test_ext(1).execute_with(|| { + let netuid: u16 = 1; + let parent_old = U256::from(1); + let coldkey = U256::from(2); + let child = U256::from(3); + let parent_new = U256::from(4); + add_network(netuid, 0, 0); + SubtensorModule::create_account_if_non_existent(&coldkey, &parent_old); + + // Set child and verify state maps + assert_ok!(SubtensorModule::do_set_children( + RuntimeOrigin::signed(coldkey), + parent_old, + netuid, + vec![(u64::MAX, child)] + )); + assert_eq!( + ParentKeys::::get(child, netuid), + vec![(u64::MAX, parent_old)] + ); + assert_eq!( + ChildKeys::::get(parent_old, netuid), + vec![(u64::MAX, child)] + ); + + // Swap + let mut weight = Weight::zero(); + assert_ok!(SubtensorModule::perform_hotkey_swap( + &parent_old, + &parent_new, + &coldkey, + &mut weight + )); + + // Verify parent and child keys updates + assert_eq!( + ParentKeys::::get(child, netuid), + vec![(u64::MAX, parent_new)] + ); + assert_eq!( + ChildKeys::::get(parent_new, netuid), + vec![(u64::MAX, child)] + ); + }) +} + +#[test] +fn test_swap_child_hotkey_childkey_maps() { + new_test_ext(1).execute_with(|| { + let netuid: u16 = 1; + let parent = U256::from(1); + let coldkey = U256::from(2); + let child_old = U256::from(3); + let child_new = U256::from(4); + add_network(netuid, 0, 0); + SubtensorModule::create_account_if_non_existent(&coldkey, &child_old); + SubtensorModule::create_account_if_non_existent(&coldkey, &parent); + + // Set child and verify state maps + assert_ok!(SubtensorModule::do_set_children( + RuntimeOrigin::signed(coldkey), + parent, + netuid, + vec![(u64::MAX, child_old)] + )); + assert_eq!( + ParentKeys::::get(child_old, netuid), + vec![(u64::MAX, parent)] + ); + assert_eq!( + ChildKeys::::get(parent, netuid), + vec![(u64::MAX, child_old)] + ); + + // Swap + let mut weight = Weight::zero(); + assert_ok!(SubtensorModule::perform_hotkey_swap( + &child_old, + &child_new, + &coldkey, + &mut weight + )); + + // Verify parent and child keys updates + assert_eq!( + ParentKeys::::get(child_new, netuid), + vec![(u64::MAX, parent)] + ); + assert_eq!( + ChildKeys::::get(parent, netuid), + vec![(u64::MAX, child_new)] + ); + }) +}