diff --git a/pallets/subtensor/src/swap/swap_hotkey.rs b/pallets/subtensor/src/swap/swap_hotkey.rs index fb3c33e4d8..366494d117 100644 --- a/pallets/subtensor/src/swap/swap_hotkey.rs +++ b/pallets/subtensor/src/swap/swap_hotkey.rs @@ -307,6 +307,41 @@ impl Pallet { weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 1)); } + // 11. Swap ChildKeys. + // 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::::get(old_hotkey, netuid); + // 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); + } + + // 12. Swap ParentKeys. + // ParentKeys( child, netuid ) --> Vec<(proportion,parent)> -- the parent keys of the child. + for netuid in Self::get_all_subnet_netuids() { + // Get the parents of the old hotkey for this subnet + let parents: Vec<(u64, T::AccountId)> = ParentKeys::::get(old_hotkey, netuid); + // Remove the old hotkey's parent entries + ParentKeys::::remove(old_hotkey, netuid); + // Insert the same parent entries for the new hotkey + ParentKeys::::insert(new_hotkey, netuid, parents.clone()); + for (_, parent_key_i) in parents { + // For each parent, update their children list + let mut parent_children: Vec<(u64, T::AccountId)> = + ChildKeys::::get(parent_key_i.clone(), netuid); + for child in parent_children.iter_mut() { + // If the child is the old hotkey, replace it with the new hotkey + if child.1 == *old_hotkey { + child.1 = new_hotkey.clone(); + } + } + // Update the parent's children list + ChildKeys::::insert(parent_key_i, netuid, parent_children); + } + } + // Return successful after swapping all the relevant terms. Ok(()) } diff --git a/pallets/subtensor/tests/swap_hotkey.rs b/pallets/subtensor/tests/swap_hotkey.rs index c6a05f2b6f..bff738b86d 100644 --- a/pallets/subtensor/tests/swap_hotkey.rs +++ b/pallets/subtensor/tests/swap_hotkey.rs @@ -959,3 +959,159 @@ fn test_swap_hotkey_error_cases() { assert_eq!(Balances::free_balance(coldkey), initial_balance - swap_cost); }); } + +// SKIP_WASM_BUILD=1 RUST_LOG=debug cargo test --test swap_hotkey -- test_swap_child_keys --exact --nocapture +#[test] +fn test_swap_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 children = vec![(100u64, U256::from(4)), (200u64, U256::from(5))]; + let mut weight = Weight::zero(); + + // Initialize ChildKeys for old_hotkey + add_network(netuid, 1, 0); + ChildKeys::::insert(old_hotkey, netuid, children.clone()); + + // Perform the swap + SubtensorModule::perform_hotkey_swap(&old_hotkey, &new_hotkey, &coldkey, &mut weight); + + // Verify the swap + assert_eq!(ChildKeys::::get(new_hotkey, netuid), children); + assert!(ChildKeys::::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() { + 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 parents = vec![(100u64, U256::from(4)), (200u64, U256::from(5))]; + let mut weight = Weight::zero(); + + // Initialize ParentKeys for old_hotkey + add_network(netuid, 1, 0); + ParentKeys::::insert(old_hotkey, netuid, parents.clone()); + + // Initialize ChildKeys for parent + ChildKeys::::insert(U256::from(4), netuid, vec![(100u64, old_hotkey)]); + ChildKeys::::insert(U256::from(5), netuid, vec![(200u64, old_hotkey)]); + + // Perform the swap + SubtensorModule::perform_hotkey_swap(&old_hotkey, &new_hotkey, &coldkey, &mut weight); + + // Verify ParentKeys swap + assert_eq!(ParentKeys::::get(new_hotkey, netuid), parents); + assert!(ParentKeys::::get(old_hotkey, netuid).is_empty()); + + // Verify ChildKeys update for parents + assert_eq!( + ChildKeys::::get(U256::from(4), netuid), + vec![(100u64, new_hotkey)] + ); + assert_eq!( + ChildKeys::::get(U256::from(5), netuid), + vec![(200u64, new_hotkey)] + ); + }); +} + +// SKIP_WASM_BUILD=1 RUST_LOG=debug cargo test --test swap_hotkey -- test_swap_multiple_subnets --exact --nocapture +#[test] +fn test_swap_multiple_subnets() { + new_test_ext(1).execute_with(|| { + let old_hotkey = U256::from(1); + let new_hotkey = U256::from(2); + let coldkey = U256::from(3); + let netuid1 = 0u16; + let netuid2 = 1u16; + let children1 = vec![(100u64, U256::from(4)), (200u64, U256::from(5))]; + let children2 = vec![(300u64, U256::from(6))]; + let mut weight = Weight::zero(); + + add_network(netuid1, 1, 0); + add_network(netuid2, 1, 0); + + // Initialize ChildKeys for old_hotkey in multiple subnets + ChildKeys::::insert(old_hotkey, netuid1, children1.clone()); + ChildKeys::::insert(old_hotkey, netuid2, children2.clone()); + + // Perform the swap + SubtensorModule::perform_hotkey_swap(&old_hotkey, &new_hotkey, &coldkey, &mut weight); + + // Verify the swap for both subnets + assert_eq!(ChildKeys::::get(new_hotkey, netuid1), children1); + assert_eq!(ChildKeys::::get(new_hotkey, netuid2), children2); + assert!(ChildKeys::::get(old_hotkey, netuid1).is_empty()); + assert!(ChildKeys::::get(old_hotkey, netuid2).is_empty()); + }); +} + +// SKIP_WASM_BUILD=1 RUST_LOG=debug cargo test --test swap_hotkey -- test_swap_complex_parent_child_structure --exact --nocapture +#[test] +fn test_swap_complex_parent_child_structure() { + 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 parent1 = U256::from(4); + let parent2 = U256::from(5); + let child1 = U256::from(6); + let child2 = U256::from(7); + let mut weight = Weight::zero(); + + add_network(netuid, 1, 0); + + // Set up complex parent-child structure + ParentKeys::::insert( + old_hotkey, + netuid, + vec![(100u64, parent1), (200u64, parent2)], + ); + ChildKeys::::insert(old_hotkey, netuid, vec![(300u64, child1), (400u64, child2)]); + ChildKeys::::insert( + parent1, + netuid, + vec![(100u64, old_hotkey), (500u64, U256::from(8))], + ); + ChildKeys::::insert( + parent2, + netuid, + vec![(200u64, old_hotkey), (600u64, U256::from(9))], + ); + + // Perform the swap + SubtensorModule::perform_hotkey_swap(&old_hotkey, &new_hotkey, &coldkey, &mut weight); + + // Verify ParentKeys swap + assert_eq!( + ParentKeys::::get(new_hotkey, netuid), + vec![(100u64, parent1), (200u64, parent2)] + ); + assert!(ParentKeys::::get(old_hotkey, netuid).is_empty()); + + // Verify ChildKeys swap + assert_eq!( + ChildKeys::::get(new_hotkey, netuid), + vec![(300u64, child1), (400u64, child2)] + ); + assert!(ChildKeys::::get(old_hotkey, netuid).is_empty()); + + // Verify parent's ChildKeys update + assert_eq!( + ChildKeys::::get(parent1, netuid), + vec![(100u64, new_hotkey), (500u64, U256::from(8))] + ); + assert_eq!( + ChildKeys::::get(parent2, netuid), + vec![(200u64, new_hotkey), (600u64, U256::from(9))] + ); + }); +}