Enable _tweak_price for 3-coin Cryptopools

curvesim/pool/cryptoswap/pool.py

@@ -1,3 +1,4 @@
+# pylint: disable=too-many-lines
 """
 Mainly a module to house the `CurveCryptoPool`, a cryptoswap implementation in Python.
 """
@@ -294,20 +295,26 @@ def _tweak_price(  # noqa: complexity: 12
         A: int,
         gamma: int,
         _xp: List[int],
-        i: int,
+        i: Optional[int],
         p_i: Optional[int],
-        new_D: int,
+        new_D: Optional[int],
         K0_prev: int = 0,
     ) -> None:
         """
         Applies several kinds of updates:
             - EMA price update: price_oracle
             - Profit counters: D, virtual_price, xcp_profit
             - price adjustment: price_scale
-                - If p_i is None, the spot price will be used as the last price
 
         Also claims admin fees if appropriate (enough profit and price scale
-        and oracle is close enough).
+        and oracle are close enough).
+
+        Note
+        -----
+        - Always pass in p_i and i as None if self.n == 3. i and p_i are only
+        meaningful when self.n == 2 because all last prices are calculated at once when
+        self.n == 3.
+        - If p_i is None, the spot price(s) will be used as the last price(s).
         """
         price_oracle: List[int] = self._price_oracle
         last_prices: List[int] = self.last_prices
@@ -342,9 +349,12 @@ def _tweak_price(  # noqa: complexity: 12
             self._price_oracle = price_oracle
             self.last_prices_timestamp = block_timestamp
 
-        D_unadjusted: int = new_D  # Withdrawal methods know new D already
-        if new_D == 0:
-            D_unadjusted = newton_D(A, gamma, _xp, K0_prev)
+        if new_D is None:
+            D_unadjusted: int = newton_D(A, gamma, _xp, K0_prev)
+        elif new_D > 0:
+            D_unadjusted = new_D  # Withdrawal methods know new D already
+        else:
+            raise CalculationError(f"new_D cannot be {new_D}.")
 
         if p_i is None:
             if n_coins == 2:
@@ -361,20 +371,29 @@ def _tweak_price(  # noqa: complexity: 12
                     )
                     for k in range(1, n_coins)
                 ]
-            else:
+            elif n_coins == 3:
                 last_prices = get_p(_xp, D_unadjusted, A, gamma)
                 last_prices = [
                     last_p * p // 10**18
                     for last_p, p in zip(last_prices, price_scale)
                 ]
         elif p_i > 0:
-            # Save the last price
-            if i > 0:
-                last_prices[i - 1] = p_i
-            else:
-                # If 0th price changed - change all prices instead
-                for k in range(n_coins - 1):
-                    last_prices[k] = last_prices[k] * 10**18 // p_i
+            if n_coins == 2:
+                if i is None:
+                    raise CalculationError(
+                        "i cannot be None when p_i is provided for 2-coin pools."
+                    )
+                if i > 0:
+                    # Save the last price
+                    last_prices[i - 1] = p_i
+                else:
+                    # If 0th price changed - change all prices instead
+                    for k in range(n_coins - 1):
+                        last_prices[k] = last_prices[k] * 10**18 // p_i
+            elif n_coins == 3:
+                raise CalculationError(
+                    "Always pass p_i (last price) as None for 3-coin pools."
+                )
         else:
             raise CalculationError(f"p_i (last price) cannot be {p_i}.")
 
