@@ -21,6 +21,10 @@ impl<T: Config> Pallet<T> {
2121 SubnetAlphaIn :: < T > :: get ( netuid) . saturating_add ( SubnetAlphaOut :: < T > :: get ( netuid) )
2222 }
2323
24+ pub fn get_protocol_tao ( netuid : NetUid ) -> TaoCurrency {
25+ T :: SwapInterface :: get_protocol_tao ( netuid)
26+ }
27+
2428 pub fn get_moving_alpha_price ( netuid : NetUid ) -> U96F32 {
2529 let one = U96F32 :: saturating_from_num ( 1.0 ) ;
2630 if netuid. is_root ( ) {
@@ -688,6 +692,9 @@ impl<T: Config> Pallet<T> {
688692 price_limit : TaoCurrency ,
689693 drop_fees : bool ,
690694 ) -> Result < TaoCurrency , DispatchError > {
695+ // Record the protocol TAO before the swap.
696+ let protocol_tao = Self :: get_protocol_tao ( netuid) ;
697+
691698 // Decrease alpha on subnet
692699 let actual_alpha_decrease =
693700 Self :: decrease_stake_for_hotkey_and_coldkey_on_subnet ( hotkey, coldkey, netuid, alpha) ;
@@ -696,6 +703,13 @@ impl<T: Config> Pallet<T> {
696703 let swap_result =
697704 Self :: swap_alpha_for_tao ( netuid, actual_alpha_decrease, price_limit, drop_fees) ?;
698705
706+ // Record the protocol TAO after the swap.
707+ let protocol_tao_after = Self :: get_protocol_tao ( netuid) ;
708+ // This should decrease as we are removing TAO from the protocol.
709+ let protocol_tao_delta: TaoCurrency = protocol_tao. saturating_sub ( protocol_tao_after) ;
710+ // Use max to overstate the TAO flow from the protocol.
711+ let tao_flow = protocol_tao_delta. max ( swap_result. amount_paid_out . into ( ) ) ;
712+
699713 // Refund the unused alpha (in case if limit price is hit)
700714 let refund = actual_alpha_decrease. saturating_sub (
701715 swap_result
@@ -722,7 +736,7 @@ impl<T: Config> Pallet<T> {
722736 // }
723737
724738 // Record TAO outflow
725- Self :: record_tao_outflow ( netuid, swap_result . amount_paid_out . into ( ) ) ;
739+ Self :: record_tao_outflow ( netuid, tao_flow ) ;
726740
727741 LastColdkeyHotkeyStakeBlock :: < T > :: insert ( coldkey, hotkey, Self :: get_current_block_as_u64 ( ) ) ;
728742
@@ -761,9 +775,20 @@ impl<T: Config> Pallet<T> {
761775 set_limit : bool ,
762776 drop_fees : bool ,
763777 ) -> Result < AlphaCurrency , DispatchError > {
778+ // Record the protocol TAO before the swap.
779+ let protocol_tao = Self :: get_protocol_tao ( netuid) ;
780+
764781 // Swap the tao to alpha.
765782 let swap_result = Self :: swap_tao_for_alpha ( netuid, tao, price_limit, drop_fees) ?;
766783
784+ // Record the protocol TAO after the swap.
785+ let protocol_tao_after = Self :: get_protocol_tao ( netuid) ;
786+
787+ // This should increase as we are adding TAO to the protocol.
788+ let protocol_tao_delta: TaoCurrency = protocol_tao_after. saturating_sub ( protocol_tao) ;
789+ // Use min to understate the TAO flow into the protocol.
790+ let tao_flow = protocol_tao_delta. min ( tao) ;
791+
767792 ensure ! (
768793 !swap_result. amount_paid_out. is_zero( ) ,
769794 Error :: <T >:: AmountTooLow
@@ -799,7 +824,7 @@ impl<T: Config> Pallet<T> {
799824 }
800825
801826 // Record TAO inflow
802- Self :: record_tao_inflow ( netuid, swap_result . amount_paid_in . into ( ) ) ;
827+ Self :: record_tao_inflow ( netuid, tao_flow ) ;
803828
804829 LastColdkeyHotkeyStakeBlock :: < T > :: insert ( coldkey, hotkey, Self :: get_current_block_as_u64 ( ) ) ;
805830
0 commit comments