Skip to content

Commit 47bd6e7

Browse files
sam0x17sam0x17
authored
Merge pull request #2112 from opentensor/fix/add-auto-stake-dest-coldkeys-map-for-hk-swap
Fix/add auto stake dest coldkeys map for hk swap
2 parents 969730d + 0a6ae62 commit 47bd6e7

File tree

11 files changed

+193
-4
lines changed

11 files changed

+193
-4
lines changed

‎pallets/subtensor/src/lib.rs‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1150,6 +1150,16 @@ pub mod pallet {
11501150
T::AccountId,
11511151
OptionQuery,
11521152
>;
1153+
#[pallet::storage] // --- DMAP ( hot, netuid )--> Vec<cold> | Returns a list of coldkeys that are autostaking to a hotkey.
1154+
pub type AutoStakeDestinationColdkeys<T: Config> = StorageDoubleMap<
1155+
_,
1156+
Blake2_128Concat,
1157+
T::AccountId,
1158+
Identity,
1159+
NetUid,
1160+
Vec<T::AccountId>,
1161+
ValueQuery,
1162+
>;
11531163

11541164
#[pallet::storage] // --- DMAP ( cold ) --> (block_expected, new_coldkey) | Maps coldkey to the block to swap at and new coldkey.
11551165
pub type ColdkeySwapScheduled<T: Config> = StorageMap<

‎pallets/subtensor/src/macros/dispatches.rs‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2313,9 +2313,20 @@ mod dispatches {
23132313
current_hotkey != hotkey,
23142314
Error::<T>::SameAutoStakeHotkeyAlreadySet
23152315
);
2316+
2317+
// Remove the coldkey from the old hotkey (if present)
2318+
AutoStakeDestinationColdkeys::<T>::mutate(current_hotkey.clone(), netuid, |v| {
2319+
v.retain(|c| c != &coldkey);
2320+
});
23162321
}
23172322

2323+
// Add the coldkey to the new hotkey (if not already present)
23182324
AutoStakeDestination::<T>::insert(coldkey.clone(), netuid, hotkey.clone());
2325+
AutoStakeDestinationColdkeys::<T>::mutate(hotkey.clone(), netuid, |v| {
2326+
if !v.contains(&coldkey) {
2327+
v.push(coldkey.clone());
2328+
}
2329+
});
23192330

23202331
Self::deposit_event(Event::AutoStakeDestinationSet {
23212332
coldkey,

‎pallets/subtensor/src/macros/hooks.rs‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -149,7 +149,9 @@ mod hooks {
149149
// Migrate subnet locked balances
150150
.saturating_add(migrations::migrate_subnet_locked::migrate_restore_subnet_locked::<T>())
151151
// Migrate subnet burn cost to 2500
152-
.saturating_add(migrations::migrate_network_lock_cost_2500::migrate_network_lock_cost_2500::<T>());
152+
.saturating_add(migrations::migrate_network_lock_cost_2500::migrate_network_lock_cost_2500::<T>())
153+
// Migrate AutoStakeDestinationColdkeys
154+
.saturating_add(migrations::migrate_auto_stake_destination::migrate_auto_stake_destination::<T>());
153155
weight
154156
}
155157

‎pallets/subtensor/src/migrations/migrate_auto_stake_destination.rs‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,11 @@ pub fn migrate_auto_stake_destination<T: Config>() -> Weight {
5757
continue;
5858
}
5959
AutoStakeDestination::<T>::insert(coldkey, netuid, hotkey.clone());
60+
AutoStakeDestinationColdkeys::<T>::mutate(hotkey.clone(), netuid, |v| {
61+
if !v.contains(coldkey) {
62+
v.push(coldkey.clone());
63+
}
64+
});
6065
}
6166

6267
old::AutoStakeDestination::<T>::remove(coldkey);

‎pallets/subtensor/src/swap/swap_coldkey.rs‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -159,7 +159,16 @@ impl<T: Config> Pallet<T> {
159159
if let Some(old_auto_stake_hotkey) = AutoStakeDestination::<T>::get(old_coldkey, netuid)
160160
{
161161
AutoStakeDestination::<T>::remove(old_coldkey, netuid);
162-
AutoStakeDestination::<T>::insert(new_coldkey, netuid, old_auto_stake_hotkey);
162+
AutoStakeDestination::<T>::insert(
163+
new_coldkey,
164+
netuid,
165+
old_auto_stake_hotkey.clone(),
166+
);
167+
AutoStakeDestinationColdkeys::<T>::mutate(old_auto_stake_hotkey, netuid, |v| {
168+
// Remove old/new coldkeys (avoid duplicates), then add the new one.
169+
v.retain(|c| *c != *old_coldkey && *c != *new_coldkey);
170+
v.push(new_coldkey.clone());
171+
});
163172
}
164173
}
165174

‎pallets/subtensor/src/swap/swap_hotkey.rs‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -524,6 +524,18 @@ impl<T: Config> Pallet<T> {
524524
}
525525
}
526526

