[WIP] Feat/balancer swap updates

bittensor_cli/cli.py

@@ -866,7 +866,7 @@ def __init__(self):
         self.subnet_mechanisms_app = typer.Typer(epilog=_epilog)
         self.weights_app = typer.Typer(epilog=_epilog)
         self.view_app = typer.Typer(epilog=_epilog)
-        self.liquidity_app = typer.Typer(epilog=_epilog)
+        self.liquidity_app = typer.Typer(epilog=_epilog, hidden=True)
         self.crowd_app = typer.Typer(epilog=_epilog)
         self.utils_app = typer.Typer(epilog=_epilog)
         self.axon_app = typer.Typer(epilog=_epilog)
@@ -8557,6 +8557,9 @@ def liquidity_add(
         json_output: bool = Options.json_output,
     ):
         """Add liquidity to the swap (as a combination of TAO + Alpha)."""
+        console.print_error("User liquidity is currently disabled on Bittensor.")
+        raise typer.Exit()
+
         self.verbosity_handler(quiet, verbose, json_output, prompt, decline)
         proxy = self.is_valid_proxy_name_or_ss58(proxy, announce_only)
         if not netuid:
@@ -8634,6 +8637,9 @@ def liquidity_list(
         json_output: bool = Options.json_output,
     ):
         """Displays liquidity positions in given subnet."""
+        console.print_error("User liquidity is currently disabled on Bittensor.")
+        raise typer.Exit()
+
         self.verbosity_handler(quiet, verbose, json_output, prompt=False)
         if not netuid:
             netuid = IntPrompt.ask(
@@ -8687,6 +8693,8 @@ def liquidity_remove(
     ):
         """Remove liquidity from the swap (as a combination of TAO + Alpha)."""
 
+        console.print_error("User liquidity is currently disabled on Bittensor.")
+        raise typer.Exit()
         self.verbosity_handler(quiet, verbose, json_output, prompt, decline)
         proxy = self.is_valid_proxy_name_or_ss58(proxy, announce_only)
         if all_liquidity_ids and position_id:
@@ -8763,6 +8771,8 @@ def liquidity_modify(
         json_output: bool = Options.json_output,
     ):
         """Modifies the liquidity position for the given subnet."""
+        console.print_error("User liquidity is currently disabled on Bittensor.")
+        raise typer.Exit()
         self.verbosity_handler(quiet, verbose, json_output, prompt, decline)
         proxy = self.is_valid_proxy_name_or_ss58(proxy, announce_only)
         if not netuid:

bittensor_cli/src/bittensor/chain_data.py

@@ -800,88 +800,6 @@ def tao_to_alpha(self, tao: Balance) -> Balance:
     def alpha_to_tao(self, alpha: Balance) -> Balance:
         return Balance.from_tao(alpha.tao * self.price.tao)
 
-    def tao_to_alpha_with_slippage(
-        self, tao: Balance
-    ) -> tuple[Balance, Balance, float]:
-        """
-        Returns an estimate of how much Alpha a staker would receive if they stake their tao using the current pool
-            state.
-
-        Args:
-            tao: Amount of TAO to stake.
-        Returns:
-            Tuple of balances where the first part is the amount of Alpha received, and the
-            second part (slippage) is the difference between the estimated amount and ideal
-            amount as if there was no slippage
-        """
-        if self.is_dynamic:
-            new_tao_in = self.tao_in + tao
-            if new_tao_in == 0:
-                return tao, Balance.from_rao(0)
-            new_alpha_in = self.k / new_tao_in
-
-            # Amount of alpha given to the staker
-            alpha_returned = Balance.from_rao(
-                self.alpha_in.rao - new_alpha_in.rao
-            ).set_unit(self.netuid)
-
-            # Ideal conversion as if there is no slippage, just price
-            alpha_ideal = self.tao_to_alpha(tao)
-
-            if alpha_ideal.tao > alpha_returned.tao:
-                slippage = Balance.from_tao(
-                    alpha_ideal.tao - alpha_returned.tao
-                ).set_unit(self.netuid)
-            else:
-                slippage = Balance.from_tao(0)
-        else:
-            alpha_returned = tao.set_unit(self.netuid)
-            slippage = Balance.from_tao(0)
-
-        slippage_pct_float = (
-            100 * float(slippage) / float(slippage + alpha_returned)
-            if slippage + alpha_returned != 0
-            else 0
-        )
-        return alpha_returned, slippage, slippage_pct_float
-
-    def alpha_to_tao_with_slippage(
-        self, alpha: Balance
-    ) -> tuple[Balance, Balance, float]:
-        """
-        Returns an estimate of how much TAO a staker would receive if they unstake their alpha using the current pool
-            state.
-
-        Args:
-            alpha: Amount of Alpha to stake.
-        Returns:
-            Tuple of balances where the first part is the amount of TAO received, and the
-            second part (slippage) is the difference between the estimated amount and ideal
-            amount as if there was no slippage
-        """
-        if self.is_dynamic:
-            new_alpha_in = self.alpha_in + alpha
-            new_tao_reserve = self.k / new_alpha_in
-            # Amount of TAO given to the unstaker
-            tao_returned = Balance.from_rao(self.tao_in - new_tao_reserve)
-
-            # Ideal conversion as if there is no slippage, just price
-            tao_ideal = self.alpha_to_tao(alpha)
-
-            if tao_ideal > tao_returned:
-                slippage = Balance.from_tao(tao_ideal.tao - tao_returned.tao)
-            else:
-                slippage = Balance.from_tao(0)
-        else:
-            tao_returned = alpha.set_unit(0)
-            slippage = Balance.from_tao(0)
-        slippage_pct_float = (
-            100 * float(slippage) / float(slippage + tao_returned)
-            if slippage + tao_returned != 0
-            else 0
-        )
-        return tao_returned, slippage, slippage_pct_float
-
 
 @dataclass
 class ColdkeySwapAnnouncementInfo(InfoBase):
@@ -1228,6 +1146,78 @@ def from_dict(cls, d: dict, netuid: int) -> "SimSwapResult":
             alpha_fee=Balance.from_rao(d["alpha_fee"]).set_unit(netuid),
         )
 
