Enable pool fees when rebalancing Aerodrome AMO

contracts/contracts/interfaces/aerodrome/IAMOStrategy.sol

@@ -59,4 +59,9 @@ interface IAMOStrategy {
     function claimGovernance() external;
 
     function transferGovernance(address _governor) external;
+
+    function getPositionPrincipal()
+        external
+        view
+        returns (uint256 _amountWeth, uint256 _amountOethb);
 }

contracts/contracts/interfaces/aerodrome/ISugarHelper.sol

@@ -29,7 +29,7 @@ interface ISugarHelper {
         uint160 sqrtRatioX96,
         uint160 sqrtRatioAX96,
         uint160 sqrtRatioBX96
-    ) external pure returns (uint256 liquidity);
+    ) external pure returns (uint128 liquidity);
 
     /// @notice Computes the amount of token0 for a given amount of token1 and price range
     /// @param amount1 Amount of token1 to estimate liquidity

contracts/contracts/strategies/aerodrome/AerodromeAMOStrategy.sol

@@ -351,7 +351,7 @@
      */
     function depositAll() external override onlyVault nonReentrant {
         uint256 _wethBalance = IERC20(WETH).balanceOf(address(this));
-        if (_wethBalance > 0) {
+        if (_wethBalance > 1e12) {
             _deposit(WETH, _wethBalance);
         }
     }
@@ -426,11 +426,13 @@
 
         /**
          * When rebalance is called for the first time there is no strategy
-         * liquidity in the pool yet. The full liquidity removal is thus skipped.
+         * liquidity in the pool yet. The liquidity removal is thus skipped.
+         * Also execute this function when WETH is required for the swap.
          */
-        if (tokenId != 0) {
-            _removeLiquidity(1e18);
+        if (tokenId != 0 && _swapWeth && _amountToSwap > 0) {
+            _ensureWETHBalance(_amountToSwap);
         }
+
         // in some cases we will just want to add liquidity and not issue a swap to move the
         // active trading position within the pool
         if (_amountToSwap > 0) {
@@ -531,6 +533,8 @@
             _amountOethbCollected,
             underlyingAssets
         );
+
+        _burnOethbOnTheContract();
     }
 
     /**
@@ -545,6 +549,8 @@
         uint256 _balance = _tokenToSwap.balanceOf(address(this));
 
         if (_balance < _amountToSwap) {
+            // This should never trigger since _ensureWETHBalance will already
+            // throw an error if there is not enough WETH
             if (_swapWeth) {
                 revert NotEnoughWethForSwap(_balance, _amountToSwap);
             }
@@ -576,6 +582,14 @@
                     : sqrtRatioX96TickHigher
             })
         );
+
+        /**
+         * In the interest of each function in _rebalance to leave the contract state as
+         * clean as possible the OETHb tokens here are burned. This decreases the
+         * dependence where `_swapToDesiredPosition` function relies on later functions
+         * (`addLiquidity`) to burn the OETHb. Reducing the risk of error introduction.
+         */
+        _burnOethbOnTheContract();
     }
 
     /**
@@ -590,7 +604,10 @@
     function _addLiquidity() internal gaugeUnstakeAndRestake {
         uint256 _wethBalance = IERC20(WETH).balanceOf(address(this));
         uint256 _oethbBalance = IERC20(OETHb).balanceOf(address(this));
-        require(_wethBalance > 0, "Must add some WETH");
+        // don't deposit small liquidity amounts
+        if (_wethBalance <= 1e12) {
+            return;
+        }
 
         uint160 _currentPrice = getPoolX96Price();
         /**
@@ -709,6 +726,7 @@
      */
     function _checkForExpectedPoolPrice(bool throwException)
         internal
