Refactor allocation

src/MetaMorpho.sol

@@ -283,7 +283,33 @@ contract MetaMorpho is ERC4626, ERC20Permit, Ownable2Step, Multicall, IMetaMorph
         external
         onlyAllocator
     {
-        _reallocate(withdrawn, supplied);
+        uint256 balanceBefore = ERC20(asset()).balanceOf(address(this));
+
+        uint256 nbWithdrawn = withdrawn.length;
+
+        for (uint256 i; i < nbWithdrawn; ++i) {
+            MarketAllocation memory allocation = withdrawn[i];
+
+            MORPHO.withdraw(allocation.marketParams, allocation.assets, allocation.shares, address(this), address(this));
+        }
+
+        uint256 nbSupplied = supplied.length;
+
+        for (uint256 i; i < nbSupplied; ++i) {
+            MarketAllocation memory allocation = supplied[i];
+
+            MORPHO.supply(allocation.marketParams, allocation.assets, allocation.shares, address(this), hex"");
+
+            require(
+                _supplyBalance(allocation.marketParams) <= config[allocation.marketParams.id()].cap,
+                ErrorsLib.SUPPLY_CAP_EXCEEDED
+            );
+        }
+
+        uint256 balanceAfter = ERC20(asset()).balanceOf(address(this));
+
+        if (balanceAfter > balanceBefore) idle += balanceAfter - balanceBefore;
+        else idle -= balanceBefore - balanceAfter;
     }
 
     /* EXTERNAL */
@@ -522,29 +548,6 @@ contract MetaMorpho is ERC4626, ERC20Permit, Ownable2Step, Multicall, IMetaMorph
 
     /* LIQUIDITY ALLOCATION */
 
