diff --git a/chain-extensions/src/lib.rs b/chain-extensions/src/lib.rs index 14ea23d9c8..484b2af35d 100644 --- a/chain-extensions/src/lib.rs +++ b/chain-extensions/src/lib.rs @@ -9,11 +9,13 @@ 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::{ 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}; @@ -523,6 +525,180 @@ where Ok(RetVal::Converging(Output::Success as u32)) } + FunctionId::RecycleAlphaV1 => { + let weight = + <::WeightInfo as SubtensorWeightInfo>::recycle_alpha(); + + env.charge_weight(weight)?; + + let (hotkey, amount, netuid): (T::AccountId, AlphaBalance, NetUid) = + env.read_as()?; + + let caller = env.caller(); + + let call_result = pallet_subtensor::Pallet::::do_recycle_alpha( + RawOrigin::Signed(caller).into(), + hotkey, + amount, + netuid, + ); + + match call_result { + Ok(real_amount) => { + env.write_output(&real_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 = + <::WeightInfo as SubtensorWeightInfo>::burn_alpha(); + + env.charge_weight(weight)?; + + let (hotkey, amount, netuid): (T::AccountId, AlphaBalance, NetUid) = + env.read_as()?; + + let caller = env.caller(); + + let call_result = pallet_subtensor::Pallet::::do_burn_alpha( + RawOrigin::Signed(caller).into(), + hotkey, + amount, + netuid, + ); + + match call_result { + Ok(real_amount) => { + env.write_output(&real_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 => { + let add_stake_weight = + <::WeightInfo as SubtensorWeightInfo>::add_stake( + ); + let recycle_weight = + <::WeightInfo as SubtensorWeightInfo>::recycle_alpha(); + + env.charge_weight(add_stake_weight)?; + + let (hotkey, netuid, tao_amount): (T::AccountId, NetUid, TaoBalance) = + env.read_as()?; + + let caller = env.caller(); + + 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) => { + 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)) + } + } + } + FunctionId::AddStakeBurnV1 => { + let add_stake_weight = + <::WeightInfo as SubtensorWeightInfo>::add_stake( + ); + let burn_weight = + <::WeightInfo as SubtensorWeightInfo>::burn_alpha(); + + env.charge_weight(add_stake_weight)?; + + let (hotkey, netuid, tao_amount): (T::AccountId, NetUid, TaoBalance) = + env.read_as()?; + + let caller = env.caller(); + + 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) => { + 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)) + } + } + } } } } diff --git a/chain-extensions/src/tests.rs b/chain-extensions/src/tests.rs index b8956e8659..e749d64ffc 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; @@ -726,6 +727,533 @@ 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 = + <::WeightInfo as SubtensorWeightInfo>::recycle_alpha(); + + 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 = + <::WeightInfo as SubtensorWeightInfo>::recycle_alpha(); + + 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 = + <::WeightInfo as SubtensorWeightInfo>::burn_alpha(); + + 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 = + <::WeightInfo as SubtensorWeightInfo>::burn_alpha(); + + 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 = + <::WeightInfo as SubtensorWeightInfo>::add_stake() + .saturating_add( + <::WeightInfo as SubtensorWeightInfo>::recycle_alpha(), + ); + + 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 = + <::WeightInfo as SubtensorWeightInfo>::add_stake() + .saturating_add( + <::WeightInfo as SubtensorWeightInfo>::burn_alpha(), + ); + + 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 + + // 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(); + + 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"), + } + assert_eq!(env.charged_weight(), Some(expected_weight)); + }); +} + +#[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 = + <::WeightInfo as SubtensorWeightInfo>::recycle_alpha(); + + 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 return actual clamped amount, not requested amount" + ); + + 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 = + <::WeightInfo as SubtensorWeightInfo>::burn_alpha(); + + 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 = + <::WeightInfo as SubtensorWeightInfo>::burn_alpha(); + + 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 return actual clamped amount, not requested amount" + ); + + 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 { @@ -758,14 +1286,20 @@ impl SubtensorExtensionEnv for MockEnv { } fn charge_weight(&mut self, weight: Weight) -> Result<(), DispatchError> { + 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 - && 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(()) } @@ -793,6 +1327,148 @@ 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 = + <::WeightInfo as SubtensorWeightInfo>::add_stake() + .saturating_add( + <::WeightInfo as SubtensorWeightInfo>::recycle_alpha(), + ); + + 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 = + <::WeightInfo as SubtensorWeightInfo>::add_stake() + .saturating_add( + <::WeightInfo as SubtensorWeightInfo>::burn_alpha(), + ); + + 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(|| { 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 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,