+    def tao_to_alpha_slippage(
+        self,
+        tao_amount: Balance,
+        current_price: float,
+        netuid: int,
+    ) -> tuple[Balance, Balance, float]:
+        """
+        Calculate slippage for a TAO -> Alpha swap.
+
+        Args:
+            tao_amount: Amount of TAO provided as input.
+            current_price: Current alpha price in TAO (TAO per 1 alpha).
+            netuid: Target subnet netuid (used for unit tagging).
+
+        Returns:
+            A tuple of:
+                received_alpha (Balance): Simulated alpha received.
+                slippage_alpha (Balance): Shortfall vs ideal at current_price.
+                slippage_pct_float (float): Slippage percentage (0 - 100).
+        """
+        if current_price <= 0:
+            zero = Balance.from_tao(0).set_unit(netuid)
+            return zero, zero, 0.0
+
+        ideal_amount = Balance.from_tao(tao_amount.tao / current_price).set_unit(netuid)
+        received_amount = self.alpha_amount
+
+        if ideal_amount.tao == 0:
+            zero = Balance.from_tao(0).set_unit(netuid)
+            return received_amount, zero, 0.0
+
+        slippage_amount = max(ideal_amount.tao - received_amount.tao, 0)
+        slippage_amount_balance = Balance.from_tao(slippage_amount).set_unit(netuid)
+        slippage_pct = 100 * slippage_amount / ideal_amount.tao
+
+        return received_amount, slippage_amount_balance, slippage_pct
+
+    def alpha_to_tao_slippage(
+        self,
+        alpha_amount: Balance,
+        current_price: float,
+    ) -> tuple[Balance, Balance, float]:
+        """
+        Calculate slippage for an Alpha -> TAO swap.
+
+        Args:
+            alpha_amount: Amount of Alpha provided as input.
+            current_price: Current alpha price in TAO (TAO per 1 alpha).
+
+        Returns:
+            A tuple of:
+                received_tao (Balance): Simulated TAO received.
+                slippage_tao (Balance): Shortfall vs ideal at current_price.
+                slippage_pct_float (float): Slippage percentage (0 - 100).
+        """
+        if current_price <= 0:
+            zero = Balance.from_tao(0).set_unit(0)
+            return zero, zero, 0.0
+
+        ideal_amount = Balance.from_tao(alpha_amount.tao * current_price).set_unit(0)
+        received_amount = self.tao_amount
+
+        if ideal_amount.tao == 0:
+            zero = Balance.from_tao(0).set_unit(0)
+            return received_amount, zero, 0.0
+
+        slippage_amount = max(ideal_amount.tao - received_amount.tao, 0)
+        slippage_amount_balance = Balance.from_tao(slippage_amount).set_unit(0)
+        slippage_pct = 100 * slippage_amount / ideal_amount.tao
+
+        return received_amount, slippage_amount_balance, slippage_pct
+
 
 @dataclass
 class CrowdloanData(InfoBase):

bittensor_cli/src/bittensor/subtensor_interface.py

