From dda99da52fc85e65e59216226d623cb3ec2bc553 Mon Sep 17 00:00:00 2001 From: Landyn Date: Wed, 1 Apr 2026 11:07:56 -0500 Subject: [PATCH 1/9] Burn and Recycle Chain Exts --- chain-extensions/src/lib.rs | 193 ++++++++++++++++- chain-extensions/src/tests.rs | 359 ++++++++++++++++++++++++++++++++ chain-extensions/src/types.rs | 10 + contract-tests/bittensor/lib.rs | 84 ++++++++ docs/wasm-contracts.md | 7 + 5 files changed, 652 insertions(+), 1 deletion(-) diff --git a/chain-extensions/src/lib.rs b/chain-extensions/src/lib.rs index 14ea23d9c8..0cd1523cac 100644 --- a/chain-extensions/src/lib.rs +++ b/chain-extensions/src/lib.rs @@ -66,7 +66,10 @@ where Env: SubtensorExtensionEnv, <::Lookup as StaticLookup>::Source: From<::AccountId>, { - let func_id: FunctionId = env.func_id().try_into().map_err(|_| { + let raw_func_id = env.func_id(); + log::info!("chain_ext: dispatch called with raw func_id={raw_func_id}"); + let func_id: FunctionId = raw_func_id.try_into().map_err(|_| { + log::error!("chain_ext: invalid func_id={raw_func_id}, not in FunctionId enum"); DispatchError::Other( "Invalid function id - does not correspond to any registered function", ) @@ -523,6 +526,194 @@ where Ok(RetVal::Converging(Output::Success as u32)) } + FunctionId::RecycleAlphaV1 => { + let weight = Weight::from_parts(113_400_000, 0) + .saturating_add(T::DbWeight::get().reads(10)) + .saturating_add(T::DbWeight::get().writes(4)); + + env.charge_weight(weight)?; + + let (hotkey, amount, netuid): (T::AccountId, AlphaBalance, NetUid) = + env.read_as()?; + + let caller = env.caller(); + + let alpha_available = + pallet_subtensor::Pallet::::get_stake_for_hotkey_and_coldkey_on_subnet( + &hotkey, &caller, netuid, + ); + let actual_amount = amount.min(alpha_available); + + let call_result = pallet_subtensor::Pallet::::recycle_alpha( + RawOrigin::Signed(caller).into(), + hotkey, + actual_amount, + netuid, + ); + + match call_result { + Ok(_) => { + env.write_output(&actual_amount.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)) + } + } + } + FunctionId::BurnAlphaV1 => { + let weight = Weight::from_parts(112_200_000, 0) + .saturating_add(T::DbWeight::get().reads(10)) + .saturating_add(T::DbWeight::get().writes(3)); + + env.charge_weight(weight)?; + + let (hotkey, amount, netuid): (T::AccountId, AlphaBalance, NetUid) = + env.read_as()?; + + let caller = env.caller(); + + let alpha_available = + pallet_subtensor::Pallet::::get_stake_for_hotkey_and_coldkey_on_subnet( + &hotkey, &caller, netuid, + ); + let actual_amount = amount.min(alpha_available); + + let call_result = pallet_subtensor::Pallet::::burn_alpha( + RawOrigin::Signed(caller).into(), + hotkey, + actual_amount, + netuid, + ); + + match call_result { + Ok(_) => { + env.write_output(&actual_amount.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)) + } + } + } + FunctionId::AddStakeRecycleV1 => { + log::info!("chain_ext: AddStakeRecycleV1 called"); + + let weight = Weight::from_parts(454_200_000, 0) + .saturating_add(T::DbWeight::get().reads(33)) + .saturating_add(T::DbWeight::get().writes(19)); + + if let Err(e) = env.charge_weight(weight) { + log::error!("chain_ext: AddStakeRecycleV1 charge_weight failed: {e:?}"); + return Err(e); + } + + let input: Result<(T::AccountId, NetUid, TaoBalance), _> = env.read_as(); + let (hotkey, netuid, tao_amount) = match input { + Ok(v) => v, + Err(e) => { + log::error!("chain_ext: AddStakeRecycleV1 read_as failed: {e:?}"); + return Err(e); + } + }; + + let caller = env.caller(); + log::info!( + "chain_ext: AddStakeRecycleV1 caller={caller:?} hotkey={hotkey:?} netuid={netuid:?} tao={tao_amount:?}" + ); + + let alpha = pallet_subtensor::Pallet::::do_add_stake( + RawOrigin::Signed(caller.clone()).into(), + hotkey.clone(), + netuid, + tao_amount, + ); + + match alpha { + Ok(alpha) => { + log::info!( + "chain_ext: AddStakeRecycleV1 do_add_stake ok, alpha={alpha:?}" + ); + let recycle_result = pallet_subtensor::Pallet::::recycle_alpha( + RawOrigin::Signed(caller).into(), + hotkey, + alpha, + netuid, + ); + + match recycle_result { + Ok(_) => { + log::info!("chain_ext: AddStakeRecycleV1 recycle ok"); + env.write_output(&alpha.encode()) + .map_err(|_| DispatchError::Other("Failed to write output"))?; + Ok(RetVal::Converging(Output::Success as u32)) + } + Err(e) => { + log::error!( + "chain_ext: AddStakeRecycleV1 recycle failed: {e:?}" + ); + let error_code = Output::from(e) as u32; + Ok(RetVal::Converging(error_code)) + } + } + } + Err(e) => { + log::error!("chain_ext: AddStakeRecycleV1 do_add_stake failed: {e:?}"); + let error_code = Output::from(e) as u32; + Ok(RetVal::Converging(error_code)) + } + } + } + FunctionId::AddStakeBurnV1 => { + let weight = Weight::from_parts(453_000_000, 0) + .saturating_add(T::DbWeight::get().reads(33)) + .saturating_add(T::DbWeight::get().writes(18)); + + env.charge_weight(weight)?; + + let (hotkey, netuid, tao_amount): (T::AccountId, NetUid, TaoBalance) = + env.read_as()?; + + let caller = env.caller(); + + let alpha = pallet_subtensor::Pallet::::do_add_stake( + RawOrigin::Signed(caller.clone()).into(), + hotkey.clone(), + netuid, + tao_amount, + ); + + match alpha { + Ok(alpha) => { + let burn_result = pallet_subtensor::Pallet::::burn_alpha( + RawOrigin::Signed(caller).into(), + hotkey, + alpha, + netuid, + ); + + match burn_result { + Ok(_) => { + env.write_output(&alpha.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)) + } + } + } + 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 b8956e8659..284e852f24 100644 --- a/chain-extensions/src/tests.rs +++ b/chain-extensions/src/tests.rs @@ -726,6 +726,365 @@ fn remove_proxy_success_removes_proxy_relationship() { }); } +#[test] +fn recycle_alpha_success_reduces_stake_and_returns_actual_amount() { + mock::new_test_ext(1).execute_with(|| { + let owner_hotkey = U256::from(9001); + let owner_coldkey = U256::from(9002); + let coldkey = U256::from(9101); + let hotkey = U256::from(9102); + let min_stake = DefaultMinStake::::get(); + let stake_amount_raw = min_stake.to_u64().saturating_mul(200); + + let netuid = mock::add_dynamic_network(&owner_hotkey, &owner_coldkey); + mock::setup_reserves( + netuid, + TaoBalance::from(130_000_000_000_u64), + AlphaBalance::from(110_000_000_000_u64), + ); + + mock::register_ok_neuron(netuid, hotkey, coldkey, 0); + + pallet_subtensor::Pallet::::add_balance_to_coldkey_account( + &coldkey, + TaoBalance::from(stake_amount_raw.saturating_add(1_000_000_000)), + ); + + assert_ok!(pallet_subtensor::Pallet::::add_stake( + RawOrigin::Signed(coldkey).into(), + hotkey, + netuid, + stake_amount_raw.into(), + )); + + let alpha_before = + pallet_subtensor::Pallet::::get_stake_for_hotkey_and_coldkey_on_subnet( + &hotkey, &coldkey, netuid, + ); + assert!(alpha_before > AlphaBalance::ZERO); + + let alpha_out_before = pallet_subtensor::SubnetAlphaOut::::get(netuid); + + let recycle_amount: AlphaBalance = (alpha_before.to_u64() / 2).into(); + + let expected_weight = Weight::from_parts(113_400_000, 0) + .saturating_add(::DbWeight::get().reads(10)) + .saturating_add(::DbWeight::get().writes(4)); + + let mut env = MockEnv::new( + FunctionId::RecycleAlphaV1, + coldkey, + (hotkey, recycle_amount, netuid).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 returned_amount = AlphaBalance::decode(&mut env.output()).unwrap(); + assert_eq!(returned_amount, recycle_amount); + + let alpha_after = + pallet_subtensor::Pallet::::get_stake_for_hotkey_and_coldkey_on_subnet( + &hotkey, &coldkey, netuid, + ); + assert!(alpha_after < alpha_before); + + let alpha_out_after = pallet_subtensor::SubnetAlphaOut::::get(netuid); + assert!(alpha_out_after < alpha_out_before); + }); +} + +#[test] +fn recycle_alpha_on_root_subnet_returns_error() { + mock::new_test_ext(1).execute_with(|| { + let coldkey = U256::from(9201); + let hotkey = U256::from(9202); + + pallet_subtensor::Owner::::insert(hotkey, coldkey); + + let expected_weight = Weight::from_parts(113_400_000, 0) + .saturating_add(::DbWeight::get().reads(10)) + .saturating_add(::DbWeight::get().writes(4)); + + let mut env = MockEnv::new( + FunctionId::RecycleAlphaV1, + coldkey, + (hotkey, AlphaBalance::from(1_000u64), NetUid::ROOT).encode(), + ) + .with_expected_weight(expected_weight); + + let ret = SubtensorChainExtension::::dispatch(&mut env).unwrap(); + match ret { + RetVal::Converging(code) => { + assert_ne!( + code, + Output::Success as u32, + "should not succeed on root subnet" + ) + } + _ => panic!("unexpected return value"), + } + }); +} + +#[test] +fn burn_alpha_success_reduces_stake_and_returns_actual_amount() { + mock::new_test_ext(1).execute_with(|| { + let owner_hotkey = U256::from(9301); + let owner_coldkey = U256::from(9302); + let coldkey = U256::from(9401); + let hotkey = U256::from(9402); + let min_stake = DefaultMinStake::::get(); + let stake_amount_raw = min_stake.to_u64().saturating_mul(200); + + let netuid = mock::add_dynamic_network(&owner_hotkey, &owner_coldkey); + mock::setup_reserves( + netuid, + TaoBalance::from(130_000_000_000_u64), + AlphaBalance::from(110_000_000_000_u64), + ); + + mock::register_ok_neuron(netuid, hotkey, coldkey, 0); + + pallet_subtensor::Pallet::::add_balance_to_coldkey_account( + &coldkey, + TaoBalance::from(stake_amount_raw.saturating_add(1_000_000_000)), + ); + + assert_ok!(pallet_subtensor::Pallet::::add_stake( + RawOrigin::Signed(coldkey).into(), + hotkey, + netuid, + stake_amount_raw.into(), + )); + + let alpha_before = + pallet_subtensor::Pallet::::get_stake_for_hotkey_and_coldkey_on_subnet( + &hotkey, &coldkey, netuid, + ); + assert!(alpha_before > AlphaBalance::ZERO); + + let alpha_out_before = pallet_subtensor::SubnetAlphaOut::::get(netuid); + + let burn_amount: AlphaBalance = (alpha_before.to_u64() / 2).into(); + + let expected_weight = Weight::from_parts(112_200_000, 0) + .saturating_add(::DbWeight::get().reads(10)) + .saturating_add(::DbWeight::get().writes(3)); + + let mut env = MockEnv::new( + FunctionId::BurnAlphaV1, + coldkey, + (hotkey, burn_amount, netuid).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 returned_amount = AlphaBalance::decode(&mut env.output()).unwrap(); + assert_eq!(returned_amount, burn_amount); + + let alpha_after = + pallet_subtensor::Pallet::::get_stake_for_hotkey_and_coldkey_on_subnet( + &hotkey, &coldkey, netuid, + ); + assert!(alpha_after < alpha_before); + + // Burn should NOT decrease SubnetAlphaOut (unlike recycle) + let alpha_out_after = pallet_subtensor::SubnetAlphaOut::::get(netuid); + assert_eq!(alpha_out_after, alpha_out_before); + }); +} + +#[test] +fn burn_alpha_on_nonexistent_subnet_returns_error() { + mock::new_test_ext(1).execute_with(|| { + let coldkey = U256::from(9501); + let hotkey = U256::from(9502); + + let expected_weight = Weight::from_parts(112_200_000, 0) + .saturating_add(::DbWeight::get().reads(10)) + .saturating_add(::DbWeight::get().writes(3)); + + let mut env = MockEnv::new( + FunctionId::BurnAlphaV1, + coldkey, + (hotkey, AlphaBalance::from(1_000u64), NetUid::from(999u16)).encode(), + ) + .with_expected_weight(expected_weight); + + let ret = SubtensorChainExtension::::dispatch(&mut env).unwrap(); + match ret { + RetVal::Converging(code) => { + assert_eq!( + code, + Output::SubnetNotExists as u32, + "expected subnet not exists error" + ) + } + _ => panic!("unexpected return value"), + } + }); +} + +#[test] +fn add_stake_recycle_success_atomically_stakes_and_recycles() { + mock::new_test_ext(1).execute_with(|| { + let owner_hotkey = U256::from(9601); + let owner_coldkey = U256::from(9602); + let coldkey = U256::from(9701); + let hotkey = U256::from(9702); + let min_stake = DefaultMinStake::::get(); + let tao_amount_raw = min_stake.to_u64().saturating_mul(200); + + let netuid = mock::add_dynamic_network(&owner_hotkey, &owner_coldkey); + mock::setup_reserves( + netuid, + TaoBalance::from(130_000_000_000_u64), + AlphaBalance::from(110_000_000_000_u64), + ); + + mock::register_ok_neuron(netuid, hotkey, coldkey, 0); + + pallet_subtensor::Pallet::::add_balance_to_coldkey_account( + &coldkey, + TaoBalance::from(tao_amount_raw.saturating_add(1_000_000_000)), + ); + + let alpha_out_before = pallet_subtensor::SubnetAlphaOut::::get(netuid); + + let expected_weight = Weight::from_parts(454_200_000, 0) + .saturating_add(::DbWeight::get().reads(33)) + .saturating_add(::DbWeight::get().writes(19)); + + let mut env = MockEnv::new( + FunctionId::AddStakeRecycleV1, + coldkey, + (hotkey, netuid, TaoBalance::from(tao_amount_raw)).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 returned_alpha = AlphaBalance::decode(&mut env.output()).unwrap(); + assert!(returned_alpha > AlphaBalance::ZERO); + + // After atomic add+recycle, the stake should be zero (we recycled everything we added) + let alpha_after = + pallet_subtensor::Pallet::::get_stake_for_hotkey_and_coldkey_on_subnet( + &hotkey, &coldkey, netuid, + ); + assert!(alpha_after.is_zero()); + + // SubnetAlphaOut should not have increased (recycle cancels out the add) + let alpha_out_after = pallet_subtensor::SubnetAlphaOut::::get(netuid); + assert!(alpha_out_after <= alpha_out_before); + }); +} + +#[test] +fn add_stake_burn_success_atomically_stakes_and_burns() { + mock::new_test_ext(1).execute_with(|| { + let owner_hotkey = U256::from(9801); + let owner_coldkey = U256::from(9802); + let coldkey = U256::from(9901); + let hotkey = U256::from(9902); + let min_stake = DefaultMinStake::::get(); + let tao_amount_raw = min_stake.to_u64().saturating_mul(200); + + let netuid = mock::add_dynamic_network(&owner_hotkey, &owner_coldkey); + mock::setup_reserves( + netuid, + TaoBalance::from(130_000_000_000_u64), + AlphaBalance::from(110_000_000_000_u64), + ); + + mock::register_ok_neuron(netuid, hotkey, coldkey, 0); + + pallet_subtensor::Pallet::::add_balance_to_coldkey_account( + &coldkey, + TaoBalance::from(tao_amount_raw.saturating_add(1_000_000_000)), + ); + + let alpha_out_before = pallet_subtensor::SubnetAlphaOut::::get(netuid); + + let expected_weight = Weight::from_parts(453_000_000, 0) + .saturating_add(::DbWeight::get().reads(33)) + .saturating_add(::DbWeight::get().writes(18)); + + let mut env = MockEnv::new( + FunctionId::AddStakeBurnV1, + coldkey, + (hotkey, netuid, TaoBalance::from(tao_amount_raw)).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 returned_alpha = AlphaBalance::decode(&mut env.output()).unwrap(); + assert!(returned_alpha > AlphaBalance::ZERO); + + // After atomic add+burn, the stake should be zero + let alpha_after = + pallet_subtensor::Pallet::::get_stake_for_hotkey_and_coldkey_on_subnet( + &hotkey, &coldkey, netuid, + ); + assert!(alpha_after.is_zero()); + + // SubnetAlphaOut should have increased (burn does NOT reduce AlphaOut) + let alpha_out_after = pallet_subtensor::SubnetAlphaOut::::get(netuid); + assert!(alpha_out_after > alpha_out_before); + }); +} + +#[test] +fn add_stake_recycle_with_insufficient_balance_returns_error() { + mock::new_test_ext(1).execute_with(|| { + let owner_hotkey = U256::from(10001); + let owner_coldkey = U256::from(10002); + let coldkey = U256::from(10101); + let hotkey = U256::from(10102); + + let netuid = mock::add_dynamic_network(&owner_hotkey, &owner_coldkey); + mock::setup_reserves( + netuid, + TaoBalance::from(130_000_000_000_u64), + AlphaBalance::from(110_000_000_000_u64), + ); + + mock::register_ok_neuron(netuid, hotkey, coldkey, 0); + + // Don't fund the coldkey - should fail with balance error + + let expected_weight = Weight::from_parts(454_200_000, 0) + .saturating_add(::DbWeight::get().reads(33)) + .saturating_add(::DbWeight::get().writes(19)); + + let mut env = MockEnv::new( + FunctionId::AddStakeRecycleV1, + coldkey, + (hotkey, netuid, TaoBalance::from(100_000_000_000_u64)).encode(), + ) + .with_expected_weight(expected_weight); + + let ret = SubtensorChainExtension::::dispatch(&mut env).unwrap(); + match ret { + RetVal::Converging(code) => { + assert_ne!(code, Output::Success as u32, "should not succeed") + } + _ => panic!("unexpected return value"), + } + }); +} + impl MockEnv { fn new(func_id: FunctionId, caller: AccountId, input: Vec) -> Self { Self { diff --git a/chain-extensions/src/types.rs b/chain-extensions/src/types.rs index ee6298ad5b..424b4848d9 100644 --- a/chain-extensions/src/types.rs +++ b/chain-extensions/src/types.rs @@ -21,6 +21,10 @@ pub enum FunctionId { AddProxyV1 = 13, RemoveProxyV1 = 14, GetAlphaPriceV1 = 15, + RecycleAlphaV1 = 16, + BurnAlphaV1 = 17, + AddStakeRecycleV1 = 18, + AddStakeBurnV1 = 19, } #[derive(PartialEq, Eq, Copy, Clone, Encode, Decode, Debug)] @@ -66,6 +70,10 @@ pub enum Output { ProxyNoSelfProxy = 18, /// Proxy relationship not found ProxyNotFound = 19, + /// Cannot burn or recycle on root subnet + CannotBurnOrRecycleOnRootSubnet = 20, + /// Subtoken is disabled for this subnet + SubtokenDisabled = 21, } impl From for Output { @@ -93,6 +101,8 @@ impl From for Output { Some("Duplicate") => Output::ProxyDuplicate, Some("NoSelfProxy") => Output::ProxyNoSelfProxy, Some("NotFound") => Output::ProxyNotFound, + Some("CannotBurnOrRecycleOnRootSubnet") => Output::CannotBurnOrRecycleOnRootSubnet, + Some("SubtokenDisabled") => Output::SubtokenDisabled, _ => Output::RuntimeError, } } diff --git a/contract-tests/bittensor/lib.rs b/contract-tests/bittensor/lib.rs index 8867d017d8..a81066d5e3 100755 --- a/contract-tests/bittensor/lib.rs +++ b/contract-tests/bittensor/lib.rs @@ -22,6 +22,10 @@ pub enum FunctionId { AddProxyV1 = 13, RemoveProxyV1 = 14, GetAlphaPriceV1 = 15, + RecycleAlphaV1 = 16, + BurnAlphaV1 = 17, + AddStakeRecycleV1 = 18, + AddStakeBurnV1 = 19, } #[ink::chain_extension(extension = 0x1000)] @@ -130,6 +134,34 @@ pub trait RuntimeReadWrite { #[ink(function = 15)] fn get_alpha_price(netuid: u16) -> u64; + + #[ink(function = 16)] + fn recycle_alpha( + hotkey: ::AccountId, + amount: u64, + netuid: u16, + ) -> u64; + + #[ink(function = 17)] + fn burn_alpha( + hotkey: ::AccountId, + amount: u64, + netuid: u16, + ) -> u64; + + #[ink(function = 18)] + fn add_stake_recycle( + hotkey: ::AccountId, + netuid: u16, + amount: u64, + ) -> u64; + + #[ink(function = 19)] + fn add_stake_burn( + hotkey: ::AccountId, + netuid: u16, + amount: u64, + ) -> u64; } #[ink::scale_derive(Encode, Decode, TypeInfo)] @@ -412,5 +444,57 @@ mod bittensor { .get_alpha_price(netuid) .map_err(|_e| ReadWriteErrorCode::ReadFailed) } + + #[ink(message)] + pub fn recycle_alpha( + &self, + hotkey: [u8; 32], + amount: u64, + netuid: u16, + ) -> Result { + self.env() + .extension() + .recycle_alpha(hotkey.into(), amount, netuid) + .map_err(|_e| ReadWriteErrorCode::WriteFailed) + } + + #[ink(message)] + pub fn burn_alpha( + &self, + hotkey: [u8; 32], + amount: u64, + netuid: u16, + ) -> Result { + self.env() + .extension() + .burn_alpha(hotkey.into(), amount, netuid) + .map_err(|_e| ReadWriteErrorCode::WriteFailed) + } + + #[ink(message)] + pub fn add_stake_recycle( + &self, + hotkey: [u8; 32], + netuid: u16, + amount: u64, + ) -> Result { + self.env() + .extension() + .add_stake_recycle(hotkey.into(), netuid, amount) + .map_err(|_e| ReadWriteErrorCode::WriteFailed) + } + + #[ink(message)] + pub fn add_stake_burn( + &self, + hotkey: [u8; 32], + netuid: u16, + amount: u64, + ) -> Result { + self.env() + .extension() + .add_stake_burn(hotkey.into(), netuid, amount) + .map_err(|_e| ReadWriteErrorCode::WriteFailed) + } } } diff --git a/docs/wasm-contracts.md b/docs/wasm-contracts.md index d3a6b5637f..5e3f202688 100644 --- a/docs/wasm-contracts.md +++ b/docs/wasm-contracts.md @@ -43,6 +43,11 @@ Subtensor provides a custom chain extension that allows smart contracts to inter | 12 | `set_coldkey_auto_stake_hotkey` | Configure automatic stake destination | `(NetUid, AccountId)` | Error code | | 13 | `add_proxy` | Add a staking proxy for the caller | `(AccountId)` | Error code | | 14 | `remove_proxy` | Remove a staking proxy for the caller | `(AccountId)` | Error code | +| 15 | `get_alpha_price` | Query the current alpha price for a subnet | `(NetUid)` | `u64` (price × 10⁹) | +| 16 | `recycle_alpha` | Recycle alpha stake, reducing SubnetAlphaOut (supply reduction) | `(AccountId, AlphaBalance, NetUid)` | `u64` (actual amount recycled) | +| 17 | `burn_alpha` | Burn alpha stake without reducing SubnetAlphaOut (supply neutral) | `(AccountId, AlphaBalance, NetUid)` | `u64` (actual amount burned) | +| 18 | `add_stake_recycle` | Atomically add stake then recycle the resulting alpha | `(AccountId, NetUid, TaoBalance)` | `u64` (alpha amount recycled) | +| 19 | `add_stake_burn` | Atomically add stake then burn the resulting alpha | `(AccountId, NetUid, TaoBalance)` | `u64` (alpha amount burned) | Example usage in your ink! contract: ```rust @@ -85,6 +90,8 @@ Chain extension functions that modify state return error codes as `u32` values. | 17 | `ProxyDuplicate` | Proxy already exists | | 18 | `ProxyNoSelfProxy` | Cannot add self as proxy | | 19 | `ProxyNotFound` | Proxy relationship not found | +| 20 | `CannotBurnOrRecycleOnRootSubnet` | Cannot burn or recycle on the root subnet | +| 21 | `SubtokenDisabled` | Subtoken is not enabled for the specified subnet | ### Call Filter From 8ebf23439c775309ee69cc9594ba50baa7d508c6 Mon Sep 17 00:00:00 2001 From: Landyn Date: Wed, 1 Apr 2026 11:26:42 -0500 Subject: [PATCH 2/9] Remove debug logging from chain extension dispatch Strip development log::info!/log::error! calls from dispatch entry and AddStakeRecycleV1 handler. Normalize AddStakeRecycleV1 to use the same concise ? pattern as all other handlers. --- chain-extensions/src/lib.rs | 33 ++++----------------------------- 1 file changed, 4 insertions(+), 29 deletions(-) diff --git a/chain-extensions/src/lib.rs b/chain-extensions/src/lib.rs index 0cd1523cac..c4bb55f4de 100644 --- a/chain-extensions/src/lib.rs +++ b/chain-extensions/src/lib.rs @@ -66,10 +66,7 @@ where Env: SubtensorExtensionEnv, <::Lookup as StaticLookup>::Source: From<::AccountId>, { - let raw_func_id = env.func_id(); - log::info!("chain_ext: dispatch called with raw func_id={raw_func_id}"); - let func_id: FunctionId = raw_func_id.try_into().map_err(|_| { - log::error!("chain_ext: invalid func_id={raw_func_id}, not in FunctionId enum"); + let func_id: FunctionId = env.func_id().try_into().map_err(|_| { DispatchError::Other( "Invalid function id - does not correspond to any registered function", ) @@ -601,30 +598,16 @@ where } } FunctionId::AddStakeRecycleV1 => { - log::info!("chain_ext: AddStakeRecycleV1 called"); - let weight = Weight::from_parts(454_200_000, 0) .saturating_add(T::DbWeight::get().reads(33)) .saturating_add(T::DbWeight::get().writes(19)); - if let Err(e) = env.charge_weight(weight) { - log::error!("chain_ext: AddStakeRecycleV1 charge_weight failed: {e:?}"); - return Err(e); - } + env.charge_weight(weight)?; - let input: Result<(T::AccountId, NetUid, TaoBalance), _> = env.read_as(); - let (hotkey, netuid, tao_amount) = match input { - Ok(v) => v, - Err(e) => { - log::error!("chain_ext: AddStakeRecycleV1 read_as failed: {e:?}"); - return Err(e); - } - }; + let (hotkey, netuid, tao_amount): (T::AccountId, NetUid, TaoBalance) = + env.read_as()?; let caller = env.caller(); - log::info!( - "chain_ext: AddStakeRecycleV1 caller={caller:?} hotkey={hotkey:?} netuid={netuid:?} tao={tao_amount:?}" - ); let alpha = pallet_subtensor::Pallet::::do_add_stake( RawOrigin::Signed(caller.clone()).into(), @@ -635,9 +618,6 @@ where match alpha { Ok(alpha) => { - log::info!( - "chain_ext: AddStakeRecycleV1 do_add_stake ok, alpha={alpha:?}" - ); let recycle_result = pallet_subtensor::Pallet::::recycle_alpha( RawOrigin::Signed(caller).into(), hotkey, @@ -647,22 +627,17 @@ where match recycle_result { Ok(_) => { - log::info!("chain_ext: AddStakeRecycleV1 recycle ok"); env.write_output(&alpha.encode()) .map_err(|_| DispatchError::Other("Failed to write output"))?; Ok(RetVal::Converging(Output::Success as u32)) } Err(e) => { - log::error!( - "chain_ext: AddStakeRecycleV1 recycle failed: {e:?}" - ); let error_code = Output::from(e) as u32; Ok(RetVal::Converging(error_code)) } } } Err(e) => { - log::error!("chain_ext: AddStakeRecycleV1 do_add_stake failed: {e:?}"); let error_code = Output::from(e) as u32; Ok(RetVal::Converging(error_code)) } From 94796b0fea0c3341058f37977ef09874bb0ae490 Mon Sep 17 00:00:00 2001 From: Landyn Date: Fri, 3 Apr 2026 11:12:44 -0500 Subject: [PATCH 3/9] Additional unit tests: Root burn test, Zero amt test, Amount clamping test --- chain-extensions/src/tests.rs | 163 ++++++++++++++++++++++++++++++++++ 1 file changed, 163 insertions(+) diff --git a/chain-extensions/src/tests.rs b/chain-extensions/src/tests.rs index 284e852f24..c997986503 100644 --- a/chain-extensions/src/tests.rs +++ b/chain-extensions/src/tests.rs @@ -1085,6 +1085,169 @@ fn add_stake_recycle_with_insufficient_balance_returns_error() { }); } +#[test] +fn recycle_alpha_clamps_to_available_when_amount_exceeds_stake() { + mock::new_test_ext(1).execute_with(|| { + let owner_hotkey = U256::from(11001); + let owner_coldkey = U256::from(11002); + let coldkey = U256::from(11101); + let hotkey = U256::from(11102); + let min_stake = DefaultMinStake::::get(); + let stake_amount_raw = min_stake.to_u64().saturating_mul(200); + + let netuid = mock::add_dynamic_network(&owner_hotkey, &owner_coldkey); + mock::setup_reserves( + netuid, + TaoBalance::from(130_000_000_000_u64), + AlphaBalance::from(110_000_000_000_u64), + ); + + mock::register_ok_neuron(netuid, hotkey, coldkey, 0); + + pallet_subtensor::Pallet::::add_balance_to_coldkey_account( + &coldkey, + TaoBalance::from(stake_amount_raw.saturating_add(1_000_000_000)), + ); + + assert_ok!(pallet_subtensor::Pallet::::add_stake( + RawOrigin::Signed(coldkey).into(), + hotkey, + netuid, + stake_amount_raw.into(), + )); + + let alpha_before = + pallet_subtensor::Pallet::::get_stake_for_hotkey_and_coldkey_on_subnet( + &hotkey, &coldkey, netuid, + ); + assert!(alpha_before > AlphaBalance::ZERO); + + // Request way more than available — should clamp to alpha_before + let huge_amount = AlphaBalance::from(u64::MAX); + + let expected_weight = Weight::from_parts(113_400_000, 0) + .saturating_add(::DbWeight::get().reads(10)) + .saturating_add(::DbWeight::get().writes(4)); + + let mut env = MockEnv::new( + FunctionId::RecycleAlphaV1, + coldkey, + (hotkey, huge_amount, netuid).encode(), + ) + .with_expected_weight(expected_weight); + + let ret = SubtensorChainExtension::::dispatch(&mut env).unwrap(); + assert_success(ret); + + let returned_amount = AlphaBalance::decode(&mut env.output()).unwrap(); + assert_eq!(returned_amount, alpha_before, "should clamp to available alpha"); + + let alpha_after = + pallet_subtensor::Pallet::::get_stake_for_hotkey_and_coldkey_on_subnet( + &hotkey, &coldkey, netuid, + ); + assert!(alpha_after.is_zero(), "all alpha should be recycled"); + }); +} + +#[test] +fn burn_alpha_on_root_subnet_returns_error() { + mock::new_test_ext(1).execute_with(|| { + let coldkey = U256::from(11201); + let hotkey = U256::from(11202); + + pallet_subtensor::Owner::::insert(hotkey, coldkey); + + let expected_weight = Weight::from_parts(112_200_000, 0) + .saturating_add(::DbWeight::get().reads(10)) + .saturating_add(::DbWeight::get().writes(3)); + + let mut env = MockEnv::new( + FunctionId::BurnAlphaV1, + coldkey, + (hotkey, AlphaBalance::from(1_000u64), NetUid::ROOT).encode(), + ) + .with_expected_weight(expected_weight); + + let ret = SubtensorChainExtension::::dispatch(&mut env).unwrap(); + match ret { + RetVal::Converging(code) => { + assert_ne!( + code, + Output::Success as u32, + "should not succeed on root subnet" + ) + } + _ => panic!("unexpected return value"), + } + }); +} + +#[test] +fn burn_alpha_clamps_to_available_when_amount_exceeds_stake() { + mock::new_test_ext(1).execute_with(|| { + let owner_hotkey = U256::from(11301); + let owner_coldkey = U256::from(11302); + let coldkey = U256::from(11401); + let hotkey = U256::from(11402); + let min_stake = DefaultMinStake::::get(); + let stake_amount_raw = min_stake.to_u64().saturating_mul(200); + + let netuid = mock::add_dynamic_network(&owner_hotkey, &owner_coldkey); + mock::setup_reserves( + netuid, + TaoBalance::from(130_000_000_000_u64), + AlphaBalance::from(110_000_000_000_u64), + ); + + mock::register_ok_neuron(netuid, hotkey, coldkey, 0); + + pallet_subtensor::Pallet::::add_balance_to_coldkey_account( + &coldkey, + TaoBalance::from(stake_amount_raw.saturating_add(1_000_000_000)), + ); + + assert_ok!(pallet_subtensor::Pallet::::add_stake( + RawOrigin::Signed(coldkey).into(), + hotkey, + netuid, + stake_amount_raw.into(), + )); + + let alpha_before = + pallet_subtensor::Pallet::::get_stake_for_hotkey_and_coldkey_on_subnet( + &hotkey, &coldkey, netuid, + ); + assert!(alpha_before > AlphaBalance::ZERO); + + // Request way more than available — should clamp to alpha_before + let huge_amount = AlphaBalance::from(u64::MAX); + + let expected_weight = Weight::from_parts(112_200_000, 0) + .saturating_add(::DbWeight::get().reads(10)) + .saturating_add(::DbWeight::get().writes(3)); + + let mut env = MockEnv::new( + FunctionId::BurnAlphaV1, + coldkey, + (hotkey, huge_amount, netuid).encode(), + ) + .with_expected_weight(expected_weight); + + let ret = SubtensorChainExtension::::dispatch(&mut env).unwrap(); + assert_success(ret); + + let returned_amount = AlphaBalance::decode(&mut env.output()).unwrap(); + assert_eq!(returned_amount, alpha_before, "should clamp to available alpha"); + + let alpha_after = + pallet_subtensor::Pallet::::get_stake_for_hotkey_and_coldkey_on_subnet( + &hotkey, &coldkey, netuid, + ); + assert!(alpha_after.is_zero(), "all alpha should be burned"); + }); +} + impl MockEnv { fn new(func_id: FunctionId, caller: AccountId, input: Vec) -> Self { Self { From 1c4f180f8dd5a10fc0125481ae22a1c7f183c3ef Mon Sep 17 00:00:00 2001 From: Landyn Date: Mon, 6 Apr 2026 19:40:01 -0500 Subject: [PATCH 4/9] cargo fmt fix --- chain-extensions/src/tests.rs | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/chain-extensions/src/tests.rs b/chain-extensions/src/tests.rs index c997986503..fd08232dda 100644 --- a/chain-extensions/src/tests.rs +++ b/chain-extensions/src/tests.rs @@ -1140,7 +1140,10 @@ fn recycle_alpha_clamps_to_available_when_amount_exceeds_stake() { assert_success(ret); let returned_amount = AlphaBalance::decode(&mut env.output()).unwrap(); - assert_eq!(returned_amount, alpha_before, "should clamp to available alpha"); + assert_eq!( + returned_amount, alpha_before, + "should clamp to available alpha" + ); let alpha_after = pallet_subtensor::Pallet::::get_stake_for_hotkey_and_coldkey_on_subnet( @@ -1238,7 +1241,10 @@ fn burn_alpha_clamps_to_available_when_amount_exceeds_stake() { assert_success(ret); let returned_amount = AlphaBalance::decode(&mut env.output()).unwrap(); - assert_eq!(returned_amount, alpha_before, "should clamp to available alpha"); + assert_eq!( + returned_amount, alpha_before, + "should clamp to available alpha" + ); let alpha_after = pallet_subtensor::Pallet::::get_stake_for_hotkey_and_coldkey_on_subnet( From 65357d303af1317dd5513cde74d5c8c86454e32c Mon Sep 17 00:00:00 2001 From: Landyn Date: Tue, 7 Apr 2026 09:26:08 -0500 Subject: [PATCH 5/9] Address review: remove redundant clamping, add atomicity --- chain-extensions/src/lib.rs | 121 +++++++++++++--------------- chain-extensions/src/tests.rs | 146 +++++++++++++++++++++++++++++++++- 2 files changed, 197 insertions(+), 70 deletions(-) diff --git a/chain-extensions/src/lib.rs b/chain-extensions/src/lib.rs index c4bb55f4de..1df272fc1d 100644 --- a/chain-extensions/src/lib.rs +++ b/chain-extensions/src/lib.rs @@ -9,6 +9,7 @@ pub mod types; use crate::types::{FunctionId, Output}; use codec::{Decode, Encode, MaxEncodedLen}; +use frame_support::storage::{TransactionOutcome, transactional}; use frame_support::{DebugNoBound, traits::Get}; use frame_system::RawOrigin; use pallet_contracts::chain_extension::{ @@ -535,22 +536,16 @@ where let caller = env.caller(); - let alpha_available = - pallet_subtensor::Pallet::::get_stake_for_hotkey_and_coldkey_on_subnet( - &hotkey, &caller, netuid, - ); - let actual_amount = amount.min(alpha_available); - let call_result = pallet_subtensor::Pallet::::recycle_alpha( RawOrigin::Signed(caller).into(), hotkey, - actual_amount, + amount, netuid, ); match call_result { Ok(_) => { - env.write_output(&actual_amount.encode()) + env.write_output(&amount.encode()) .map_err(|_| DispatchError::Other("Failed to write output"))?; Ok(RetVal::Converging(Output::Success as u32)) } @@ -572,22 +567,16 @@ where let caller = env.caller(); - let alpha_available = - pallet_subtensor::Pallet::::get_stake_for_hotkey_and_coldkey_on_subnet( - &hotkey, &caller, netuid, - ); - let actual_amount = amount.min(alpha_available); - let call_result = pallet_subtensor::Pallet::::burn_alpha( RawOrigin::Signed(caller).into(), hotkey, - actual_amount, + amount, netuid, ); match call_result { Ok(_) => { - env.write_output(&actual_amount.encode()) + env.write_output(&amount.encode()) .map_err(|_| DispatchError::Other("Failed to write output"))?; Ok(RetVal::Converging(Output::Success as u32)) } @@ -609,33 +598,33 @@ where let caller = env.caller(); - let alpha = pallet_subtensor::Pallet::::do_add_stake( - RawOrigin::Signed(caller.clone()).into(), - hotkey.clone(), - netuid, - tao_amount, - ); + let result = transactional::with_transaction(|| { + let alpha = match pallet_subtensor::Pallet::::do_add_stake( + RawOrigin::Signed(caller.clone()).into(), + hotkey.clone(), + netuid, + tao_amount, + ) { + Ok(a) => a, + Err(e) => return TransactionOutcome::Rollback(Err(e)), + }; + + match pallet_subtensor::Pallet::::recycle_alpha( + RawOrigin::Signed(caller).into(), + hotkey, + alpha, + netuid, + ) { + Ok(_) => TransactionOutcome::Commit(Ok(alpha)), + Err(e) => TransactionOutcome::Rollback(Err(e)), + } + }); - match alpha { + match result { Ok(alpha) => { - let recycle_result = pallet_subtensor::Pallet::::recycle_alpha( - RawOrigin::Signed(caller).into(), - hotkey, - alpha, - netuid, - ); - - match recycle_result { - Ok(_) => { - env.write_output(&alpha.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)) - } - } + env.write_output(&alpha.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; @@ -655,33 +644,33 @@ where let caller = env.caller(); - let alpha = pallet_subtensor::Pallet::::do_add_stake( - RawOrigin::Signed(caller.clone()).into(), - hotkey.clone(), - netuid, - tao_amount, - ); + let result = transactional::with_transaction(|| { + let alpha = match pallet_subtensor::Pallet::::do_add_stake( + RawOrigin::Signed(caller.clone()).into(), + hotkey.clone(), + netuid, + tao_amount, + ) { + Ok(a) => a, + Err(e) => return TransactionOutcome::Rollback(Err(e)), + }; + + match pallet_subtensor::Pallet::::burn_alpha( + RawOrigin::Signed(caller).into(), + hotkey, + alpha, + netuid, + ) { + Ok(_) => TransactionOutcome::Commit(Ok(alpha)), + Err(e) => TransactionOutcome::Rollback(Err(e)), + } + }); - match alpha { + match result { Ok(alpha) => { - let burn_result = pallet_subtensor::Pallet::::burn_alpha( - RawOrigin::Signed(caller).into(), - hotkey, - alpha, - netuid, - ); - - match burn_result { - Ok(_) => { - env.write_output(&alpha.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)) - } - } + env.write_output(&alpha.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; diff --git a/chain-extensions/src/tests.rs b/chain-extensions/src/tests.rs index fd08232dda..55db1d2c2e 100644 --- a/chain-extensions/src/tests.rs +++ b/chain-extensions/src/tests.rs @@ -1141,8 +1141,8 @@ fn recycle_alpha_clamps_to_available_when_amount_exceeds_stake() { let returned_amount = AlphaBalance::decode(&mut env.output()).unwrap(); assert_eq!( - returned_amount, alpha_before, - "should clamp to available alpha" + returned_amount, huge_amount, + "should return requested amount" ); let alpha_after = @@ -1242,8 +1242,8 @@ fn burn_alpha_clamps_to_available_when_amount_exceeds_stake() { let returned_amount = AlphaBalance::decode(&mut env.output()).unwrap(); assert_eq!( - returned_amount, alpha_before, - "should clamp to available alpha" + returned_amount, huge_amount, + "should return requested amount" ); let alpha_after = @@ -1321,6 +1321,144 @@ fn assert_success(ret: RetVal) { } } +#[test] +fn add_stake_recycle_rollback_on_recycle_failure() { + mock::new_test_ext(1).execute_with(|| { + let owner_hotkey = U256::from(12001); + let owner_coldkey = U256::from(12002); + let coldkey = U256::from(12101); + let hotkey = U256::from(12102); + let min_stake = DefaultMinStake::::get(); + let tao_amount_raw = min_stake.to_u64().saturating_mul(200); + + let netuid = mock::add_dynamic_network(&owner_hotkey, &owner_coldkey); + + // Set up very low reserves so recycle will fail with InsufficientLiquidity + mock::setup_reserves( + netuid, + TaoBalance::from(1_000_u64), + AlphaBalance::from(1_000_u64), + ); + + mock::register_ok_neuron(netuid, hotkey, coldkey, 0); + + pallet_subtensor::Pallet::::add_balance_to_coldkey_account( + &coldkey, + TaoBalance::from(tao_amount_raw.saturating_add(1_000_000_000)), + ); + + let balance_before = pallet_subtensor::Pallet::::get_coldkey_balance(&coldkey); + let alpha_before = + pallet_subtensor::Pallet::::get_stake_for_hotkey_and_coldkey_on_subnet( + &hotkey, &coldkey, netuid, + ); + + let expected_weight = Weight::from_parts(454_200_000, 0) + .saturating_add(::DbWeight::get().reads(33)) + .saturating_add(::DbWeight::get().writes(19)); + + let mut env = MockEnv::new( + FunctionId::AddStakeRecycleV1, + coldkey, + (hotkey, netuid, TaoBalance::from(tao_amount_raw)).encode(), + ) + .with_expected_weight(expected_weight); + + let ret = SubtensorChainExtension::::dispatch(&mut env).unwrap(); + match ret { + RetVal::Converging(code) => { + assert_ne!(code, Output::Success as u32, "should not succeed") + } + _ => panic!("unexpected return value"), + } + + // Verify full rollback: balance and stake unchanged + let balance_after = pallet_subtensor::Pallet::::get_coldkey_balance(&coldkey); + let alpha_after = + pallet_subtensor::Pallet::::get_stake_for_hotkey_and_coldkey_on_subnet( + &hotkey, &coldkey, netuid, + ); + + assert_eq!( + balance_before, balance_after, + "balance should be unchanged after rollback" + ); + assert_eq!( + alpha_before, alpha_after, + "stake should be unchanged after rollback" + ); + }); +} + +#[test] +fn add_stake_burn_rollback_on_burn_failure() { + mock::new_test_ext(1).execute_with(|| { + let owner_hotkey = U256::from(12201); + let owner_coldkey = U256::from(12202); + let coldkey = U256::from(12301); + let hotkey = U256::from(12302); + let min_stake = DefaultMinStake::::get(); + let tao_amount_raw = min_stake.to_u64().saturating_mul(200); + + let netuid = mock::add_dynamic_network(&owner_hotkey, &owner_coldkey); + + // Set up very low reserves so burn will fail with InsufficientLiquidity + mock::setup_reserves( + netuid, + TaoBalance::from(1_000_u64), + AlphaBalance::from(1_000_u64), + ); + + mock::register_ok_neuron(netuid, hotkey, coldkey, 0); + + pallet_subtensor::Pallet::::add_balance_to_coldkey_account( + &coldkey, + TaoBalance::from(tao_amount_raw.saturating_add(1_000_000_000)), + ); + + let balance_before = pallet_subtensor::Pallet::::get_coldkey_balance(&coldkey); + let alpha_before = + pallet_subtensor::Pallet::::get_stake_for_hotkey_and_coldkey_on_subnet( + &hotkey, &coldkey, netuid, + ); + + let expected_weight = Weight::from_parts(453_000_000, 0) + .saturating_add(::DbWeight::get().reads(33)) + .saturating_add(::DbWeight::get().writes(18)); + + let mut env = MockEnv::new( + FunctionId::AddStakeBurnV1, + coldkey, + (hotkey, netuid, TaoBalance::from(tao_amount_raw)).encode(), + ) + .with_expected_weight(expected_weight); + + let ret = SubtensorChainExtension::::dispatch(&mut env).unwrap(); + match ret { + RetVal::Converging(code) => { + assert_ne!(code, Output::Success as u32, "should not succeed") + } + _ => panic!("unexpected return value"), + } + + // Verify full rollback: balance and stake unchanged + let balance_after = pallet_subtensor::Pallet::::get_coldkey_balance(&coldkey); + let alpha_after = + pallet_subtensor::Pallet::::get_stake_for_hotkey_and_coldkey_on_subnet( + &hotkey, &coldkey, netuid, + ); + + assert_eq!( + balance_before, balance_after, + "balance should be unchanged after rollback" + ); + assert_eq!( + alpha_before, alpha_after, + "stake should be unchanged after rollback" + ); + }); +} + #[test] fn get_stake_info_returns_encoded_runtime_value() { mock::new_test_ext(1).execute_with(|| { From 60dffb39a6a67135e7fb263b09f9cc45f45af4f7 Mon Sep 17 00:00:00 2001 From: Landyn Date: Wed, 8 Apr 2026 11:14:28 -0500 Subject: [PATCH 6/9] Return actual recycled/burned alpha amount from pallet --- chain-extensions/src/lib.rs | 20 +++++++++---------- chain-extensions/src/tests.rs | 8 ++++---- pallets/subtensor/src/macros/dispatches.rs | 4 ++-- .../subtensor/src/staking/recycle_alpha.rs | 16 +++++++-------- 4 files changed, 24 insertions(+), 24 deletions(-) diff --git a/chain-extensions/src/lib.rs b/chain-extensions/src/lib.rs index 1df272fc1d..954176952f 100644 --- a/chain-extensions/src/lib.rs +++ b/chain-extensions/src/lib.rs @@ -536,7 +536,7 @@ where let caller = env.caller(); - let call_result = pallet_subtensor::Pallet::::recycle_alpha( + let call_result = pallet_subtensor::Pallet::::do_recycle_alpha( RawOrigin::Signed(caller).into(), hotkey, amount, @@ -544,8 +544,8 @@ where ); match call_result { - Ok(_) => { - env.write_output(&amount.encode()) + Ok(real_amount) => { + env.write_output(&real_amount.encode()) .map_err(|_| DispatchError::Other("Failed to write output"))?; Ok(RetVal::Converging(Output::Success as u32)) } @@ -567,7 +567,7 @@ where let caller = env.caller(); - let call_result = pallet_subtensor::Pallet::::burn_alpha( + let call_result = pallet_subtensor::Pallet::::do_burn_alpha( RawOrigin::Signed(caller).into(), hotkey, amount, @@ -575,8 +575,8 @@ where ); match call_result { - Ok(_) => { - env.write_output(&amount.encode()) + Ok(real_amount) => { + env.write_output(&real_amount.encode()) .map_err(|_| DispatchError::Other("Failed to write output"))?; Ok(RetVal::Converging(Output::Success as u32)) } @@ -609,13 +609,13 @@ where Err(e) => return TransactionOutcome::Rollback(Err(e)), }; - match pallet_subtensor::Pallet::::recycle_alpha( + match pallet_subtensor::Pallet::::do_recycle_alpha( RawOrigin::Signed(caller).into(), hotkey, alpha, netuid, ) { - Ok(_) => TransactionOutcome::Commit(Ok(alpha)), + Ok(real_alpha) => TransactionOutcome::Commit(Ok(real_alpha)), Err(e) => TransactionOutcome::Rollback(Err(e)), } }); @@ -655,13 +655,13 @@ where Err(e) => return TransactionOutcome::Rollback(Err(e)), }; - match pallet_subtensor::Pallet::::burn_alpha( + match pallet_subtensor::Pallet::::do_burn_alpha( RawOrigin::Signed(caller).into(), hotkey, alpha, netuid, ) { - Ok(_) => TransactionOutcome::Commit(Ok(alpha)), + Ok(real_alpha) => TransactionOutcome::Commit(Ok(real_alpha)), Err(e) => TransactionOutcome::Rollback(Err(e)), } }); diff --git a/chain-extensions/src/tests.rs b/chain-extensions/src/tests.rs index 55db1d2c2e..e468270daf 100644 --- a/chain-extensions/src/tests.rs +++ b/chain-extensions/src/tests.rs @@ -1141,8 +1141,8 @@ fn recycle_alpha_clamps_to_available_when_amount_exceeds_stake() { let returned_amount = AlphaBalance::decode(&mut env.output()).unwrap(); assert_eq!( - returned_amount, huge_amount, - "should return requested amount" + returned_amount, alpha_before, + "should return actual clamped amount, not requested amount" ); let alpha_after = @@ -1242,8 +1242,8 @@ fn burn_alpha_clamps_to_available_when_amount_exceeds_stake() { let returned_amount = AlphaBalance::decode(&mut env.output()).unwrap(); assert_eq!( - returned_amount, huge_amount, - "should return requested amount" + returned_amount, alpha_before, + "should return actual clamped amount, not requested amount" ); let alpha_after = diff --git a/pallets/subtensor/src/macros/dispatches.rs b/pallets/subtensor/src/macros/dispatches.rs index aa5ca6ee98..27bcdf1577 100644 --- a/pallets/subtensor/src/macros/dispatches.rs +++ b/pallets/subtensor/src/macros/dispatches.rs @@ -1857,7 +1857,7 @@ mod dispatches { amount: AlphaBalance, netuid: NetUid, ) -> DispatchResult { - Self::do_recycle_alpha(origin, hotkey, amount, netuid) + Self::do_recycle_alpha(origin, hotkey, amount, netuid).map(|_| ()) } /// Burns alpha from a cold/hot key pair without reducing `AlphaOut` @@ -1878,7 +1878,7 @@ mod dispatches { amount: AlphaBalance, netuid: NetUid, ) -> DispatchResult { - Self::do_burn_alpha(origin, hotkey, amount, netuid) + Self::do_burn_alpha(origin, hotkey, amount, netuid).map(|_| ()) } /// Sets the pending childkey cooldown (in blocks). Root only. diff --git a/pallets/subtensor/src/staking/recycle_alpha.rs b/pallets/subtensor/src/staking/recycle_alpha.rs index bb93c12818..96633eaebf 100644 --- a/pallets/subtensor/src/staking/recycle_alpha.rs +++ b/pallets/subtensor/src/staking/recycle_alpha.rs @@ -14,13 +14,13 @@ impl Pallet { /// /// # Returns /// - /// * `DispatchResult` - Success or error - pub(crate) fn do_recycle_alpha( + /// * `Result` - The actual amount recycled, or error + pub fn do_recycle_alpha( origin: OriginFor, hotkey: T::AccountId, amount: AlphaBalance, netuid: NetUid, - ) -> DispatchResult { + ) -> Result { let coldkey: T::AccountId = ensure_signed(origin)?; ensure!(Self::if_subnet_exist(netuid), Error::::SubnetNotExists); @@ -58,7 +58,7 @@ impl Pallet { Self::deposit_event(Event::AlphaRecycled(coldkey, hotkey, amount, netuid)); - Ok(()) + Ok(amount) } /// Burns alpha from a cold/hot key pair without reducing AlphaOut @@ -72,13 +72,13 @@ impl Pallet { /// /// # Returns /// - /// * `DispatchResult` - Success or error - pub(crate) fn do_burn_alpha( + /// * `Result` - The actual amount burned, or error + pub fn do_burn_alpha( origin: OriginFor, hotkey: T::AccountId, amount: AlphaBalance, netuid: NetUid, - ) -> DispatchResult { + ) -> Result { let coldkey = ensure_signed(origin)?; ensure!(Self::if_subnet_exist(netuid), Error::::SubnetNotExists); @@ -116,7 +116,7 @@ impl Pallet { // Deposit event Self::deposit_event(Event::AlphaBurned(coldkey, hotkey, amount, netuid)); - Ok(()) + Ok(amount) } pub(crate) fn do_add_stake_burn( origin: OriginFor, From 757fe295d0bc0adc0a24396d251d3d09619e83bc Mon Sep 17 00:00:00 2001 From: Landyn Date: Mon, 13 Apr 2026 19:48:27 -0500 Subject: [PATCH 7/9] Use pallet WeightInfo trait for recycle/burn chain extensions Replaces hardcoded Weight::from_parts values with calls to pallet_subtensor::weights::WeightInfo, matching the pattern junius introduced in PR 2550. Weights now auto-track benchmark updates. AddStakeRecycleV1 sums add_stake() + recycle_alpha() since the pallet has no add_stake_recycle weight; AddStakeBurnV1 uses add_stake_burn() directly (1:1 with pallet's do_add_stake_burn). --- chain-extensions/src/lib.rs | 24 ++++++------ chain-extensions/src/tests.rs | 70 +++++++++++++++++------------------ 2 files changed, 46 insertions(+), 48 deletions(-) diff --git a/chain-extensions/src/lib.rs b/chain-extensions/src/lib.rs index 954176952f..f2ef5bd96a 100644 --- a/chain-extensions/src/lib.rs +++ b/chain-extensions/src/lib.rs @@ -15,6 +15,7 @@ use frame_system::RawOrigin; use pallet_contracts::chain_extension::{ BufInBufOutState, ChainExtension, Environment, Ext, InitState, RetVal, SysConfig, }; +use pallet_subtensor::weights::WeightInfo as SubtensorWeightInfo; use pallet_subtensor_proxy as pallet_proxy; use pallet_subtensor_proxy::WeightInfo; use sp_runtime::{DispatchError, Weight, traits::StaticLookup}; @@ -525,9 +526,8 @@ where Ok(RetVal::Converging(Output::Success as u32)) } FunctionId::RecycleAlphaV1 => { - let weight = Weight::from_parts(113_400_000, 0) - .saturating_add(T::DbWeight::get().reads(10)) - .saturating_add(T::DbWeight::get().writes(4)); + let weight = + <::WeightInfo as SubtensorWeightInfo>::recycle_alpha(); env.charge_weight(weight)?; @@ -556,9 +556,8 @@ where } } FunctionId::BurnAlphaV1 => { - let weight = Weight::from_parts(112_200_000, 0) - .saturating_add(T::DbWeight::get().reads(10)) - .saturating_add(T::DbWeight::get().writes(3)); + let weight = + <::WeightInfo as SubtensorWeightInfo>::burn_alpha(); env.charge_weight(weight)?; @@ -587,9 +586,11 @@ where } } FunctionId::AddStakeRecycleV1 => { - let weight = Weight::from_parts(454_200_000, 0) - .saturating_add(T::DbWeight::get().reads(33)) - .saturating_add(T::DbWeight::get().writes(19)); + let weight = + <::WeightInfo as SubtensorWeightInfo>::add_stake() + .saturating_add( + <::WeightInfo as SubtensorWeightInfo>::recycle_alpha(), + ); env.charge_weight(weight)?; @@ -633,9 +634,8 @@ where } } FunctionId::AddStakeBurnV1 => { - let weight = Weight::from_parts(453_000_000, 0) - .saturating_add(T::DbWeight::get().reads(33)) - .saturating_add(T::DbWeight::get().writes(18)); + let weight = + <::WeightInfo as SubtensorWeightInfo>::add_stake_burn(); env.charge_weight(weight)?; diff --git a/chain-extensions/src/tests.rs b/chain-extensions/src/tests.rs index e468270daf..aef12f1e0d 100644 --- a/chain-extensions/src/tests.rs +++ b/chain-extensions/src/tests.rs @@ -8,6 +8,7 @@ use frame_support::{assert_ok, weights::Weight}; use frame_system::RawOrigin; use pallet_contracts::chain_extension::RetVal; use pallet_subtensor::DefaultMinStake; +use pallet_subtensor::weights::WeightInfo as SubtensorWeightInfo; use sp_core::Get; use sp_core::U256; use sp_runtime::DispatchError; @@ -767,9 +768,8 @@ fn recycle_alpha_success_reduces_stake_and_returns_actual_amount() { let recycle_amount: AlphaBalance = (alpha_before.to_u64() / 2).into(); - let expected_weight = Weight::from_parts(113_400_000, 0) - .saturating_add(::DbWeight::get().reads(10)) - .saturating_add(::DbWeight::get().writes(4)); + let expected_weight = + <::WeightInfo as SubtensorWeightInfo>::recycle_alpha(); let mut env = MockEnv::new( FunctionId::RecycleAlphaV1, @@ -804,9 +804,8 @@ fn recycle_alpha_on_root_subnet_returns_error() { pallet_subtensor::Owner::::insert(hotkey, coldkey); - let expected_weight = Weight::from_parts(113_400_000, 0) - .saturating_add(::DbWeight::get().reads(10)) - .saturating_add(::DbWeight::get().writes(4)); + let expected_weight = + <::WeightInfo as SubtensorWeightInfo>::recycle_alpha(); let mut env = MockEnv::new( FunctionId::RecycleAlphaV1, @@ -870,9 +869,8 @@ fn burn_alpha_success_reduces_stake_and_returns_actual_amount() { let burn_amount: AlphaBalance = (alpha_before.to_u64() / 2).into(); - let expected_weight = Weight::from_parts(112_200_000, 0) - .saturating_add(::DbWeight::get().reads(10)) - .saturating_add(::DbWeight::get().writes(3)); + let expected_weight = + <::WeightInfo as SubtensorWeightInfo>::burn_alpha(); let mut env = MockEnv::new( FunctionId::BurnAlphaV1, @@ -906,9 +904,8 @@ fn burn_alpha_on_nonexistent_subnet_returns_error() { let coldkey = U256::from(9501); let hotkey = U256::from(9502); - let expected_weight = Weight::from_parts(112_200_000, 0) - .saturating_add(::DbWeight::get().reads(10)) - .saturating_add(::DbWeight::get().writes(3)); + let expected_weight = + <::WeightInfo as SubtensorWeightInfo>::burn_alpha(); let mut env = MockEnv::new( FunctionId::BurnAlphaV1, @@ -957,9 +954,11 @@ fn add_stake_recycle_success_atomically_stakes_and_recycles() { let alpha_out_before = pallet_subtensor::SubnetAlphaOut::::get(netuid); - let expected_weight = Weight::from_parts(454_200_000, 0) - .saturating_add(::DbWeight::get().reads(33)) - .saturating_add(::DbWeight::get().writes(19)); + let expected_weight = + <::WeightInfo as SubtensorWeightInfo>::add_stake() + .saturating_add( + <::WeightInfo as SubtensorWeightInfo>::recycle_alpha(), + ); let mut env = MockEnv::new( FunctionId::AddStakeRecycleV1, @@ -1014,9 +1013,8 @@ fn add_stake_burn_success_atomically_stakes_and_burns() { let alpha_out_before = pallet_subtensor::SubnetAlphaOut::::get(netuid); - let expected_weight = Weight::from_parts(453_000_000, 0) - .saturating_add(::DbWeight::get().reads(33)) - .saturating_add(::DbWeight::get().writes(18)); + let expected_weight = + <::WeightInfo as SubtensorWeightInfo>::add_stake_burn(); let mut env = MockEnv::new( FunctionId::AddStakeBurnV1, @@ -1064,9 +1062,11 @@ fn add_stake_recycle_with_insufficient_balance_returns_error() { // Don't fund the coldkey - should fail with balance error - let expected_weight = Weight::from_parts(454_200_000, 0) - .saturating_add(::DbWeight::get().reads(33)) - .saturating_add(::DbWeight::get().writes(19)); + let expected_weight = + <::WeightInfo as SubtensorWeightInfo>::add_stake() + .saturating_add( + <::WeightInfo as SubtensorWeightInfo>::recycle_alpha(), + ); let mut env = MockEnv::new( FunctionId::AddStakeRecycleV1, @@ -1125,9 +1125,8 @@ fn recycle_alpha_clamps_to_available_when_amount_exceeds_stake() { // Request way more than available — should clamp to alpha_before let huge_amount = AlphaBalance::from(u64::MAX); - let expected_weight = Weight::from_parts(113_400_000, 0) - .saturating_add(::DbWeight::get().reads(10)) - .saturating_add(::DbWeight::get().writes(4)); + let expected_weight = + <::WeightInfo as SubtensorWeightInfo>::recycle_alpha(); let mut env = MockEnv::new( FunctionId::RecycleAlphaV1, @@ -1161,9 +1160,8 @@ fn burn_alpha_on_root_subnet_returns_error() { pallet_subtensor::Owner::::insert(hotkey, coldkey); - let expected_weight = Weight::from_parts(112_200_000, 0) - .saturating_add(::DbWeight::get().reads(10)) - .saturating_add(::DbWeight::get().writes(3)); + let expected_weight = + <::WeightInfo as SubtensorWeightInfo>::burn_alpha(); let mut env = MockEnv::new( FunctionId::BurnAlphaV1, @@ -1226,9 +1224,8 @@ fn burn_alpha_clamps_to_available_when_amount_exceeds_stake() { // Request way more than available — should clamp to alpha_before let huge_amount = AlphaBalance::from(u64::MAX); - let expected_weight = Weight::from_parts(112_200_000, 0) - .saturating_add(::DbWeight::get().reads(10)) - .saturating_add(::DbWeight::get().writes(3)); + let expected_weight = + <::WeightInfo as SubtensorWeightInfo>::burn_alpha(); let mut env = MockEnv::new( FunctionId::BurnAlphaV1, @@ -1353,9 +1350,11 @@ fn add_stake_recycle_rollback_on_recycle_failure() { &hotkey, &coldkey, netuid, ); - let expected_weight = Weight::from_parts(454_200_000, 0) - .saturating_add(::DbWeight::get().reads(33)) - .saturating_add(::DbWeight::get().writes(19)); + let expected_weight = + <::WeightInfo as SubtensorWeightInfo>::add_stake() + .saturating_add( + <::WeightInfo as SubtensorWeightInfo>::recycle_alpha(), + ); let mut env = MockEnv::new( FunctionId::AddStakeRecycleV1, @@ -1422,9 +1421,8 @@ fn add_stake_burn_rollback_on_burn_failure() { &hotkey, &coldkey, netuid, ); - let expected_weight = Weight::from_parts(453_000_000, 0) - .saturating_add(::DbWeight::get().reads(33)) - .saturating_add(::DbWeight::get().writes(18)); + let expected_weight = + <::WeightInfo as SubtensorWeightInfo>::add_stake_burn(); let mut env = MockEnv::new( FunctionId::AddStakeBurnV1, From 0216e94884a053611a198b898b79d002181fe026 Mon Sep 17 00:00:00 2001 From: Landyn Date: Mon, 13 Apr 2026 21:44:22 -0500 Subject: [PATCH 8/9] Charge second-stage weight only when stage is reached For AddStakeRecycleV1 and AddStakeBurnV1, charge add_stake() upfront and only charge the second-stage weight (recycle_alpha()/burn_alpha()) after do_add_stake returns Ok. Atomicity is preserved by keeping both ops inside with_transaction and tracking attempt state via a stack flag. --- chain-extensions/src/lib.rs | 121 ++++++++++++++++++++-------------- chain-extensions/src/tests.rs | 27 +++++--- 2 files changed, 89 insertions(+), 59 deletions(-) diff --git a/chain-extensions/src/lib.rs b/chain-extensions/src/lib.rs index f2ef5bd96a..3b9d408086 100644 --- a/chain-extensions/src/lib.rs +++ b/chain-extensions/src/lib.rs @@ -586,40 +586,48 @@ where } } FunctionId::AddStakeRecycleV1 => { - let weight = - <::WeightInfo as SubtensorWeightInfo>::add_stake() - .saturating_add( - <::WeightInfo as SubtensorWeightInfo>::recycle_alpha(), - ); + let add_stake_weight = + <::WeightInfo as SubtensorWeightInfo>::add_stake(); + let recycle_weight = + <::WeightInfo as SubtensorWeightInfo>::recycle_alpha(); - env.charge_weight(weight)?; + env.charge_weight(add_stake_weight)?; let (hotkey, netuid, tao_amount): (T::AccountId, NetUid, TaoBalance) = env.read_as()?; let caller = env.caller(); - let result = transactional::with_transaction(|| { - let alpha = match pallet_subtensor::Pallet::::do_add_stake( - RawOrigin::Signed(caller.clone()).into(), - hotkey.clone(), - netuid, - tao_amount, - ) { - Ok(a) => a, - Err(e) => return TransactionOutcome::Rollback(Err(e)), - }; - - match pallet_subtensor::Pallet::::do_recycle_alpha( - RawOrigin::Signed(caller).into(), - hotkey, - alpha, - netuid, - ) { - Ok(real_alpha) => TransactionOutcome::Commit(Ok(real_alpha)), - Err(e) => TransactionOutcome::Rollback(Err(e)), - } - }); + let mut recycle_attempted = false; + + let result: Result = + transactional::with_transaction(|| { + let alpha = match pallet_subtensor::Pallet::::do_add_stake( + RawOrigin::Signed(caller.clone()).into(), + hotkey.clone(), + netuid, + tao_amount, + ) { + Ok(a) => a, + Err(e) => return TransactionOutcome::Rollback(Err(e)), + }; + + recycle_attempted = true; + + match pallet_subtensor::Pallet::::do_recycle_alpha( + RawOrigin::Signed(caller).into(), + hotkey, + alpha, + netuid, + ) { + Ok(real_alpha) => TransactionOutcome::Commit(Ok(real_alpha)), + Err(e) => TransactionOutcome::Rollback(Err(e)), + } + }); + + if recycle_attempted { + env.charge_weight(recycle_weight)?; + } match result { Ok(alpha) => { @@ -634,37 +642,48 @@ where } } FunctionId::AddStakeBurnV1 => { - let weight = - <::WeightInfo as SubtensorWeightInfo>::add_stake_burn(); + let add_stake_weight = + <::WeightInfo as SubtensorWeightInfo>::add_stake(); + let burn_weight = + <::WeightInfo as SubtensorWeightInfo>::burn_alpha(); - env.charge_weight(weight)?; + env.charge_weight(add_stake_weight)?; let (hotkey, netuid, tao_amount): (T::AccountId, NetUid, TaoBalance) = env.read_as()?; let caller = env.caller(); - let result = transactional::with_transaction(|| { - let alpha = match pallet_subtensor::Pallet::::do_add_stake( - RawOrigin::Signed(caller.clone()).into(), - hotkey.clone(), - netuid, - tao_amount, - ) { - Ok(a) => a, - Err(e) => return TransactionOutcome::Rollback(Err(e)), - }; - - match pallet_subtensor::Pallet::::do_burn_alpha( - RawOrigin::Signed(caller).into(), - hotkey, - alpha, - netuid, - ) { - Ok(real_alpha) => TransactionOutcome::Commit(Ok(real_alpha)), - Err(e) => TransactionOutcome::Rollback(Err(e)), - } - }); + let mut burn_attempted = false; + + let result: Result = + transactional::with_transaction(|| { + let alpha = match pallet_subtensor::Pallet::::do_add_stake( + RawOrigin::Signed(caller.clone()).into(), + hotkey.clone(), + netuid, + tao_amount, + ) { + Ok(a) => a, + Err(e) => return TransactionOutcome::Rollback(Err(e)), + }; + + burn_attempted = true; + + match pallet_subtensor::Pallet::::do_burn_alpha( + RawOrigin::Signed(caller).into(), + hotkey, + alpha, + netuid, + ) { + Ok(real_alpha) => TransactionOutcome::Commit(Ok(real_alpha)), + Err(e) => TransactionOutcome::Rollback(Err(e)), + } + }); + + if burn_attempted { + env.charge_weight(burn_weight)?; + } match result { Ok(alpha) => { diff --git a/chain-extensions/src/tests.rs b/chain-extensions/src/tests.rs index aef12f1e0d..2b9da523b2 100644 --- a/chain-extensions/src/tests.rs +++ b/chain-extensions/src/tests.rs @@ -1014,7 +1014,10 @@ fn add_stake_burn_success_atomically_stakes_and_burns() { let alpha_out_before = pallet_subtensor::SubnetAlphaOut::::get(netuid); let expected_weight = - <::WeightInfo as SubtensorWeightInfo>::add_stake_burn(); + <::WeightInfo as SubtensorWeightInfo>::add_stake() + .saturating_add( + <::WeightInfo as SubtensorWeightInfo>::burn_alpha(), + ); let mut env = MockEnv::new( FunctionId::AddStakeBurnV1, @@ -1062,11 +1065,10 @@ fn add_stake_recycle_with_insufficient_balance_returns_error() { // Don't fund the coldkey - should fail with balance error + // add_stake fails early, so only add_stake weight should be charged — + // recycle_alpha weight is not charged because that stage is never reached. let expected_weight = - <::WeightInfo as SubtensorWeightInfo>::add_stake() - .saturating_add( - <::WeightInfo as SubtensorWeightInfo>::recycle_alpha(), - ); + <::WeightInfo as SubtensorWeightInfo>::add_stake(); let mut env = MockEnv::new( FunctionId::AddStakeRecycleV1, @@ -1082,6 +1084,7 @@ fn add_stake_recycle_with_insufficient_balance_returns_error() { } _ => panic!("unexpected return value"), } + assert_eq!(env.charged_weight(), Some(expected_weight)); }); } @@ -1283,14 +1286,19 @@ impl SubtensorExtensionEnv for MockEnv { } fn charge_weight(&mut self, weight: Weight) -> Result<(), DispatchError> { + let cumulative = self + .charged_weight + .unwrap_or_default() + .saturating_add(weight); if let Some(expected) = self.expected_weight - && weight != expected + && (cumulative.ref_time() > expected.ref_time() + || cumulative.proof_size() > expected.proof_size()) { return Err(DispatchError::Other( "unexpected weight charged by mock env", )); } - self.charged_weight = Some(weight); + self.charged_weight = Some(cumulative); Ok(()) } @@ -1422,7 +1430,10 @@ fn add_stake_burn_rollback_on_burn_failure() { ); let expected_weight = - <::WeightInfo as SubtensorWeightInfo>::add_stake_burn(); + <::WeightInfo as SubtensorWeightInfo>::add_stake() + .saturating_add( + <::WeightInfo as SubtensorWeightInfo>::burn_alpha(), + ); let mut env = MockEnv::new( FunctionId::AddStakeBurnV1, From 089a24aa95d697b5dd06e177e21456e271d9148b Mon Sep 17 00:00:00 2001 From: Landyn Date: Mon, 13 Apr 2026 22:57:53 -0500 Subject: [PATCH 9/9] Fix CI: rustfmt wrap and forbid-saturating-math in mock - cargo fmt wraps the add_stake() call in AddStakeRecycleV1/AddStakeBurnV1 - mock charge_weight replaces saturating_add with checked_add().unwrap() to satisfy the ForbidSaturatingMath custom lint --- chain-extensions/src/lib.rs | 6 ++++-- chain-extensions/src/tests.rs | 9 +++++---- 2 files changed, 9 insertions(+), 6 deletions(-) diff --git a/chain-extensions/src/lib.rs b/chain-extensions/src/lib.rs index 3b9d408086..484b2af35d 100644 --- a/chain-extensions/src/lib.rs +++ b/chain-extensions/src/lib.rs @@ -587,7 +587,8 @@ where } FunctionId::AddStakeRecycleV1 => { let add_stake_weight = - <::WeightInfo as SubtensorWeightInfo>::add_stake(); + <::WeightInfo as SubtensorWeightInfo>::add_stake( + ); let recycle_weight = <::WeightInfo as SubtensorWeightInfo>::recycle_alpha(); @@ -643,7 +644,8 @@ where } FunctionId::AddStakeBurnV1 => { let add_stake_weight = - <::WeightInfo as SubtensorWeightInfo>::add_stake(); + <::WeightInfo as SubtensorWeightInfo>::add_stake( + ); let burn_weight = <::WeightInfo as SubtensorWeightInfo>::burn_alpha(); diff --git a/chain-extensions/src/tests.rs b/chain-extensions/src/tests.rs index 2b9da523b2..e749d64ffc 100644 --- a/chain-extensions/src/tests.rs +++ b/chain-extensions/src/tests.rs @@ -1286,10 +1286,11 @@ impl SubtensorExtensionEnv for MockEnv { } fn charge_weight(&mut self, weight: Weight) -> Result<(), DispatchError> { - let cumulative = self - .charged_weight - .unwrap_or_default() - .saturating_add(weight); + let prev = self.charged_weight.unwrap_or_default(); + let cumulative = Weight::from_parts( + prev.ref_time().checked_add(weight.ref_time()).unwrap(), + prev.proof_size().checked_add(weight.proof_size()).unwrap(), + ); if let Some(expected) = self.expected_weight && (cumulative.ref_time() > expected.ref_time() || cumulative.proof_size() > expected.proof_size())