+        view
         returns (bool _isExpectedRange, uint256 _wethSharePct)
     {
         require(
@@ -759,7 +777,7 @@
      */
     function _burnOethbOnTheContract() internal {
         uint256 _oethbBalance = IERC20(OETHb).balanceOf(address(this));
-        if (_oethbBalance > 0) {
+        if (_oethbBalance > 1e12) {
             IVault(vaultAddress).burnForStrategy(_oethbBalance);
         }
     }
@@ -802,6 +820,36 @@
         emit UnderlyingAssetsUpdated(underlyingAssets);
     }
 
+    /**
+     * @dev This function removes the appropriate amount of liquidity to assure that the required
+     * amount of WETH is available on the contract
+     *
+     * @param _amount  WETH balance required on the contract
+     */
+    function _ensureWETHBalance(uint256 _amount) internal {
+        uint256 _wethBalance = IERC20(WETH).balanceOf(address(this));
+        if (_wethBalance >= _amount) {
+            return;
+        }
+
+        require(tokenId != 0, "No liquidity available");
+        uint256 _additionalWethRequired = _amount - _wethBalance;
+        (uint256 _wethInThePool, ) = getPositionPrincipal();
+
+        if (_wethInThePool < _additionalWethRequired) {
+            revert NotEnoughWethLiquidity(
+                _wethInThePool,
+                _additionalWethRequired
+            );
+        }
+
+        uint256 shareOfWethToRemove = Math.min(
+            _additionalWethRequired.divPrecisely(_wethInThePool) + 1,
+            1e18
+        );
+        _removeLiquidity(shareOfWethToRemove);
+    }
+
     /**
      * @notice Withdraw an `amount` of assets from the platform and
      *         send to the `_recipient`.
@@ -817,28 +865,8 @@
         require(_asset == WETH, "Unsupported asset");
         require(_recipient == vaultAddress, "Only withdraw to vault allowed");
 
-        uint256 _wethBalance = IERC20(WETH).balanceOf(address(this));
-        if (_wethBalance < _amount) {
-            require(tokenId != 0, "No liquidity available");
-            uint256 _additionalWethRequired = _amount - _wethBalance;
-            (uint256 _wethInThePool, ) = getPositionPrincipal();
-
-            if (_wethInThePool < _additionalWethRequired) {
-                revert NotEnoughWethLiquidity(
-                    _wethInThePool,
-                    _additionalWethRequired
-                );
-            }
-
-            uint256 shareOfWethToRemove = Math.min(
-                _additionalWethRequired.divPrecisely(_wethInThePool) + 1,
-                1e18
-            );
-            _removeLiquidity(shareOfWethToRemove);
-        }
+        _ensureWETHBalance(_amount);
 
-        // burn remaining OETHb
-        _burnOethbOnTheContract();
         _withdraw(_recipient, _amount);
     }
 
@@ -854,8 +882,6 @@
         if (_balance > 0) {
             _withdraw(vaultAddress, _balance);
         }
-        // burn remaining OETHb
-        _burnOethbOnTheContract();
     }
 
     function _withdraw(address _recipient, uint256 _amount) internal {

contracts/contracts/utils/AerodromeAMOQuoter.sol

@@ -39,12 +39,10 @@
     /// --- CONSTANT & IMMUTABLE
     ////////////////////////////////////////////////////////////////
     uint256 public constant BINARY_MIN_AMOUNT = 1 wei;
-    uint256 public constant BINARY_MAX_AMOUNT_FOR_REBALANCE = 3_000 ether;
-    uint256 public constant BINARY_MAX_AMOUNT_FOR_PUSH_PRICE = 5_000_000 ether;
 
-    uint256 public constant BINARY_MAX_ITERATIONS = 100;
+    uint256 public constant BINARY_MAX_ITERATIONS = 40;
     uint256 public constant PERCENTAGE_BASE = 1e18; // 100%
-    uint256 public constant ALLOWED_VARIANCE_PERCENTAGE = 1e12; // 0.0001%
+    uint256 public constant ALLOWED_VARIANCE_PERCENTAGE = 1e15; // 0.1%
 
     ////////////////////////////////////////////////////////////////
     /// --- VARIABLES STORAGE
@@ -101,9 +99,11 @@
 
             strategy.setAllowedPoolWethShareInterval(shareStart, shareEnd);
         }
+
         uint256 iterations = 0;
         uint256 low = BINARY_MIN_AMOUNT;
-        uint256 high = BINARY_MAX_AMOUNT_FOR_REBALANCE;
+        uint256 high;
+        (high, ) = strategy.getPositionPrincipal();
         int24 lowerTick = strategy.lowerTick();
         int24 upperTick = strategy.upperTick();
         bool swapWETHForOETHB = getSwapDirectionForRebalance();
@@ -338,6 +338,14 @@
         return currentPrice > targetPrice;
     }
 
+    // returns total amount in the position principal of the Aerodrome AMO strategy. Needed as a
+    // separate function because of the limitation in local variable count in getAmountToSwapToReachPrice
+    function getTotalStrategyPosition() internal returns (uint256) {
+        (uint256 wethAmount, uint256 oethBalance) = strategy
+            .getPositionPrincipal();
+        return wethAmount + oethBalance;
+    }
+
     /// @notice Get the amount of tokens to swap to reach the target price.
     /// @dev This act like a quoter, i.e. the transaction is not performed.
     /// @dev Because the amount to swap can be largely overestimated, because CLAMM alow partial orders,
@@ -359,34 +367,53 @@
     {
         uint256 iterations = 0;
         uint256 low = BINARY_MIN_AMOUNT;
-        uint256 high = BINARY_MAX_AMOUNT_FOR_PUSH_PRICE;
+        // high search start is twice the position principle of Aerodrome AMO strategy.
+        // should be more than enough
+        uint256 high = getTotalStrategyPosition() * 2;
         bool swapWETHForOETHB = getSwapDirection(sqrtPriceTargetX96);
 
         while (low <= high && iterations < BINARY_MAX_ITERATIONS) {
             uint256 mid = (low + high) / 2;
 
             // Call QuoterV2 from SugarHelper
-            (, uint160 sqrtPriceX96After, , ) = quoterV2.quoteExactInputSingle(
-                IQuoterV2.QuoteExactInputSingleParams({
-                    tokenIn: swapWETHForOETHB
-                        ? clPool.token0()
-                        : clPool.token1(),
-                    tokenOut: swapWETHForOETHB
-                        ? clPool.token1()
-                        : clPool.token0(),
-                    amountIn: mid,
-                    tickSpacing: strategy.tickSpacing(),
-                    sqrtPriceLimitX96: sqrtPriceTargetX96
-                })
-            );
+            (uint256 amountOut, uint160 sqrtPriceX96After, , ) = quoterV2
+                .quoteExactInputSingle(
+                    IQuoterV2.QuoteExactInputSingleParams({
+                        tokenIn: swapWETHForOETHB
+                            ? clPool.token0()
+                            : clPool.token1(),
+                        tokenOut: swapWETHForOETHB
+                            ? clPool.token1()
+                            : clPool.token0(),
+                        amountIn: mid,
+                        tickSpacing: strategy.tickSpacing(),
+                        sqrtPriceLimitX96: sqrtPriceTargetX96
+                    })
+                );
 
             if (
                 isWithinAllowedVariance(sqrtPriceX96After, sqrtPriceTargetX96)
             ) {
-                return (mid, iterations, swapWETHForOETHB, sqrtPriceX96After);
+                /** Very important to return `amountOut` instead of `mid` as the first return parameter.
+                 * The issues was that when quoting we impose a swap price limit (sqrtPriceLimitX96: sqrtPriceTargetX96)
+                 * and in that case the `amountIn` acts like a maximum amount to swap. And we don't know how much
+                 * of that amount was actually consumed. For that reason we "estimate" it by returning the
+                 * amountOut since that is only going to be a couple of basis point away from amountIn in the
+                 * worst cases.
+                 *
+                 * Note: we could be returning mid instead of amountOut in cases when those values are only basis
+                 * points apart (assuming that complete balance of amountIn has been consumed) but that might increase
+                 * complexity too much in an already complex contract.
+                 */
+                return (
+                    amountOut,
+                    iterations,
+                    swapWETHForOETHB,
+                    sqrtPriceX96After
+                );
             } else if (low == high) {
                 // target swap amount not found.
-                // try increasing BINARY_MAX_AMOUNT_FOR_PUSH_PRICE
+                // might be that "high" amount is too low on start
                 revert("SwapAmountNotFound");
             } else if (
                 swapWETHForOETHB
@@ -539,7 +566,6 @@
             revert("Previous call should only revert, it cannot succeed");
         } catch (bytes memory reason) {
             bytes4 receivedSelector = bytes4(reason);
-
             if (receivedSelector == QuoterHelper.ValidAmount.selector) {
                 uint256 value;
                 uint256 iterations;

contracts/deploy/base/018_upgrade_amo.js

@@ -0,0 +1,28 @@
+const { deployOnBaseWithGuardian } = require("../../utils/deploy-l2");
+const {
+  deployBaseAerodromeAMOStrategyImplementation,
+} = require("../deployActions");
+
+module.exports = deployOnBaseWithGuardian(
+  {
+    deployName: "018_upgrade_amo",
+  },
+  async ({ ethers }) => {
+    const cAMOStrategyProxy = await ethers.getContract(
+      "AerodromeAMOStrategyProxy"
+    );
+    const cAMOStrategyImpl =
+      await deployBaseAerodromeAMOStrategyImplementation();
+
+    return {
+      actions: [
+        {
+          // 1. Upgrade AMO
+          contract: cAMOStrategyProxy,
+          signature: "upgradeTo(address)",
+          args: [cAMOStrategyImpl.address],
+        },
+      ],
+    };
+  }
+);

contracts/test/_fixture-base.js

@@ -148,14 +148,15 @@ const defaultBaseFixture = deployments.createFixture(async () => {
   );
   const oethVaultSigner = await impersonateAccount(oethbVault.address);
 
-  let strategist;
+  let strategist, harvesterSigner;
   if (isFork) {
     // Impersonate strategist on Fork
     strategist = await impersonateAndFund(strategistAddr);
     strategist.address = strategistAddr;
 
     await impersonateAndFund(governor.address);
     await impersonateAndFund(timelock.address);
+    harvesterSigner = await impersonateAndFund(harvester.address);
 
     // configure Vault to not automatically deposit to strategy
     await oethbVault.connect(governor).setVaultBuffer(oethUnits("1"));
@@ -208,6 +209,7 @@ const defaultBaseFixture = deployments.createFixture(async () => {
     wOETHb,
     zapper,
     harvester,
+    harvesterSigner,
     dripper,
 
     // Bridged WOETH