Reallocate by withdrawals + sink

src/PublicAllocator.sol

@@ -18,7 +18,13 @@ import {UtilsLib} from "./libraries/UtilsLib.sol";
 
 import {ErrorsLib} from "./libraries/ErrorsLib.sol";
 import {EventsLib} from "./libraries/EventsLib.sol";
-import {FlowCap, FlowConfig, SupplyConfig, IPublicAllocatorStaticTyping} from "./interfaces/IPublicAllocator.sol";
+import {
+    FlowCap,
+    FlowConfig,
+    SupplyConfig,
+    Withdrawal,
+    IPublicAllocatorStaticTyping
+} from "./interfaces/IPublicAllocator.sol";
 
 contract PublicAllocator is IPublicAllocatorStaticTyping {
     using MorphoLib for IMorpho;
@@ -59,36 +65,43 @@ contract PublicAllocator is IPublicAllocatorStaticTyping {
 
     /// PUBLIC ///
 
-    function reallocate(MarketAllocation[] calldata allocations) external payable {
+    function withdrawTo(Withdrawal[] calldata withdrawals, MarketParams calldata depositMarketParams)
+        external
+        payable
+    {
         if (msg.value != fee) revert ErrorsLib.IncorrectFee(msg.value);
 
-        uint256[] memory assets = new uint256[](allocations.length);
-        for (uint256 i = 0; i < allocations.length; ++i) {
-            // Do not compute interest twice for every market
-            MORPHO.accrueInterest(allocations[i].marketParams);
-            assets[i] = MORPHO.expectedSupplyAssets(allocations[i].marketParams, address(VAULT));
-        }
+        MarketAllocation[] memory allocations = new MarketAllocation[](withdrawals.length + 1);
+        allocations[withdrawals.length].marketParams = depositMarketParams;
+        allocations[withdrawals.length].assets = type(uint256).max;
 
-        VAULT.reallocate(allocations);
+        uint128 totalWithdrawn;
 
-        MarketParams memory marketParams;
-        for (uint256 i = 0; i < allocations.length; ++i) {
-            marketParams = allocations[i].marketParams;
-            Id id = marketParams.id();
-            uint256 newAssets = MORPHO.expectedSupplyAssets(marketParams, address(VAULT));
-            if (newAssets >= assets[i]) {
-                if (newAssets > supplyCap[id]) revert ErrorsLib.PublicAllocatorSupplyCapExceeded(id);
-                uint128 inflow = (newAssets - assets[i]).toUint128();
-                flowCap[id].maxIn -= inflow;
-                flowCap[id].maxOut = (flowCap[id].maxOut).saturatingAdd(inflow);
-            } else {
-                uint128 outflow = (assets[i] - newAssets).toUint128();
-                flowCap[id].maxIn = (flowCap[id].maxIn).saturatingAdd(outflow);
-                flowCap[id].maxOut -= outflow;
+        for (uint256 i = 0; i < withdrawals.length; ++i) {
+            allocations[i].marketParams = withdrawals[i].marketParams;
+            Id id = withdrawals[i].marketParams.id();
+            uint256 assets = MORPHO.expectedSupplyAssets(withdrawals[i].marketParams, address(VAULT));
+            uint128 withdrawnAssets = withdrawals[i].amount;
+            // Clamp at 0 if withdrawnAssets is too big
+            if (withdrawnAssets > assets) {
+                withdrawnAssets = assets.toUint128();
             }
+
+            totalWithdrawn += withdrawnAssets;
+            allocations[i].assets = assets - withdrawnAssets;
+            flowCap[id].maxIn = (flowCap[id].maxIn).saturatingAdd(withdrawnAssets);
+            flowCap[id].maxOut -= withdrawnAssets;
+            emit EventsLib.PublicWithdrawal(id, withdrawnAssets);
         }
 
-        emit EventsLib.PublicReallocate(msg.sender, msg.value);
+        VAULT.reallocate(allocations);
+
+        Id depositMarketId = depositMarketParams.id();
+        uint256 depositAssets = MORPHO.expectedSupplyAssets(depositMarketParams, address(VAULT));
+        if (depositAssets > supplyCap[depositMarketId]) revert ErrorsLib.PublicAllocatorSupplyCapExceeded(depositMarketId);
+        flowCap[depositMarketId].maxIn -= totalWithdrawn;
+        flowCap[depositMarketId].maxOut = (flowCap[depositMarketId].maxOut).saturatingAdd(totalWithdrawn);
+        emit EventsLib.PublicReallocateTo(msg.sender, fee, depositMarketId, totalWithdrawn);
     }
 
     /// OWNER ONLY ///

src/interfaces/IPublicAllocator.sol

@@ -1,7 +1,7 @@
 // SPDX-License-Identifier: GPL-2.0-or-later
 pragma solidity ^0.8.21;
 
-import {IMetaMorpho, IMorpho, MarketAllocation, Id} from "../../lib/metamorpho/src/interfaces/IMetaMorpho.sol";
+import {IMetaMorpho, IMorpho, MarketAllocation, Id, MarketParams} from "../../lib/metamorpho/src/interfaces/IMetaMorpho.sol";
 
 struct FlowCap {
     uint128 maxIn;
@@ -18,6 +18,11 @@ struct SupplyConfig {
     uint256 cap;
 }
 
+struct Withdrawal {
+    MarketParams marketParams;
+    uint128 amount;
+}
+
 /// @dev This interface is used for factorizing IPublicAllocatorStaticTyping and IPublicAllocator.
 /// @dev Consider using the IPublicAllocator interface instead of this one.
 interface IPublicAllocatorBase {
@@ -28,7 +33,9 @@ interface IPublicAllocatorBase {
     function fee() external view returns (uint256);
     function supplyCap(Id) external view returns (uint256);
 
-    function reallocate(MarketAllocation[] calldata allocations) external payable;
+    function withdrawTo(Withdrawal[] calldata withdrawals, MarketParams calldata depositMarketParams)
+        external
+        payable;
     function setFee(uint256 _fee) external;
     function transferFee(address payable feeRecipient) external;
     function setFlowCaps(FlowConfig[] calldata _flowCaps) external;

src/libraries/EventsLib.sol

@@ -1,16 +1,19 @@
 // SPDX-License-Identifier: GPL-2.0-or-later
 pragma solidity ^0.8.0;
 
-import {FlowConfig, SupplyConfig} from "../interfaces/IPublicAllocator.sol";
+import {FlowConfig, SupplyConfig, Id} from "../interfaces/IPublicAllocator.sol";
 
 /// @title EventsLib
 /// @author Morpho Labs
 /// @custom:contact [email protected]
 /// @notice Library exposing events.
 library EventsLib {
-    /// @notice Emitted when the public reallocation is triggered.
-    event PublicReallocate(address sender, uint256 fee);
-    /// @notice Emitted when the owner changes the `fee`.
+    /// @notice Emitted during a public reallocation for each withdrawn-from market
+    event PublicWithdrawal(Id id, uint256 withdrawnAssets);
+    /// @notice Emitted at the end of a public reallocation
+    event PublicReallocateTo(address sender, uint256 fee, Id depositMarketId, uint256 depositedAssets);
+
+    /// @notice Emitted when the owner changes the `fee`
     event SetFee(uint256 fee);
     /// @notice Emitted when the owner transfers the fee.
     event TransferFee(uint256 amount);