527+
// 6.4 Swap AutoStakeDestination
528+
if let Ok(old_auto_stake_coldkeys) =
529+
AutoStakeDestinationColdkeys::<T>::try_get(old_hotkey, netuid)
530+
{
531+
// Move the vector from old hotkey to new hotkey.
532+
for coldkey in &old_auto_stake_coldkeys {
533+
AutoStakeDestination::<T>::insert(coldkey, netuid, new_hotkey);
534+
}
535+
AutoStakeDestinationColdkeys::<T>::remove(old_hotkey, netuid);
536+
AutoStakeDestinationColdkeys::<T>::insert(new_hotkey, netuid, old_auto_stake_coldkeys);
537+
}
538+
527539
// 7. Swap SubnetOwnerHotkey
528540
// SubnetOwnerHotkey( netuid ) --> hotkey -- the hotkey that is the owner of the subnet.
529541
if let Ok(old_subnet_owner_hotkey) = SubnetOwnerHotkey::<T>::try_get(netuid) {

‎pallets/subtensor/src/tests/auto_stake_hotkey.rs‎

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -110,3 +110,64 @@ fn test_set_coldkey_auto_stake_hotkey_same_hotkey_again() {
110110
);
111111
});
112112
}
113+
114+
#[test]
115+
fn test_set_coldkey_auto_stake_hotkey_change_hotkey() {
116+
new_test_ext(1).execute_with(|| {
117+
let subnet_owner_ck = U256::from(0);
118+
let subnet_owner_hk = U256::from(1);
119+
120+
let coldkey = U256::from(10);
121+
let hotkey = U256::from(11);
122+
let new_hotkey = U256::from(12);
123+
124+
Owner::<Test>::insert(hotkey, coldkey);
125+
OwnedHotkeys::<Test>::insert(coldkey, vec![hotkey]);
126+
127+
let netuid = add_dynamic_network(&subnet_owner_hk, &subnet_owner_ck);
128+
Uids::<Test>::insert(netuid, hotkey, 1);
129+
Uids::<Test>::insert(netuid, new_hotkey, 2);
130+
131+
// First call should succeed
132+
assert_ok!(SubtensorModule::set_coldkey_auto_stake_hotkey(
133+
RuntimeOrigin::signed(coldkey),
134+
netuid,
135+
hotkey,
136+
));
137+
138+
// Check maps
139+
assert_eq!(
140+
AutoStakeDestination::<Test>::get(coldkey, netuid),
141+
Some(hotkey)
142+
);
143+
assert_eq!(
144+
AutoStakeDestinationColdkeys::<Test>::get(hotkey, netuid),
145+
vec![coldkey]
146+
);
147+
assert_eq!(
148+
AutoStakeDestinationColdkeys::<Test>::get(new_hotkey, netuid),
149+
vec![]
150+
);
151+
152+
// Second call with new hotkey should succeed
153+
assert_ok!(SubtensorModule::set_coldkey_auto_stake_hotkey(
154+
RuntimeOrigin::signed(coldkey),
155+
netuid,
156+
new_hotkey,
157+
));
158+
159+
// Check maps again
160+
assert_eq!(
161+
AutoStakeDestination::<Test>::get(coldkey, netuid),
162+
Some(new_hotkey)
163+
);
164+
assert_eq!(
165+
AutoStakeDestinationColdkeys::<Test>::get(hotkey, netuid),
166+
vec![]
167+
);
168+
assert_eq!(
169+
AutoStakeDestinationColdkeys::<Test>::get(new_hotkey, netuid),
170+
vec![coldkey]
171+
);
172+
});
173+
}

