diff --git a/chain-extensions/src/lib.rs b/chain-extensions/src/lib.rs index 3fe67c687b..e2f9a56d81 100644 --- a/chain-extensions/src/lib.rs +++ b/chain-extensions/src/lib.rs @@ -19,6 +19,7 @@ use pallet_subtensor_proxy as pallet_proxy; use pallet_subtensor_proxy::WeightInfo; use sp_runtime::{DispatchError, Weight, traits::StaticLookup}; use sp_std::marker::PhantomData; +use sp_std::vec; use substrate_fixed::types::U64F64; use subtensor_runtime_common::{AlphaBalance, NetUid, ProxyType, TaoBalance}; use subtensor_swap_interface::SwapHandler; @@ -964,6 +965,37 @@ where } } } + FunctionId::ClaimRootWithHotkeyV1 => { + let hotkey: T::AccountId = env + .read_as() + .map_err(|_| DispatchError::Other("Failed to decode input parameters"))?; + + let weight = pallet_subtensor::Pallet::::root_claim_declared_weight(); + env.charge_weight(weight)?; + + if !pallet_subtensor::Pallet::::root_claim_fits_declared_budget( + core::slice::from_ref(&hotkey), + ) { + return Ok(RetVal::Converging(Output::RuntimeError as u32)); + } + + let caller = env.caller(); + let call_result = + pallet_subtensor::Pallet::::do_root_claim(caller.clone(), vec![hotkey]); + + match call_result { + Ok(outcome) => { + pallet_subtensor::Pallet::::maybe_add_coldkey_index(&caller); + env.write_output(&outcome.tao.encode()) + .map_err(|_| DispatchError::Other("Failed to write output"))?; + Ok(RetVal::Converging(Output::Success as u32)) + } + Err(e) => { + let error_code = Output::from(e) as u32; + Ok(RetVal::Converging(error_code)) + } + } + } } } } diff --git a/chain-extensions/src/tests.rs b/chain-extensions/src/tests.rs index a207ca6690..032b1b1777 100644 --- a/chain-extensions/src/tests.rs +++ b/chain-extensions/src/tests.rs @@ -12,7 +12,7 @@ use pallet_subtensor::weights::WeightInfo as SubtensorWeightInfo; use sp_core::Get; use sp_core::U256; use sp_runtime::DispatchError; -use substrate_fixed::types::U64F64; +use substrate_fixed::types::{I96F32, U64F64}; use subtensor_runtime_common::{AlphaBalance, NetUid, TaoBalance, Token}; use subtensor_swap_interface::SwapHandler; @@ -1358,6 +1358,271 @@ fn assert_success(ret: RetVal) { } } +fn expected_claim_root_with_hotkey_weight() -> Weight { + let max = pallet_subtensor::MAX_ROOT_CLAIM_WORK; + let full = + <::WeightInfo as SubtensorWeightInfo>::claim_root( + max, + ); + let scan = + <::WeightInfo as SubtensorWeightInfo>::claim_root_scan( + max, + ); + full.saturating_add(scan) +} + +#[test] +fn claim_root_with_hotkey_noop_returns_zero() { + mock::new_test_ext(1).execute_with(|| { + let coldkey = U256::from(61001); + let hotkey = U256::from(61002); + + let expected_weight = expected_claim_root_with_hotkey_weight(); + + let mut env = MockEnv::new(FunctionId::ClaimRootWithHotkeyV1, coldkey, hotkey.encode()) + .with_expected_weight(expected_weight); + + let ret = SubtensorChainExtension::::dispatch(&mut env).unwrap(); + assert_success(ret); + assert_eq!(env.charged_weight(), Some(expected_weight)); + + let tao: u64 = Decode::decode(&mut &env.output()[..]).unwrap(); + assert_eq!(tao, 0); + }); +} + +#[test] +fn claim_root_with_hotkey_noop_indexes_fresh_coldkey() { + mock::new_test_ext(1).execute_with(|| { + let coldkey = U256::from(61211); + let hotkey = U256::from(61212); + + assert!(!pallet_subtensor::StakingColdkeys::::contains_key(coldkey)); + let num_before = pallet_subtensor::NumStakingColdkeys::::get(); + + let expected_weight = expected_claim_root_with_hotkey_weight(); + + let mut env = MockEnv::new(FunctionId::ClaimRootWithHotkeyV1, coldkey, hotkey.encode()) + .with_expected_weight(expected_weight); + + let ret = SubtensorChainExtension::::dispatch(&mut env).unwrap(); + assert_success(ret); + assert_eq!(env.charged_weight(), Some(expected_weight)); + + let tao: u64 = Decode::decode(&mut &env.output()[..]).unwrap(); + assert_eq!(tao, 0); + + assert!(pallet_subtensor::StakingColdkeys::::contains_key(coldkey)); + assert_eq!( + pallet_subtensor::NumStakingColdkeys::::get(), + num_before.saturating_add(1) + ); + let idx = pallet_subtensor::StakingColdkeys::::get(coldkey).unwrap(); + assert_eq!( + pallet_subtensor::StakingColdkeysByIndex::::get(idx), + Some(coldkey) + ); + }); +} + +#[test] +fn claim_root_with_hotkey_repeat_claim_preserves_index() { + mock::new_test_ext(1).execute_with(|| { + let coldkey = U256::from(61221); + let hotkey = U256::from(61222); + + pallet_subtensor::Pallet::::maybe_add_coldkey_index(&coldkey); + let num_before = pallet_subtensor::NumStakingColdkeys::::get(); + + let expected_weight = expected_claim_root_with_hotkey_weight(); + + let mut env = MockEnv::new(FunctionId::ClaimRootWithHotkeyV1, coldkey, hotkey.encode()) + .with_expected_weight(expected_weight); + + let ret = SubtensorChainExtension::::dispatch(&mut env).unwrap(); + assert_success(ret); + + let tao: u64 = Decode::decode(&mut &env.output()[..]).unwrap(); + assert_eq!(tao, 0); + + assert_eq!( + pallet_subtensor::NumStakingColdkeys::::get(), + num_before + ); + let idx = pallet_subtensor::StakingColdkeys::::get(coldkey).unwrap(); + assert_eq!( + pallet_subtensor::StakingColdkeysByIndex::::get(idx), + Some(coldkey) + ); + }); +} + +#[test] +fn claim_root_with_hotkey_rejects_basket_above_envelope() { + mock::new_test_ext(1).execute_with(|| { + let coldkey = U256::from(61101); + let hotkey = U256::from(61102); + + // Seed the validator's basket with one escrow holding per netuid, 257 rows + // > MAX_ROOT_CLAIM_WORK (256): 1 hotkey unit + 257 rows exceed the fixed + // admission budget, so the signed path refuses RootClaimTooHeavy here. + let escrow = pallet_subtensor::Pallet::::get_beta_escrow_account_id(); + for i in 0..=pallet_subtensor::MAX_ROOT_CLAIM_WORK { + pallet_subtensor::Pallet::::increase_stake_for_hotkey_and_coldkey_on_subnet( + &hotkey, + &escrow, + NetUid::from(i as u16), + 1u64.into(), + ); + } + assert!( + !pallet_subtensor::Pallet::::root_claim_fits_declared_budget( + core::slice::from_ref(&hotkey) + ) + ); + let num_before = pallet_subtensor::NumStakingColdkeys::::get(); + + let expected_weight = expected_claim_root_with_hotkey_weight(); + let mut env = MockEnv::new(FunctionId::ClaimRootWithHotkeyV1, coldkey, hotkey.encode()) + .with_expected_weight(expected_weight); + + let ret = SubtensorChainExtension::::dispatch(&mut env).unwrap(); + match ret { + RetVal::Converging(code) => assert_eq!(code, Output::RuntimeError as u32), + _ => panic!("expected converging error code"), + } + + assert_eq!(env.charged_weight(), Some(expected_weight)); + + // Rejection must not mutate the index. + assert!(!pallet_subtensor::StakingColdkeys::::contains_key(coldkey)); + assert_eq!( + pallet_subtensor::NumStakingColdkeys::::get(), + num_before + ); + }); +} + +#[test] +fn claim_root_with_hotkey_rejects_basket_at_budget_boundary() { + mock::new_test_ext(1).execute_with(|| { + let coldkey = U256::from(61111); + let hotkey = U256::from(61112); + + let escrow = pallet_subtensor::Pallet::::get_beta_escrow_account_id(); + for i in 0..pallet_subtensor::MAX_ROOT_CLAIM_WORK { + pallet_subtensor::Pallet::::increase_stake_for_hotkey_and_coldkey_on_subnet( + &hotkey, + &escrow, + NetUid::from(i as u16), + 1u64.into(), + ); + } + assert!( + !pallet_subtensor::Pallet::::root_claim_fits_declared_budget( + core::slice::from_ref(&hotkey) + ) + ); + let num_before = pallet_subtensor::NumStakingColdkeys::::get(); + + let expected_weight = expected_claim_root_with_hotkey_weight(); + let mut env = MockEnv::new(FunctionId::ClaimRootWithHotkeyV1, coldkey, hotkey.encode()) + .with_expected_weight(expected_weight); + + let ret = SubtensorChainExtension::::dispatch(&mut env).unwrap(); + match ret { + RetVal::Converging(code) => assert_eq!(code, Output::RuntimeError as u32), + _ => panic!("expected converging error code"), + } + assert_eq!(env.charged_weight(), Some(expected_weight)); + + assert!(!pallet_subtensor::StakingColdkeys::::contains_key(coldkey)); + assert_eq!( + pallet_subtensor::NumStakingColdkeys::::get(), + num_before + ); + }); +} + +#[test] +fn claim_root_with_hotkey_admits_basket_under_budget() { + mock::new_test_ext(1).execute_with(|| { + let coldkey = U256::from(61121); + let hotkey = U256::from(61122); + + let escrow = pallet_subtensor::Pallet::::get_beta_escrow_account_id(); + for i in 0..pallet_subtensor::MAX_ROOT_CLAIM_WORK.saturating_sub(1) { + pallet_subtensor::Pallet::::increase_stake_for_hotkey_and_coldkey_on_subnet( + &hotkey, + &escrow, + NetUid::from(i as u16), + 1u64.into(), + ); + } + assert!( + pallet_subtensor::Pallet::::root_claim_fits_declared_budget( + core::slice::from_ref(&hotkey) + ) + ); + + let expected_weight = expected_claim_root_with_hotkey_weight(); + let mut env = MockEnv::new(FunctionId::ClaimRootWithHotkeyV1, coldkey, hotkey.encode()) + .with_expected_weight(expected_weight); + + let ret = SubtensorChainExtension::::dispatch(&mut env).unwrap(); + assert_success(ret); + assert_eq!(env.charged_weight(), Some(expected_weight)); + + let tao: u64 = Decode::decode(&mut &env.output()[..]).unwrap(); + assert_eq!(tao, 0); + + assert!(pallet_subtensor::StakingColdkeys::::contains_key(coldkey)); + }); +} + +#[test] +fn claim_root_with_hotkey_payout_indexes_coldkey() { + mock::new_test_ext(1).execute_with(|| { + let coldkey = U256::from(61201); + let hotkey = U256::from(61202); + + pallet_subtensor::Pallet::::increase_stake_for_hotkey_and_coldkey_on_subnet( + &hotkey, + &coldkey, + NetUid::ROOT, + 1u64.into(), + ); + let escrow = pallet_subtensor::Pallet::::get_beta_escrow_account_id(); + pallet_subtensor::Pallet::::increase_stake_for_hotkey_and_coldkey_on_subnet( + &hotkey, + &escrow, + NetUid::ROOT, + 1_000_000u64.into(), + ); + pallet_subtensor::BasketShares::::insert(hotkey, 1u64); + pallet_subtensor::BasketRate::::insert(hotkey, I96F32::from_num(1)); + pallet_subtensor::RootClaimableThreshold::::insert( + NetUid::ROOT, + I96F32::from_num(0), + ); + + assert!(!pallet_subtensor::StakingColdkeys::::contains_key(coldkey)); + + let expected_weight = expected_claim_root_with_hotkey_weight(); + + let mut env = MockEnv::new(FunctionId::ClaimRootWithHotkeyV1, coldkey, hotkey.encode()) + .with_expected_weight(expected_weight); + + let ret = SubtensorChainExtension::::dispatch(&mut env).unwrap(); + assert_success(ret); + + let tao: u64 = Decode::decode(&mut &env.output()[..]).unwrap(); + assert_eq!(tao, 1_000_000); + + assert!(pallet_subtensor::StakingColdkeys::::contains_key(coldkey)); + }); +} + #[test] fn add_stake_recycle_rollback_on_recycle_failure() { mock::new_test_ext(1).execute_with(|| { diff --git a/chain-extensions/src/types.rs b/chain-extensions/src/types.rs index f98a4a1591..a91df529f7 100644 --- a/chain-extensions/src/types.rs +++ b/chain-extensions/src/types.rs @@ -46,6 +46,7 @@ pub enum FunctionId { GetStakeAvailabilityV1 = 36, MoveStakeLimitV1 = 37, CallerMoveStakeLimitV1 = 38, + ClaimRootWithHotkeyV1 = 39, } #[freeze_struct("5dc33d60abed5c08")] @@ -192,11 +193,12 @@ mod function_id_tests { assert_eq!(FunctionId::GetStakeAvailabilityV1 as u16, 36); assert_eq!(FunctionId::MoveStakeLimitV1 as u16, 37); assert_eq!(FunctionId::CallerMoveStakeLimitV1 as u16, 38); + assert_eq!(FunctionId::ClaimRootWithHotkeyV1 as u16, 39); } #[test] fn caller_ids_roundtrip_try_from_primitive() { - for id in 16u16..=38u16 { + for id in 16u16..=39u16 { let v = FunctionId::try_from_primitive(id) .unwrap_or_else(|_| panic!("try_from_primitive failed for {id}")); assert_eq!(v as u16, id); diff --git a/pallets/subtensor/src/staking/claim_root.rs b/pallets/subtensor/src/staking/claim_root.rs index 83930c0db7..8ab383c35b 100644 --- a/pallets/subtensor/src/staking/claim_root.rs +++ b/pallets/subtensor/src/staking/claim_root.rs @@ -902,7 +902,7 @@ impl Pallet { /// Pre-dispatch weight for both independently bounded dimensions: full claim work and /// scan-only work. - pub(crate) fn root_claim_declared_weight() -> Weight { + pub fn root_claim_declared_weight() -> Weight { let limit = Self::root_claim_declared_work(); ::WeightInfo::claim_root(limit).saturating_add( ::WeightInfo::claim_root_scan(limit), @@ -929,7 +929,7 @@ impl Pallet { /// True when the hotkeys plus the basket storage rows the claim will scan fit the fixed /// admission envelope. Count raw Alpha/AlphaV2 rows so legacy duplicates and malformed /// zero rows are charged conservatively, and stop as soon as the bound is exceeded. - pub(crate) fn root_claim_fits_declared_budget(hotkeys: &[T::AccountId]) -> bool { + pub fn root_claim_fits_declared_budget(hotkeys: &[T::AccountId]) -> bool { let budget = Self::root_claim_declared_work(); let mut work = u32::try_from(hotkeys.len()).unwrap_or(u32::MAX); if work > budget {