Repay/Withdraw max

src/Blue.sol

@@ -122,15 +122,22 @@ contract Blue {
     function withdraw(Market calldata market, uint256 amount, address onBehalf) external {
         Id id = market.id();
         require(lastUpdate[id] != 0, Errors.MARKET_NOT_CREATED);
-        require(amount != 0, Errors.ZERO_AMOUNT);
         require(_isSenderOrIsApproved(onBehalf), Errors.MANAGER_NOT_APPROVED);
 
         _accrueInterests(market, id);
 
-        uint256 shares = amount.toSharesUp(totalSupply[id], totalSupplyShares[id]);
+        uint256 shares;
+        if (amount == type(uint256).max) {
+            amount = supplyShare[id][onBehalf].toAssetsDown(totalSupply[id], totalSupplyShares[id]);
+            shares = supplyShare[id][onBehalf];
+        } else {
+            shares = amount.toSharesUp(totalSupply[id], totalSupplyShares[id]);
+        }
+
+        require(amount != 0, Errors.ZERO_AMOUNT);
+
         supplyShare[id][onBehalf] -= shares;
         totalSupplyShares[id] -= shares;
-
         totalSupply[id] -= amount;
 
         require(totalBorrow[id] <= totalSupply[id], Errors.INSUFFICIENT_LIQUIDITY);
@@ -163,14 +170,21 @@ contract Blue {
     function repay(Market calldata market, uint256 amount, address onBehalf) external {
         Id id = market.id();
         require(lastUpdate[id] != 0, Errors.MARKET_NOT_CREATED);
-        require(amount != 0, Errors.ZERO_AMOUNT);
 
         _accrueInterests(market, id);
 
-        uint256 shares = amount.toSharesDown(totalBorrow[id], totalBorrowShares[id]);
+        uint256 shares;
+        if (amount == type(uint256).max) {
+            amount = borrowShare[id][onBehalf].toAssetsUp(totalBorrow[id], totalBorrowShares[id]);
+            shares = borrowShare[id][onBehalf];
+        } else {
+            shares = amount.toSharesDown(totalBorrow[id], totalBorrowShares[id]);
+        }
+
+        require(amount != 0, Errors.ZERO_AMOUNT);
+
         borrowShare[id][onBehalf] -= shares;
         totalBorrowShares[id] -= shares;
-
         totalBorrow[id] -= amount;
 
         market.borrowableAsset.safeTransferFrom(msg.sender, address(this), amount);
@@ -194,6 +208,7 @@ contract Blue {
     function withdrawCollateral(Market calldata market, uint256 amount, address onBehalf) external {
         Id id = market.id();
         require(lastUpdate[id] != 0, Errors.MARKET_NOT_CREATED);
+        if (amount == type(uint256).max) amount = collateral[id][msg.sender];
         require(amount != 0, Errors.ZERO_AMOUNT);
         require(_isSenderOrIsApproved(onBehalf), Errors.MANAGER_NOT_APPROVED);
 

test/forge/Blue.t.sol

@@ -360,6 +360,20 @@ contract BlueTest is Test {
         );
     }
 
+    function testWithdrawAll(uint256 amountLent, uint256 amountBorrowed) public {
+        amountLent = bound(amountLent, 1, 2 ** 64);
+        amountBorrowed = bound(amountBorrowed, 1, 2 ** 64);
+        vm.assume(amountLent >= amountBorrowed);
+
+        borrowableAsset.setBalance(address(this), amountLent);
+        blue.supply(market, amountLent, address(this));
+        blue.withdraw(market, type(uint256).max, address(this));
+
+        assertEq(blue.supplyShare(id, address(this)), 0, "supply share");
+        assertEq(borrowableAsset.balanceOf(address(this)), amountLent, "this balance");
+        assertEq(borrowableAsset.balanceOf(address(blue)), 0, "blue balance");
+    }
+
     function testCollateralRequirements(
         uint256 amountCollateral,
         uint256 amountBorrowed,
@@ -444,6 +458,23 @@ contract BlueTest is Test {
         assertEq(borrowableAsset.balanceOf(address(blue)), amountLent - amountBorrowed + amountRepaid, "blue balance");
     }
 
+    function testRepayAll(uint256 amountLent, uint256 amountBorrowed) public {
+        amountLent = bound(amountLent, 1, 2 ** 64);
+        amountBorrowed = bound(amountBorrowed, 1, amountLent);
+
+        borrowableAsset.setBalance(address(this), amountLent);
+        blue.supply(market, amountLent, address(this));
+
+        vm.startPrank(BORROWER);
+        blue.borrow(market, amountBorrowed, BORROWER);
+        blue.repay(market, type(uint256).max, BORROWER);
+        vm.stopPrank();
+
+        assertEq(blue.borrowShare(id, BORROWER), 0, "borrow share");
+        assertEq(borrowableAsset.balanceOf(BORROWER), 0, "BORROWER balance");
+        assertEq(borrowableAsset.balanceOf(address(blue)), amountLent, "blue balance");
+    }
+
     function testSupplyCollateralOnBehalf(uint256 amount, address onBehalf) public {
         vm.assume(onBehalf != address(blue));
         amount = bound(amount, 1, 2 ** 64);
@@ -476,6 +507,18 @@ contract BlueTest is Test {
         assertEq(collateralAsset.balanceOf(address(blue)), amountDeposited - amountWithdrawn, "blue balance");
     }
 
+    function testWithdrawCollateralAll(uint256 amountDeposited) public {
+        amountDeposited = bound(amountDeposited, 1, 2 ** 64);
+
+        collateralAsset.setBalance(address(this), amountDeposited);
+        blue.supplyCollateral(market, amountDeposited, address(this));
+        blue.withdrawCollateral(market, type(uint256).max, address(this));
+
+        assertEq(blue.collateral(id, address(this)), 0, "this collateral");
+        assertEq(collateralAsset.balanceOf(address(this)), amountDeposited, "this balance");
+        assertEq(collateralAsset.balanceOf(address(blue)), 0, "blue balance");
+    }
+
     function testLiquidate(uint256 amountLent) public {
         borrowableOracle.setPrice(1e18);
         amountLent = bound(amountLent, 1000, 2 ** 64);