@@ -398,6 +417,10 @@ def _tweak_price(  # noqa: complexity: 12
             if virtual_price < old_virtual_price:
                 raise CryptoPoolError("Loss")
 
+            # 3-coin cryptopool does this, but 2-coin cryptopool doesn't
+            if n_coins == 3 and virtual_price == old_virtual_price:
+                raise CryptoPoolError("Loss")
+
             xcp_profit = old_xcp_profit * virtual_price // old_virtual_price
 
         self.xcp_profit = xcp_profit
@@ -610,7 +633,6 @@ def _exchange(
         A = self.A
         gamma = self.gamma
         xp: List[int] = self.balances.copy()
-        ix: int = j
 
         y: int = xp[j]
         xp[i] += dx
@@ -645,22 +667,22 @@ def _exchange(
         xp[j] = y
 
         p: Optional[int] = None
+        ix: Optional[int] = None
         K0_prev: int = 0
         if self.n == 2:
             if dx > 10**5 and dy > 10**5:
+                ix = j
                 _dx: int = dx * prec_i
                 _dy: int = dy * prec_j
-                if i != 0 and j != 0:
-                    p = self.last_prices[i - 1] * _dx // _dy
-                elif i == 0:
+                if i == 0:
                     p = _dx * 10**18 // _dy
                 else:  # j == 0
                     p = _dy * 10**18 // _dx
                     ix = i
         else:
             K0_prev = y_out[1]
 
-        self._tweak_price(A, gamma, xp, ix, p, 0, K0_prev)
+        self._tweak_price(A, gamma, xp, ix, p, None, K0_prev)
 
         return dy, fee
 
@@ -768,7 +790,7 @@ def add_liquidity(
             # = sum{k!=i}(p_k * (dtoken / token_supply * xx_k - dx_k))
             # only ix is nonzero
             p: Optional[int] = None
-            ix: int = -1
+            ix: Optional[int] = None
             if d_token > 10**5:
                 nonzero_indices = [i for i, a in enumerate(amounts) if a != 0]
                 if len(nonzero_indices) == 1:
@@ -896,7 +918,11 @@ def remove_liquidity_one_coin(
         self.balances[i] -= dy
         self.tokens -= token_amount
 
-        self._tweak_price(A, gamma, xp, i, p, D)
+        ix: Optional[int] = None
+        if self.n == 2:
+            ix = i
+
+        self._tweak_price(A, gamma, xp, ix, p, D)
 
         return dy
 

test/unit/test_cryptopool.py

@@ -337,6 +337,7 @@ def test_get_y(vyper_cryptopool, A, gamma, x0, x1, i, delta_perc):
 def test_tweak_price(
     vyper_cryptopool, cryptopool_lp_token, A, gamma, x0, x_perc, price_perc
 ):
+    """Test _tweak_price against vyper implementation."""
     # def test_tweak_price(vyper_cryptopool, cryptopool_lp_token):
     #     A = 4196
     #     gamma = 10000050055
@@ -392,7 +393,7 @@ def test_tweak_price(
         pool._tweak_price(A, gamma, xp, 1, 0, 0)
     except Exception as err:
         assert isinstance(err, CalculationError)
-    pool._tweak_price(A, gamma, xp, 1, last_price, 0)
+    pool._tweak_price(A, gamma, xp, 1, last_price, None)
     vyper_cryptopool.eval(f"self.tweak_price({A_gamma}, {xp}, {last_price}, 0)")
 
     assert pool.price_scale == [vyper_cryptopool.price_scale()]
@@ -412,7 +413,7 @@ def test_tweak_price(
     old_scale = pool.price_scale
     old_virtual_price = pool.virtual_price
 
-    pool._tweak_price(A, gamma, xp, 1, last_price, 0)
+    pool._tweak_price(A, gamma, xp, 1, last_price, None)
     vyper_cryptopool.eval(f"self.tweak_price({A_gamma}, {xp}, {last_price}, 0)")
 
     # check the pools are the same
@@ -440,7 +441,7 @@ def test_tweak_price(
     xp[0] = xp[0] + pool.allowed_extra_profit // 10
 
     # omitting price will calculate the spot price in `tweak_price`
-    pool._tweak_price(A, gamma, xp, 1, None, 0)
+    pool._tweak_price(A, gamma, xp, 1, None, None)
     vyper_cryptopool.eval(f"self.tweak_price({A_gamma}, {xp}, 0, 0)")
 
     assert pool.price_scale == [vyper_cryptopool.price_scale()]
@@ -459,7 +460,7 @@ def test_tweak_price(
     xp[0] = xp[0] * 115 // 100
 
     # omitting price will calculate the spot price in `tweak_price`
-    pool._tweak_price(A, gamma, xp, 1, None, 0)
+    pool._tweak_price(A, gamma, xp, 1, None, None)
     vyper_cryptopool.eval(f"self.tweak_price({A_gamma}, {xp}, 0, 0)")
 
     assert pool.price_scale == [vyper_cryptopool.price_scale()]