@@ -444,40 +444,52 @@ async def get_total_stake_for_coldkey(
 
         :return: {address: Balance objects}
         """
-        sub_stakes, dynamic_info = await asyncio.gather(
+        sub_stakes, price_map = await asyncio.gather(
             self.get_stake_for_coldkeys(list(ss58_addresses), block_hash=block_hash),
-            # Token pricing info
-            self.all_subnets(block_hash=block_hash),
+            self.get_subnet_prices(block_hash=block_hash),
         )
 
-        results = {}
+        results: dict[str, tuple[Balance, Balance]] = {}
+        dynamic_stakes: list[tuple[str, "StakeInfo"]] = []
+
         for ss58, stake_info_list in sub_stakes.items():
-            total_tao_value = Balance(0)
-            total_swapped_tao_value = Balance(0)
+            total_tao_value, total_swapped_tao_value = Balance(0), Balance(0)
             for sub_stake in stake_info_list:
                 if sub_stake.stake.rao == 0:
                     continue
-                netuid = sub_stake.netuid
-                pool = dynamic_info[netuid]
 
-                alpha_value = Balance.from_rao(int(sub_stake.stake.rao)).set_unit(
-                    netuid
+                netuid = sub_stake.netuid
+                price = price_map[netuid]
+                ideal_tao = Balance.from_tao(sub_stake.stake.tao * price.tao).set_unit(
+                    0
                 )
+                total_tao_value += ideal_tao
 
-                # Without slippage
-                tao_value = pool.alpha_to_tao(alpha_value)
-                total_tao_value += tao_value
-
-                # With slippage
                 if netuid == 0:
-                    swapped_tao_value = tao_value
+                    total_swapped_tao_value += ideal_tao
                 else:
-                    swapped_tao_value, _, _ = pool.alpha_to_tao_with_slippage(
-                        sub_stake.stake
-                    )
-                total_swapped_tao_value += swapped_tao_value
+                    dynamic_stakes.append((ss58, sub_stake))
 
             results[ss58] = (total_tao_value, total_swapped_tao_value)
+
+        if dynamic_stakes:
+            sim_results = await asyncio.gather(
+                *[
+                    self.sim_swap(
+                        origin_netuid=sub_stake.netuid,
+                        destination_netuid=0,
+                        amount=sub_stake.stake.rao,
+                        block_hash=block_hash,
+                    )
+                    for _, sub_stake in dynamic_stakes
+                ]
+            )
+
+            for (ss58, sub_stake), sim_result in zip(dynamic_stakes, sim_results):
+                total_tao_value, total_swapped_tao_value = results[ss58]
+                total_swapped_tao_value += sim_result.tao_amount
+                results[ss58] = (total_tao_value, total_swapped_tao_value)
+
         return results
 
     async def get_total_stake_for_hotkey(
@@ -1654,7 +1666,7 @@ async def all_subnets(self, block_hash: Optional[str] = None) -> list[DynamicInf
                 "get_all_dynamic_info",
                 block_hash=block_hash,
             ),
-            self.get_subnet_prices(block_hash=block_hash, page_size=129),
+            self.get_subnet_prices(block_hash=block_hash),
         )
         sns: list[DynamicInfo] = DynamicInfo.list_from_any(result)
         for sn in sns:
@@ -2523,43 +2535,46 @@ async def get_subnet_price(
 
         :return: The current Alpha price in TAO units for the specified subnet.
         """
-        # TODO update this to use the runtime call SwapRuntimeAPI.current_alpha_price
-        current_sqrt_price = await self.query(
-            module="Swap",
-            storage_function="AlphaSqrtPrice",
-            params=[netuid],
+        if netuid == 0:
+            return Balance.from_tao(1.0)
+
+        current_price = await self.query_runtime_api(
+            "SwapRuntimeApi",
+            "current_alpha_price",
+            params={"netuid": netuid},
             block_hash=block_hash,
         )
-
-        current_sqrt_price = fixed_to_float(current_sqrt_price)
-        current_price = current_sqrt_price * current_sqrt_price
-        return Balance.from_rao(int(current_price * 1e9))
+        return Balance.from_rao(current_price)
 
     async def get_subnet_prices(
-        self, block_hash: Optional[str] = None, page_size: int = 100
+        self, block_hash: Optional[str] = None
     ) -> dict[int, Balance]:
         """
         Gets the current Alpha prices in TAO for all subnets.
 
         :param block_hash: The hash of the block to retrieve prices from.
-        :param page_size: The page size for batch queries (default: 100).
 
         :return: A dictionary mapping netuid to the current Alpha price in TAO units.
         """
-        query = await self.substrate.query_map(
-            module="Swap",
-            storage_function="AlphaSqrtPrice",
-            page_size=page_size,
-            block_hash=block_hash,
-        )
+        all_netuids = await self.get_all_subnet_netuids(block_hash=block_hash)
 
-        map_ = {}
-        async for netuid_, current_sqrt_price in query:
-            current_sqrt_price_ = fixed_to_float(current_sqrt_price.value)
-            current_price = current_sqrt_price_**2
-            map_[netuid_] = Balance.from_rao(int(current_price * 1e9))
+        result = {0: Balance.from_tao(1.0)}
+        netuids_to_query = [netuid for netuid in all_netuids if netuid != 0]
+        prices = await asyncio.gather(
+            *[
+                self.query_runtime_api(
+                    "SwapRuntimeApi",
+                    "current_alpha_price",
+                    params={"netuid": netuid},
+                    block_hash=block_hash,
+                )
+                for netuid in netuids_to_query
+            ],
+        )
+        for netuid, current_price in zip(netuids_to_query, prices):
+            result[netuid] = Balance.from_rao(current_price)
 
-        return map_
+        return result
 
     async def get_all_subnet_ema_tao_inflow(
         self,