-    function _reallocate(MarketAllocation[] memory withdrawn, MarketAllocation[] memory supplied) internal {
-        uint256 nbWithdrawn = withdrawn.length;
-
-        for (uint256 i; i < nbWithdrawn; ++i) {
-            MarketAllocation memory allocation = withdrawn[i];
-
-            MORPHO.withdraw(allocation.marketParams, allocation.assets, 0, address(this), address(this));
-        }
-
-        uint256 nbSupplied = supplied.length;
-
-        for (uint256 i; i < nbSupplied; ++i) {
-            MarketAllocation memory allocation = supplied[i];
-
-            require(
-                _suppliable(allocation.marketParams, allocation.marketParams.id()) >= allocation.assets,
-                ErrorsLib.SUPPLY_CAP_EXCEEDED
-            );
-
-            MORPHO.supply(allocation.marketParams, allocation.assets, 0, address(this), hex"");
-        }
-    }
-
     function _supplyMorpho(uint256 assets) internal {
         uint256 nbMarkets = supplyQueue.length;
 

src/interfaces/IMetaMorpho.sol

@@ -22,6 +22,7 @@ struct PendingAddress {
 struct MarketAllocation {
     MarketParams marketParams;
     uint256 assets;
+    uint256 shares;
 }
 
 interface IMetaMorpho is IERC4626 {

test/forge/ReallocateIdleTest.sol

@@ -0,0 +1,56 @@
+// SPDX-License-Identifier: GPL-2.0-or-later
+pragma solidity ^0.8.0;
+
+import {SharesMathLib} from "@morpho-blue/libraries/SharesMathLib.sol";
+
+import "./helpers/BaseTest.sol";
+
+uint256 constant CAP2 = 100e18;
+uint256 constant INITIAL_DEPOSIT = 4 * CAP2;
+
+contract ReallocateTest is BaseTest {
+    using MarketParamsLib for MarketParams;
+    using MorphoLib for IMorpho;
+
+    MarketAllocation[] internal withdrawn;
+    MarketAllocation[] internal supplied;
+
+    function setUp() public override {
+        super.setUp();
+
+        _setCap(allMarkets[0], CAP2);
+        _setCap(allMarkets[1], CAP2);
+        _setCap(allMarkets[2], CAP2);
+
+        vm.prank(ALLOCATOR);
+        vault.setSupplyQueue(new Id[](0));
+
+        loanToken.setBalance(SUPPLIER, INITIAL_DEPOSIT);
+
+        vm.prank(SUPPLIER);
+        vault.deposit(INITIAL_DEPOSIT, ONBEHALF);
+    }
+
+    function testReallocateSupplyIdle(uint256[3] memory suppliedShares) public {
+        suppliedShares[0] = bound(suppliedShares[0], SharesMathLib.VIRTUAL_SHARES, CAP2 * SharesMathLib.VIRTUAL_SHARES);
+        suppliedShares[1] = bound(suppliedShares[1], SharesMathLib.VIRTUAL_SHARES, CAP2 * SharesMathLib.VIRTUAL_SHARES);
+        suppliedShares[2] = bound(suppliedShares[2], SharesMathLib.VIRTUAL_SHARES, CAP2 * SharesMathLib.VIRTUAL_SHARES);
+
+        supplied.push(MarketAllocation(allMarkets[0], 0, suppliedShares[0]));
+        supplied.push(MarketAllocation(allMarkets[1], 0, suppliedShares[1]));
+        supplied.push(MarketAllocation(allMarkets[2], 0, suppliedShares[2]));
+
+        uint256 idleBefore = vault.idle();
+
+        vm.prank(ALLOCATOR);
+        vault.reallocate(withdrawn, supplied);
+
+        assertEq(morpho.supplyShares(allMarkets[0].id(), address(vault)), suppliedShares[0], "morpho.supplyShares(0)");
+        assertEq(morpho.supplyShares(allMarkets[1].id(), address(vault)), suppliedShares[1], "morpho.supplyShares(1)");
+        assertEq(morpho.supplyShares(allMarkets[2].id(), address(vault)), suppliedShares[2], "morpho.supplyShares(2)");
+
+        uint256 expectedIdle = idleBefore - suppliedShares[0] / SharesMathLib.VIRTUAL_SHARES
+            - suppliedShares[1] / SharesMathLib.VIRTUAL_SHARES - suppliedShares[2] / SharesMathLib.VIRTUAL_SHARES;
+        assertApproxEqAbs(vault.idle(), expectedIdle, 3, "vault.idle() 1");
+    }
+}

test/forge/ReallocateWithdrawTest.sol

@@ -0,0 +1,126 @@
+// SPDX-License-Identifier: GPL-2.0-or-later
+pragma solidity ^0.8.0;
+
+import {UtilsLib} from "@morpho-blue/libraries/UtilsLib.sol";
+import {SharesMathLib} from "@morpho-blue/libraries/SharesMathLib.sol";
+
+import "./helpers/BaseTest.sol";
+
+uint256 constant CAP2 = 100e18;
+uint256 constant INITIAL_DEPOSIT = 4 * CAP2;
+
+contract ReallocateWithdrawTest is BaseTest {
+    using MarketParamsLib for MarketParams;
+    using MorphoBalancesLib for IMorpho;
+    using MorphoLib for IMorpho;
+    using SharesMathLib for uint256;
+    using UtilsLib for uint256;
+
+    MarketAllocation[] internal withdrawn;
+    MarketAllocation[] internal supplied;
+
+    function setUp() public override {
+        super.setUp();
+
+        _setCap(allMarkets[0], CAP2);
+        _setCap(allMarkets[1], CAP2);
+        _setCap(allMarkets[2], CAP2);
+
+        loanToken.setBalance(SUPPLIER, INITIAL_DEPOSIT);
+
+        vm.prank(SUPPLIER);
+        vault.deposit(INITIAL_DEPOSIT, ONBEHALF);
+    }
+
+    function testReallocateWithdrawAll() public {
+        withdrawn.push(MarketAllocation(allMarkets[0], 0, morpho.supplyShares(allMarkets[0].id(), address(vault))));
+        withdrawn.push(MarketAllocation(allMarkets[1], 0, morpho.supplyShares(allMarkets[1].id(), address(vault))));
+        withdrawn.push(MarketAllocation(allMarkets[2], 0, morpho.supplyShares(allMarkets[2].id(), address(vault))));
+
+        vm.prank(ALLOCATOR);
+        vault.reallocate(withdrawn, supplied);
+
+        assertEq(morpho.supplyShares(allMarkets[0].id(), address(vault)), 0, "morpho.supplyShares(0)");
+        assertEq(morpho.supplyShares(allMarkets[1].id(), address(vault)), 0, "morpho.supplyShares(1)");
+        assertEq(morpho.supplyShares(allMarkets[2].id(), address(vault)), 0, "morpho.supplyShares(2)");
+        assertEq(vault.idle(), INITIAL_DEPOSIT, "vault.idle() 1");
+    }
+
+    function testReallocateWithdrawSupply(uint256[3] memory withdrawnShares, uint256[3] memory suppliedAssets) public {
+        uint256[3] memory sharesBefore = [
+            morpho.supplyShares(allMarkets[0].id(), address(vault)),
+            morpho.supplyShares(allMarkets[1].id(), address(vault)),
+            morpho.supplyShares(allMarkets[2].id(), address(vault))
+        ];
+
+        withdrawnShares[0] = bound(withdrawnShares[0], 0, sharesBefore[0]);
+        withdrawnShares[1] = bound(withdrawnShares[1], 0, sharesBefore[1]);
+        withdrawnShares[2] = bound(withdrawnShares[2], 0, sharesBefore[2]);
+
+        uint256[3] memory totalSupplyAssets;
+        uint256[3] memory totalSupplyShares;
+        (totalSupplyAssets[0], totalSupplyShares[0],,) = morpho.expectedMarketBalances(allMarkets[0]);
+        (totalSupplyAssets[1], totalSupplyShares[1],,) = morpho.expectedMarketBalances(allMarkets[1]);
+        (totalSupplyAssets[2], totalSupplyShares[2],,) = morpho.expectedMarketBalances(allMarkets[2]);
+
+        uint256[3] memory withdrawnAssets = [
+            withdrawnShares[0].toAssetsDown(totalSupplyAssets[0], totalSupplyShares[0]),
+            withdrawnShares[1].toAssetsDown(totalSupplyAssets[1], totalSupplyShares[1]),
+            withdrawnShares[2].toAssetsDown(totalSupplyAssets[2], totalSupplyShares[2])
+        ];
+
+        if (withdrawnShares[0] > 0) withdrawn.push(MarketAllocation(allMarkets[0], 0, withdrawnShares[0]));
+        if (withdrawnAssets[1] > 0) withdrawn.push(MarketAllocation(allMarkets[1], withdrawnAssets[1], 0));
+        if (withdrawnShares[2] > 0) withdrawn.push(MarketAllocation(allMarkets[2], 0, withdrawnShares[2]));
+
+        totalSupplyAssets[0] -= withdrawnAssets[0];
+        totalSupplyAssets[1] -= withdrawnAssets[1];
+        totalSupplyAssets[2] -= withdrawnAssets[2];
+
+        totalSupplyShares[0] -= withdrawnShares[0];
+        totalSupplyShares[1] -= withdrawnShares[1];
+        totalSupplyShares[2] -= withdrawnShares[2];
+
+        uint256 expectedIdle = vault.idle() + withdrawnAssets[0] + withdrawnAssets[1] + withdrawnAssets[2];
+
+        suppliedAssets[0] = bound(suppliedAssets[0], 0, withdrawnAssets[0].zeroFloorSub(CAP2).min(expectedIdle));
+        expectedIdle -= suppliedAssets[0];
+
+        suppliedAssets[1] = bound(suppliedAssets[1], 0, withdrawnAssets[1].zeroFloorSub(CAP2).min(expectedIdle));
+        expectedIdle -= suppliedAssets[1];
+
+        suppliedAssets[2] = bound(suppliedAssets[2], 0, withdrawnAssets[2].zeroFloorSub(CAP2).min(expectedIdle));
+        expectedIdle -= suppliedAssets[2];
+
+        uint256[3] memory suppliedShares = [
+            suppliedAssets[0].toSharesDown(totalSupplyAssets[0], totalSupplyShares[0]),
+            suppliedAssets[1].toSharesDown(totalSupplyAssets[1], totalSupplyShares[1]),
+            suppliedAssets[2].toSharesDown(totalSupplyAssets[2], totalSupplyShares[2])
+        ];
+
+        if (suppliedShares[0] > 0) supplied.push(MarketAllocation(allMarkets[0], suppliedAssets[0], 0));
+        if (suppliedAssets[1] > 0) supplied.push(MarketAllocation(allMarkets[1], 0, suppliedShares[1]));
+        if (suppliedShares[2] > 0) supplied.push(MarketAllocation(allMarkets[2], suppliedAssets[2], 0));
+
+        vm.prank(ALLOCATOR);
+        vault.reallocate(withdrawn, supplied);
+
+        assertEq(
+            morpho.supplyShares(allMarkets[0].id(), address(vault)),
+            sharesBefore[0] - withdrawnShares[0] + suppliedShares[0],
+            "morpho.supplyShares(0)"
+        );
+        assertApproxEqAbs(
+            morpho.supplyShares(allMarkets[1].id(), address(vault)),
+            sharesBefore[1] - withdrawnShares[1] + suppliedShares[1],
+            SharesMathLib.VIRTUAL_SHARES,
+            "morpho.supplyShares(1)"
+        );
+        assertEq(
+            morpho.supplyShares(allMarkets[2].id(), address(vault)),
+            sharesBefore[2] - withdrawnShares[2] + suppliedShares[2],
+            "morpho.supplyShares(2)"
+        );
+        assertApproxEqAbs(vault.idle(), expectedIdle, 1, "vault.idle() 1");
+    }
+}

test/hardhat/MetaMorpho.spec.ts

@@ -13,6 +13,8 @@ import MorphoArtifact from "../../lib/morpho-blue/out/Morpho.sol/Morpho.json";
 // Without the division it overflows.
 const initBalance = MaxUint256 / 10000000000000000n;
 const oraclePriceScale = 1000000000000000000000000000000000000n;
+const virtualShares = 100000n;
+const virtualAssets = 1n;
 const nbMarkets = 5;
 
 let seed = 42;
@@ -179,21 +181,39 @@ describe("MetaMorpho", () => {
           const market = await morpho.market(id);
           const position = await morpho.position(id, await metaMorpho.getAddress());
 
-          const assets = position.supplyShares
-            .mulDivDown(market.totalSupplyAssets + 1n, market.totalSupplyShares + 10n ** 6n)
-            .min(market.totalSupplyAssets - market.totalBorrowAssets);
+          const liquidity = market.totalSupplyAssets - market.totalBorrowAssets;
+          const liquidShares = liquidity.mulDivDown(
+            market.totalSupplyShares + virtualShares,
+            market.totalSupplyAssets + virtualAssets,
+          );
 
           return {
             marketParams,
-            assets,
+            market,
+            shares: position.supplyShares.min(liquidShares),
           };
         }),
       );
 
       await metaMorpho.connect(allocator).reallocate(
-        allocation.filter(({ assets }) => assets > 0n),
         allocation
-          .map(({ marketParams, assets }) => ({ marketParams, assets: (assets * 3n) / 4n }))
+          .map(({ marketParams, shares }) => ({
+            marketParams,
+            assets: 0n,
+            // Always withdraw all, up to the liquidity.
+            shares,
+          }))
+          .filter(({ shares }) => shares > 0n),
+        allocation
+          .map(({ marketParams, market, shares }) => {
+            const assets = shares.mulDivDown(
+              market.totalSupplyAssets + virtualAssets,
+              market.totalSupplyShares + virtualShares,
+            );
+
+            // Always supply 3/4 of what the vault withdrawn.
+            return { marketParams, assets: (assets * 3n) / 4n, shares: 0n };
+          })
           .filter(({ assets }) => assets > 0n),
       );