feat: subcaps

src/PublicAllocator.sol

@@ -24,6 +24,7 @@ contract PublicAllocator is Ownable2Step, Multicall, IPublicAllocatorStaticTypin
     IMorpho public immutable MORPHO;
     mapping(Id => int256) public flows;
     mapping(Id => FlowCaps) public flowCaps;
+    mapping(Id => uint256) public supplyCaps;
     // using IMorpho
 
     /// CONSTRUCTOR ///
@@ -54,6 +55,10 @@ contract PublicAllocator is Ownable2Step, Multicall, IPublicAllocatorStaticTypin
             market = MORPHO.market(id);
             uint256 newShares = MORPHO.supplyShares(id, address(VAULT));
             if (newShares >= shares[i]) {
+                // Withdrawing small enough amounts when the cap is already exceeded can result in the error below
+                if (newShares.toAssetsUp(market.totalSupplyAssets,market.totalSupplyShares) > supplyCaps[id]) {
+                    revert ErrorsLib.PublicAllocatorSupplyCapExceeded(id);
+                }
                 flows[id] +=
                     int256((newShares - shares[i]).toAssetsUp(market.totalSupplyAssets, market.totalSupplyShares));
                 if (flows[id] > int256(uint256(flowCaps[id].inflow))) {
@@ -93,4 +98,9 @@ contract PublicAllocator is Ownable2Step, Multicall, IPublicAllocatorStaticTypin
             flows[flowConfig.id] = 0;
         }
     }
+
+    // Set supply cap. Public reallocation will not be able to increase supply if it ends above its cap. 
+    function setCap(Id id, uint supplyCap) external onlyOwner {
+        supplyCaps[id] = supplyCap;
+    }
 }

src/interfaces/IPublicAllocator.sol

@@ -27,11 +27,13 @@ interface IPublicAllocatorBase {
     function VAULT() external view returns (IMetaMorpho);
     function MORPHO() external view returns (IMorpho);
     function flows(Id) external view returns (int256);
+    function supplyCaps(Id) external view returns (uint256);
 
     function reallocate(MarketAllocation[] calldata allocations) external payable;
     function setFee(uint256 _fee) external;
     function transferFee(address feeRecipient) external;
     function setFlow(FlowConfig calldata flowConfig) external;
+    function setCap(Id id, uint supplyCap) external;
 }
 
 /// @dev This interface is inherited by PublicAllocator so that function signatures are checked by the compiler.

src/libraries/ErrorsLib.sol

@@ -25,4 +25,7 @@ library ErrorsLib {
 
     /// @notice Thrown when the fee recipient fails to receive the fee
     error FeeTransferFail();
+
+    /// @notice Thrown when the supply cap has been exceeded on market `id` during a reallocation of funds.
+    error PublicAllocatorSupplyCapExceeded(Id id);
 }

test/PublicAllocator.t.sol