‎pallets/subtensor/src/tests/migration.rs‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1581,8 +1581,17 @@ fn test_migrate_auto_stake_destination() {
15811581
AutoStakeDestination::<Test>::get(coldkey2, *netuid),
15821582
Some(hotkey2)
15831583
);
1584-
}
15851584

1585+
// Verify entry for AutoStakeDestinationColdkeys
1586+
assert_eq!(
1587+
AutoStakeDestinationColdkeys::<Test>::get(hotkey1, *netuid),
1588+
vec![coldkey1]
1589+
);
1590+
assert_eq!(
1591+
AutoStakeDestinationColdkeys::<Test>::get(hotkey2, *netuid),
1592+
vec![coldkey2]
1593+
);
1594+
}
15861595
}
15871596

15881597
// Verify old format entries are cleared

‎pallets/subtensor/src/tests/swap_coldkey.rs‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2619,3 +2619,37 @@ fn test_coldkey_in_swap_schedule_prevents_critical_calls() {
26192619
);
26202620
});
26212621
}
2622+
2623+
#[test]
2624+
fn test_swap_auto_stake_destination_coldkeys() {
2625+
new_test_ext(1).execute_with(|| {
2626+
let old_coldkey = U256::from(1);
2627+
let new_coldkey = U256::from(2);
2628+
let hotkey = U256::from(3);
2629+
let netuid = NetUid::from(1u16);
2630+
let coldkeys = vec![U256::from(4), U256::from(5), old_coldkey];
2631+
2632+
add_network(netuid, 1, 0);
2633+
AutoStakeDestinationColdkeys::<Test>::insert(hotkey, netuid, coldkeys.clone());
2634+
AutoStakeDestination::<Test>::insert(old_coldkey, netuid, hotkey);
2635+
2636+
let mut weight = Weight::zero();
2637+
assert_ok!(SubtensorModule::perform_swap_coldkey(
2638+
&old_coldkey,
2639+
&new_coldkey,
2640+
&mut weight
2641+
));
2642+
2643+
let new_coldkeys = AutoStakeDestinationColdkeys::<Test>::get(hotkey, netuid);
2644+
assert!(new_coldkeys.contains(&new_coldkey));
2645+
assert!(!new_coldkeys.contains(&old_coldkey));
2646+
assert_eq!(
2647+
AutoStakeDestination::<Test>::try_get(old_coldkey, netuid),
2648+
Err(())
2649+
);
2650+
assert_eq!(
2651+
AutoStakeDestination::<Test>::try_get(new_coldkey, netuid),
2652+
Ok(hotkey)
2653+
);
2654+
});
2655+
}

‎pallets/subtensor/src/tests/swap_hotkey.rs‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1465,3 +1465,39 @@ fn test_swap_hotkey_swap_rate_limits() {
14651465
);
14661466
});
14671467
}
1468+
1469+
#[test]
1470+
fn test_swap_auto_stake_destination_coldkeys() {
1471+
new_test_ext(1).execute_with(|| {
1472+
let old_hotkey = U256::from(1);
1473+
let new_hotkey = U256::from(2);
1474+
let coldkey = U256::from(3);
1475+
let netuid = NetUid::from(2u16); // Can't be root
1476+
let coldkeys = vec![U256::from(4), U256::from(5), coldkey];
1477+
let mut weight = Weight::zero();
1478+
1479+
// Initialize ChildKeys for old_hotkey
1480+
add_network(netuid, 1, 0);
1481+
AutoStakeDestinationColdkeys::<Test>::insert(old_hotkey, netuid, coldkeys.clone());
1482+
AutoStakeDestination::<Test>::insert(coldkey, netuid, old_hotkey);
1483+
1484+
// Perform the swap
1485+
SubtensorModule::perform_hotkey_swap_on_all_subnets(
1486+
&old_hotkey,
1487+
&new_hotkey,
1488+
&coldkey,
1489+
&mut weight,
1490+
);
1491+
1492+
// Verify the swap
1493+
assert_eq!(
1494+
AutoStakeDestinationColdkeys::<Test>::get(new_hotkey, netuid),
1495+
coldkeys
1496+
);
1497+
assert!(AutoStakeDestinationColdkeys::<Test>::get(old_hotkey, netuid).is_empty());
1498+
assert_eq!(
1499+
AutoStakeDestination::<Test>::get(coldkey, netuid),
1500+
Some(new_hotkey)
1501+
);
1502+
});
1503+
}

0 commit comments

Comments
 (0)