Adding remove_liquidity code/tests

curvesim/pool/stableswap/pool.py

@@ -479,6 +479,86 @@ def remove_liquidity_one_coin(self, token_amount, i):
         self.tokens -= token_amount
         return dy, dy_fee
 
+    def remove_liquidity(self, _amount, min_amounts=None):
+        """
+        Remove liquidity (burn LP tokens) to receive back part (or all) of
+        the deposited funds.
+
+        Parameters
+        ----------
+        _amount: int
+            Amount LP tokens to burn.
+        min_amounts: List[int], optional
+            Minimum required amounts for each coin.  Default is 0 each.
+
+        Note
+        ----
+        "This withdrawal method is very safe, does no complex math"
+        """
+        min_amounts = min_amounts or [0] * self.n
+
+        total_supply = self.tokens
+        self.tokens -= _amount
+        balances = self.balances
+        amount = _amount - 1  # Make rounding errors favoring other LPs a tiny bit
+
+        for i in range(self.n):
+            d_balance = balances[i] * amount // total_supply
+            assert d_balance >= min_amounts[i]
+            self.balances[i] = balances[i] - d_balance
+
+    # pylint: disable-next=too-many-locals
+    def remove_liquidity_imbalance(self, amounts, max_burn_amount=None):
+        """
+        Withdraw an imbalanced amount of tokens from the pool. Accounts for fees.
+
+        Parameters
+        ----------
+        amounts: list of int
+            Amounts of tokens to withdraw (positive ints)
+        max_burn_amount: int, optional
+            Maximum amount of LP tokens to burn
+
+        Returns
+        -------
+        burn_amount: int
+            Amount of LP token burned in the withdrawal
+        fees: list of int
+            Amount of fees paid
+        """
+        A = self.A
+        old_balances = self.balances
+        D0 = self.get_D_mem(old_balances, A)
+
+        new_balances = self.balances[:]
+        for i in range(self.n):
+            amount = amounts[i]
+            assert amount >= 0  # must input positive ints
+            new_balances[i] -= amount
+        D1 = self.get_D_mem(new_balances, A)
+
+        fees = [0] * self.n
+        _fee = self.fee * self.n // (4 * (self.n - 1))
+        for i in range(self.n):
+            ideal_balance = D1 * old_balances[i] // D0
+            difference = abs(ideal_balance - new_balances[i])
+            fees[i] = _fee * difference // 10**10
+            admin_fee = fees[i] * self.admin_fee // 10**10
+            self.admin_balances[i] += admin_fee
+            self.balances[i] = new_balances[i] - admin_fee  # sans admin fees
+            new_balances[i] -= fees[i]
+
+        D2 = self.get_D_mem(new_balances, A)
+
+        burn_amount = self.tokens * (D0 - D2) // D0 + 1  # should be positive
+        if max_burn_amount:
+            assert burn_amount <= max_burn_amount
+        assert burn_amount >= 0
+
+        self.tokens -= burn_amount
+
+        return burn_amount, fees
+
     # pylint: disable-next=too-many-locals
     def calc_token_amount(self, amounts, use_fee=False):
         """

test/fixtures/curve/basepool.vy

@@ -507,6 +507,103 @@ def _calc_withdraw_one_coin(_token_amount: uint256, i: uint256) -> (uint256, uin
 def calc_withdraw_one_coin(_token_amount: uint256, i: uint256) -> uint256:
     return self._calc_withdraw_one_coin(_token_amount, i)[0]
 
+@external
+@nonreentrant('lock')
+def remove_liquidity(_burn_amount: uint256, _min_amounts: uint256[N_COINS]):
+    """
+    This withdrawal method is very safe, does no complex math
+    """
+    total_supply: uint256 = self.token.totalSupply()
+    amounts: uint256[N_COINS] = empty(uint256[N_COINS])
+    amount: uint256 =  _burn_amount - 1 # favor LPs a little bit
+
+    for i in range(N_COINS):
+        old_balance: uint256 = self.balances[i]
+        value: uint256 = old_balance * amount / total_supply
+        assert value >= _min_amounts[i], "Withdrawal resulted in fewer coins than expected"
+        self.balances[i] = old_balance - value
+        amounts[i] = value
+        # sim: comment-out the interaction with coin
+        # ----------------------------------------------------------
+        # response: Bytes[32] = raw_call(
+        #     self.coins[i],
+        #     concat(
+        #         method_id("transfer(address,uint256)"),
+        #         convert(_receiver, bytes32),
+        #         convert(value, bytes32),
+        #     ),
+        #     max_outsize=32,
+        # )
+        # if len(response) > 0:
+        #     assert convert(response, bool)
+
+    total_supply -= _burn_amount
+    self.token.burnFrom(msg.sender, _burn_amount)  # dev: insufficient funds
+
+    log RemoveLiquidity(msg.sender, amounts, empty(uint256[N_COINS]), total_supply)
+
+@external
+@nonreentrant('lock')
+def remove_liquidity_imbalance(
+    _amounts: uint256[N_COINS],
+    _max_burn_amount: uint256,
+) -> uint256:
+    """
+    Remove an imbalanced amount of coins from the pool.
+    """
+    amp: uint256 = self._A()
+    rates: uint256[N_COINS] = RATES
+    old_balances: uint256[N_COINS] = self.balances
+    D0: uint256 = self.get_D_mem(old_balances, amp)
+
+    new_balances: uint256[N_COINS] = old_balances
+    for i in range(N_COINS):
+        amount: uint256 = _amounts[i]
+        if amount != 0:
+            # print(new_balances[i])
+            # print(amount)
+            assert new_balances[i] > amount, 'Foo'
+            new_balances[i] -= amount
+            # sim: comment-out the interaction with coin
+            # ----------------------------------------------------------
+            # response: Bytes[32] = raw_call(
+            #     self.coins[i],
+            #     concat(
+            #         method_id("transfer(address,uint256)"),
+            #         convert(_receiver, bytes32),
+            #         convert(amount, bytes32),
+            #     ),
+            #     max_outsize=32,
+            # )
+            # if len(response) > 0:
+            #     assert convert(response, bool)
+    D1: uint256 = self.get_D_mem(new_balances, amp)
+
+    fees: uint256[N_COINS] = empty(uint256[N_COINS])
+    base_fee: uint256 = self.fee * N_COINS / (4 * (N_COINS - 1))
+    for i in range(N_COINS):
+        ideal_balance: uint256 = D1 * old_balances[i] / D0
+        difference: uint256 = 0
+        new_balance: uint256 = new_balances[i]
+        if ideal_balance > new_balance:
+            difference = ideal_balance - new_balance
+        else:
+            difference = new_balance - ideal_balance
+        fees[i] = base_fee * difference / FEE_DENOMINATOR
+        self.balances[i] = new_balance - (fees[i] * self.admin_fee / FEE_DENOMINATOR)
+        new_balances[i] -= fees[i]
+    D2: uint256 = self.get_D_mem(new_balances, amp)
+
+    total_supply: uint256 = self.token.totalSupply()
+    burn_amount: uint256 = ((D0 - D2) * total_supply / D0) + 1
+    assert burn_amount > 1  # dev: zero tokens burned
+    assert burn_amount <= _max_burn_amount
+
+    total_supply -= burn_amount
+    self.token.burnFrom(msg.sender, burn_amount)  # dev: insufficient funds
+    log RemoveLiquidityImbalance(msg.sender, _amounts, fees, D1, total_supply)
+
+    return burn_amount
 
 @external
 @nonreentrant('lock')

test/unit/test_pool.py

@@ -283,3 +283,58 @@ def test_remove_liquidity_one_coin(vyper_3pool, amount, i):
 
     assert coin_balance == expected_coin_balance
     assert lp_supply == expected_lp_supply
+
+
+@given(positive_balance)
+@settings(
+    suppress_health_check=[HealthCheck.function_scoped_fixture],
+    max_examples=5,
+    deadline=None,
+)
+def test_remove_liquidity(vyper_3pool, amount):
+    """Test `remove_liquidity` against vyper implementation."""
+    assume(amount < vyper_3pool.totalSupply())
+
+    pool = initialize_pool(vyper_3pool)
+
+    vyper_3pool.remove_liquidity(amount, [0] * pool.n)
+    expected_balances = [vyper_3pool.balances(i) for i in range(3)]
+    expected_lp_supply = vyper_3pool.totalSupply()
+    expected_D = vyper_3pool.D()
+
+    pool.remove_liquidity(amount)
+
+    assert pool.balances == expected_balances
+    assert pool.tokens == expected_lp_supply
+    assert pool.D() == expected_D
+
+
+@given(positive_balance, positive_balance, positive_balance)
+@settings(
+    suppress_health_check=[HealthCheck.function_scoped_fixture],
+    max_examples=5,
+    deadline=None,
+)
+def test_remove_liquidity_imbalance(vyper_3pool, x0, x1, x2):
+    """Test `remove_liquidity_imbalance` against vyper implementation."""
+
+    _balances = [x0, x1, x2]
+    rates = [vyper_3pool.rates(i) for i in range(len(_balances))]
+    amounts = [b * 10**18 // r for b, r in zip(_balances, rates)]
+
+    assume(amounts[0] < vyper_3pool.balances(0))
+    assume(amounts[1] < vyper_3pool.balances(1))
+    assume(amounts[2] < vyper_3pool.balances(2))
+
+    pool = initialize_pool(vyper_3pool)
+
+    vyper_3pool.remove_liquidity_imbalance(amounts, 2**256 - 1)
+    expected_balances = [vyper_3pool.balances(i) for i in range(3)]
+    expected_lp_supply = vyper_3pool.totalSupply()
+    expected_D = vyper_3pool.D()
+
+    pool.remove_liquidity_imbalance(amounts)
+
+    assert pool.balances == expected_balances
+    assert pool.tokens == expected_lp_supply
+    assert pool.D() == expected_D