@@ -33,6 +33,12 @@ contract PublicAllocatorTest is IntegrationTest {
 
         _setCap(allMarkets[0], CAP2);
         _sortSupplyQueueIdleLast();
+
+        // Remove public allocator caps by default
+        vm.prank(OWNER);
+        publicAllocator.setCap(idleParams.id(), type(uint).max);
+        vm.prank(OWNER);
+        publicAllocator.setCap(allMarkets[0].id(), type(uint).max);
     }
 
     function testOwner() public {
@@ -79,6 +85,14 @@ contract PublicAllocatorTest is IntegrationTest {
         publicAllocator.setFee(fee);
     }
 
+    function testSetCapAccess(address sender, Id id, uint cap) public {
+        vm.assume(sender != OWNER);
+        vm.prank(sender);
+        vm.expectRevert(abi.encodeWithSelector(Ownable.OwnableUnauthorizedAccount.selector, sender));
+        publicAllocator.setCap(id, cap);
+
+    }
+
     function testReallocateNetting(uint128 flow) public {
         flow = uint128(bound(flow, 0, CAP2));
         vm.assume(flow != 0);
@@ -183,4 +197,82 @@ contract PublicAllocatorTest is IntegrationTest {
     }
 
     receive() external payable {}
+
+    function testInflowGoesAboveCap(uint cap, uint128 flow) public {
+        cap = bound(cap, 0, CAP2-1);
+        flow = uint128(bound(flow, cap+1, CAP2));
+
+        vm.startPrank(OWNER);
+        publicAllocator.setCap(allMarkets[0].id(), cap);
+        publicAllocator.setFlow(FlowConfig(idleParams.id(), FlowCaps(type(uint128).max, 0), false));
+        publicAllocator.setFlow(FlowConfig(allMarkets[0].id(), FlowCaps(0, type(uint128).max), false));
+        vm.stopPrank();
+
+        // Should work at cap
+        allocations.push(MarketAllocation(idleParams, INITIAL_DEPOSIT - cap));
+        allocations.push(MarketAllocation(allMarkets[0], cap));
+        publicAllocator.reallocate(allocations);
+
+        // Should not work above cap
+        allocations.push(MarketAllocation(idleParams, INITIAL_DEPOSIT - flow));
+        allocations.push(MarketAllocation(allMarkets[0], flow));
+
+        vm.expectRevert(abi.encodeWithSelector(PAErrorsLib.PublicAllocatorSupplyCapExceeded.selector, allMarkets[0].id())); 
+        publicAllocator.reallocate(allocations);
+    }
+
+    function testInflowStartsAboveCap(uint cap, uint128 flow) public {
+        cap = bound(cap, 0, CAP2-2);
+        flow = uint128(bound(flow, cap, CAP2-1));
+
+        // Remove flow limits
+        vm.prank(OWNER);
+        publicAllocator.setFlow(FlowConfig(idleParams.id(), FlowCaps(type(uint128).max, 0), false));
+        vm.prank(OWNER);
+        publicAllocator.setFlow(FlowConfig(allMarkets[0].id(), FlowCaps(0, type(uint128).max), false));
+
+        // Set supply above future public allocator cap
+        allocations.push(MarketAllocation(idleParams, INITIAL_DEPOSIT - flow));
+        allocations.push(MarketAllocation(allMarkets[0], flow));
+        publicAllocator.reallocate(allocations);
+
+        // Set supply in market 0 > public allocation cap
+        vm.prank(OWNER);
+        publicAllocator.setCap(allMarkets[0].id(), cap);
+
+        // Increase supply even more (by 1)
+        allocations[0].assets = INITIAL_DEPOSIT - flow - 1;
+        allocations[1].assets = flow + 1;
+
+        vm.expectRevert(abi.encodeWithSelector(PAErrorsLib.PublicAllocatorSupplyCapExceeded.selector, allMarkets[0].id()));
+        publicAllocator.reallocate(allocations);
+    }
+
+    function testStrictOutflowStartsAboveCap(uint cap, uint128 flow, uint128 flow2) public {
+        cap = bound(cap, 0, CAP2-2);
+        flow = uint128(bound(flow, cap+2, CAP2));
+        flow2 = uint128(bound(flow2,cap+1,flow-1));
+
+        // Remove flow limits
+        vm.prank(OWNER);
+        publicAllocator.setFlow(FlowConfig(idleParams.id(), FlowCaps(type(uint128).max, 0), false));
+        vm.prank(OWNER);
+        publicAllocator.setFlow(FlowConfig(allMarkets[0].id(), FlowCaps(0, type(uint128).max), false));
+
+        // Set supply above future public allocator cap
+        allocations.push(MarketAllocation(idleParams, INITIAL_DEPOSIT - flow));
+        allocations.push(MarketAllocation(allMarkets[0], flow));
+        publicAllocator.reallocate(allocations);
+
+        // Set supply in market 0 > public allocation cap
+        vm.prank(OWNER);
+        publicAllocator.setCap(allMarkets[0].id(), cap);
+
+        // Strictly decrease supply
+        delete allocations;
+        allocations.push(MarketAllocation(allMarkets[0], flow2));
+        allocations.push(MarketAllocation(idleParams, INITIAL_DEPOSIT - flow2));
+
+        publicAllocator.reallocate(allocations);
+    }
 }