Update add_liquidity to support 3-coin pools

curvesim/pool/cryptoswap/pool.py

@@ -709,7 +709,7 @@ def exchange_underlying(
         """
         return self.exchange(i, j, dx, min_dy)
 
-    # pylint: disable-next=too-many-locals
+    # pylint: disable-next=too-many-locals, too-many-statements
     def add_liquidity(
         self,
         amounts: List[int],
@@ -730,7 +730,8 @@ def add_liquidity(
         int
             Amount of LP tokens minted.
         """
-        assert amounts[0] > 0 or amounts[1] > 0  # dev: no coins to add
+        assert [x for x in amounts if x >= 0] == amounts  # dev: invalid amounts
+        assert sum(amounts) > 0  # dev: no coins to add
 
         A = self.A
         gamma = self.gamma
@@ -745,7 +746,7 @@ def add_liquidity(
         amountsp: List[int] = [xp[i] - xp_old[i] for i in range(n_coins)]
 
         old_D: int = self.D
-        D: int = factory_2_coin.newton_D(A, gamma, xp)
+        D: int = newton_D(A, gamma, xp)
 
         d_token: int = 0
         token_supply: int = self.tokens
@@ -763,16 +764,15 @@ def add_liquidity(
             token_supply += d_token
             self.tokens += d_token
 
-            # Calculate price:
-            # p_i * (dx_i - dtoken / token_supply * xx_i)
-            # = sum{k!=i}(p_k * (dtoken / token_supply * xx_k - dx_k))
-            # only ix is nonzero
             p: Optional[int] = None
             ix: int = -1
-            if d_token > 10**5:
+            if n_coins == 2 and d_token > 10**5:
                 nonzero_indices = [i for i, a in enumerate(amounts) if a != 0]
                 if len(nonzero_indices) == 1:
-                    # not reached in current tests, which never have nonzero amounts
+                    # Calculate price for 2 coins:
+                    # p_i * (dx_i - dtoken / token_supply * xx_i)
+                    # = sum{k!=i}(p_k * (dtoken / token_supply * xx_k - dx_k))
+                    # only ix is nonzero
                     prec: List[int] = self.precisions
                     last_prices: List[int] = self.last_prices
                     balances: List[int] = self.balances
@@ -803,6 +803,7 @@ def add_liquidity(
             self.D = D
             self.virtual_price = 10**18
             self.xcp_profit = 10**18
+            self.xcp_profit_a = 10**18
             self.tokens += d_token
 
         assert d_token >= min_mint_amount, "Slippage"

test/fixtures/curve/tricrypto_ng.vy

@@ -582,33 +582,34 @@ def add_liquidity(
 
         if amounts[i] > 0:
 
-            if coins[i] == WETH20:
-
-                self._transfer_in(
-                    coins[i],
-                    amounts[i],
-                    0,  # <-----------------------------------
-                    msg.value,  #                             | No callbacks
-                    empty(address),  # <----------------------| for
-                    empty(bytes32),  # <----------------------| add_liquidity.
-                    msg.sender,  #                            |
-                    empty(address),  # <-----------------------
-                    use_eth
-                )
-
-            else:
-
-                self._transfer_in(
-                    coins[i],
-                    amounts[i],
-                    0,
-                    0,  # <----------------- mvalue = 0 if coin is not WETH20.
-                    empty(address),
-                    empty(bytes32),
-                    msg.sender,
-                    empty(address),
-                    False  # <-------- use_eth is False if coin is not WETH20.
-                )
+            # curvesim: comment out unneeded coin transfer logic
+            # if coins[i] == WETH20:
+
+            #     self._transfer_in(
+            #         coins[i],
+            #         amounts[i],
+            #         0,  # <-----------------------------------
+            #         msg.value,  #                             | No callbacks
+            #         empty(address),  # <----------------------| for
+            #         empty(bytes32),  # <----------------------| add_liquidity.
+            #         msg.sender,  #                            |
+            #         empty(address),  # <-----------------------
+            #         use_eth
+            #     )
+
+            # else:
+
+            #     self._transfer_in(
+            #         coins[i],
+            #         amounts[i],
+            #         0,
+            #         0,  # <----------------- mvalue = 0 if coin is not WETH20.
+            #         empty(address),
+            #         empty(bytes32),
+            #         msg.sender,
+            #         empty(address),
+            #         False  # <-------- use_eth is False if coin is not WETH20.
+            #     )
 
             amountsp[i] = xp[i] - xp_old[i]
 

test/unit/test_tricrypto.py

@@ -585,6 +585,38 @@ def test_calc_withdraw_one_coin(vyper_tricrypto, amount, i):
     assert pool.balances == expected_balances
 
 
+@given(positive_balance, positive_balance, positive_balance)
+@settings(
+    suppress_health_check=[HealthCheck.function_scoped_fixture],
+    max_examples=5,
+    deadline=None,
+)
+def test_add_liquidity(vyper_tricrypto, x0, x1, x2):
+    n_coins = 3
+    assume(0.02 < x0 / x1 < 50)
+    assume(0.02 < x0 / x2 < 50)
+    assume(0.02 < x1 / x2 < 50)
+    xp = [x0, x1, x2]
+
+    precisions = vyper_tricrypto.precisions()
+    price_scale = [vyper_tricrypto.price_scale(i) for i in range(n_coins - 1)]
+    amounts = get_real_balances(xp, precisions, price_scale)
+
+    pool = initialize_pool(vyper_tricrypto)
+
+    expected_lp_amount = vyper_tricrypto.add_liquidity(amounts, 0)
+    expected_balances = [vyper_tricrypto.balances(i) for i in range(n_coins)]
+    expected_lp_supply = vyper_tricrypto.totalSupply()
+    expected_D = vyper_tricrypto.D()
+
+    lp_amount = pool.add_liquidity(amounts)
+
+    assert lp_amount == expected_lp_amount
+    assert pool.balances == expected_balances
+    assert pool.tokens == expected_lp_supply
+    assert pool.D == expected_D
+
+
 def test_claim_admin_fees(vyper_tricrypto, tricrypto_math):
     """Test admin fee claim against vyper implementation."""
     update_cached_values(vyper_tricrypto, tricrypto_math)