repay/withdraw shares input

src/Blue.sol

@@ -151,18 +151,18 @@ contract Blue is IFlashLender {
         market.borrowableAsset.safeTransferFrom(msg.sender, address(this), amount);
     }
 
-    function withdraw(Market memory market, uint256 amount, address onBehalf, address receiver) external {
+    function withdraw(Market calldata market, uint256 shares, address onBehalf, address receiver) external {
         Id id = market.id();
         require(lastUpdate[id] != 0, Errors.MARKET_NOT_CREATED);
-        require(amount != 0, Errors.ZERO_AMOUNT);
+        require(shares != 0, Errors.ZERO_AMOUNT);
         require(_isSenderAuthorized(onBehalf), Errors.UNAUTHORIZED);
 
         _accrueInterests(market, id);
 
-        uint256 shares = amount.toSharesUp(totalSupply[id], totalSupplyShares[id]);
+        uint256 amount = shares.toAssetsDown(totalSupply[id], totalSupplyShares[id]);
+
         supplyShare[id][onBehalf] -= shares;
         totalSupplyShares[id] -= shares;
-
         totalSupply[id] -= amount;
 
         require(totalBorrow[id] <= totalSupply[id], Errors.INSUFFICIENT_LIQUIDITY);
@@ -192,17 +192,17 @@ contract Blue is IFlashLender {
         market.borrowableAsset.safeTransfer(receiver, amount);
     }
 
-    function repay(Market memory market, uint256 amount, address onBehalf, bytes calldata data) external {
+    function repay(Market calldata market, uint256 shares, address onBehalf, bytes calldata data) external {
         Id id = market.id();
+        require(shares != 0, Errors.ZERO_AMOUNT);
         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 amount = shares.toAssetsUp(totalBorrow[id], totalBorrowShares[id]);
+
         borrowShare[id][onBehalf] -= shares;
         totalBorrowShares[id] -= shares;
-
         totalBorrow[id] -= amount;
 
         if (data.length > 0) IBlueRepayCallback(msg.sender).onBlueRepay(amount, data);

test/forge/Blue.t.sol

@@ -7,6 +7,7 @@ import "forge-std/console.sol";
 import {SigUtils} from "./helpers/SigUtils.sol";
 
 import "src/Blue.sol";
+import {SharesMath} from "src/libraries/SharesMath.sol";
 import {
     IBlueLiquidateCallback,
     IBlueRepayCallback,
@@ -339,12 +340,11 @@ contract BlueTest is
         assertEq(borrowableAsset.balanceOf(address(blue)), amountLent - amountBorrowed, "blue balance");
     }
 
-    function testWithdraw(uint256 amountLent, uint256 amountWithdrawn, uint256 amountBorrowed, address receiver)
+    function testWithdraw(uint256 amountLent, uint256 sharesWithdrawn, uint256 amountBorrowed, address receiver)
         public
     {
         vm.assume(receiver != address(blue));
         amountLent = bound(amountLent, 1, 2 ** 64);
-        amountWithdrawn = bound(amountWithdrawn, 1, 2 ** 64);
         amountBorrowed = bound(amountBorrowed, 1, 2 ** 64);
         vm.assume(amountLent >= amountBorrowed);
 
@@ -354,30 +354,45 @@ contract BlueTest is
         vm.prank(BORROWER);
         blue.borrow(market, amountBorrowed, BORROWER, BORROWER);
 
+        uint256 supplyShareBefore = blue.supplyShare(id, address(this));
+        sharesWithdrawn = bound(sharesWithdrawn, 1, 2 ** 64);
+        uint256 amountWithdrawn =
+            SharesMath.toAssetsDown(sharesWithdrawn, blue.totalSupply(id), blue.totalSupplyShares(id));
+
         if (amountWithdrawn > amountLent - amountBorrowed) {
-            if (amountWithdrawn > amountLent) {
+            if (sharesWithdrawn > blue.supplyShare(id, address(this))) {
                 vm.expectRevert();
             } else {
                 vm.expectRevert(bytes(Errors.INSUFFICIENT_LIQUIDITY));
             }
-            blue.withdraw(market, amountWithdrawn, address(this), receiver);
+            blue.withdraw(market, sharesWithdrawn, address(this), receiver);
             return;
         }
 
-        blue.withdraw(market, amountWithdrawn, address(this), receiver);
+        blue.withdraw(market, sharesWithdrawn, address(this), receiver);
 
-        assertApproxEqAbs(
-            blue.supplyShare(id, address(this)),
-            (amountLent - amountWithdrawn) * SharesMath.VIRTUAL_SHARES,
-            100,
-            "supply share"
-        );
+        assertEq(blue.supplyShare(id, address(this)), supplyShareBefore - sharesWithdrawn, "supply share");
         assertEq(borrowableAsset.balanceOf(receiver), amountWithdrawn, "receiver balance");
         assertEq(
             borrowableAsset.balanceOf(address(blue)), amountLent - amountBorrowed - amountWithdrawn, "blue balance"
         );
     }
 
+    function testWithdrawAll(uint256 amountLent, uint256 amountBorrowed, address receiver) public {
+        vm.assume(receiver != address(blue));
+        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), hex"");
+        blue.withdraw(market, blue.supplyShare(id, address(this)), address(this), receiver);
+
+        assertEq(blue.supplyShare(id, address(this)), 0, "supply share");
+        assertEq(borrowableAsset.balanceOf(receiver), amountLent, "receiver balance");
+        assertEq(borrowableAsset.balanceOf(address(blue)), 0, "blue balance");
+    }
+
     function testCollateralRequirements(
         uint256 amountCollateral,
         uint256 amountBorrowed,
@@ -412,56 +427,70 @@ contract BlueTest is
         }
     }
 
-    function testRepay(uint256 amountLent, uint256 amountBorrowed, uint256 amountRepaid) public {
+    function testRepay(uint256 amountLent, uint256 amountBorrowed, uint256 sharesRepaid) public {
         amountLent = bound(amountLent, 1, 2 ** 64);
         amountBorrowed = bound(amountBorrowed, 1, amountLent);
-        amountRepaid = bound(amountRepaid, 1, amountBorrowed);
 
         borrowableAsset.setBalance(address(this), amountLent);
         blue.supply(market, amountLent, address(this), hex"");
 
-        vm.startPrank(BORROWER);
+        vm.prank(BORROWER);
         blue.borrow(market, amountBorrowed, BORROWER, BORROWER);
-        blue.repay(market, amountRepaid, BORROWER, hex"");
-        vm.stopPrank();
 
-        assertApproxEqAbs(
-            blue.borrowShare(id, BORROWER),
-            (amountBorrowed - amountRepaid) * SharesMath.VIRTUAL_SHARES,
-            100,
-            "borrow share"
-        );
+        uint256 borrowShareBefore = blue.borrowShare(id, BORROWER);
+        sharesRepaid = bound(sharesRepaid, 1, borrowShareBefore);
+        uint256 amountRepaid = SharesMath.toAssetsUp(sharesRepaid, blue.totalBorrow(id), blue.totalBorrowShares(id));
+        vm.prank(BORROWER);
+        blue.repay(market, sharesRepaid, BORROWER, hex"");
+
+        assertEq(blue.borrowShare(id, BORROWER), borrowShareBefore - sharesRepaid, "borrow share");
         assertEq(borrowableAsset.balanceOf(BORROWER), amountBorrowed - amountRepaid, "BORROWER balance");
         assertEq(borrowableAsset.balanceOf(address(blue)), amountLent - amountBorrowed + amountRepaid, "blue balance");
     }
 
-    function testRepayOnBehalf(uint256 amountLent, uint256 amountBorrowed, uint256 amountRepaid, address onBehalf)
+    function testRepayOnBehalf(uint256 amountLent, uint256 amountBorrowed, uint256 sharesRepaid, address onBehalf)
         public
     {
         vm.assume(onBehalf != address(blue));
         vm.assume(onBehalf != address(this));
         amountLent = bound(amountLent, 1, 2 ** 64);
         amountBorrowed = bound(amountBorrowed, 1, amountLent);
-        amountRepaid = bound(amountRepaid, 1, amountBorrowed);
 
-        borrowableAsset.setBalance(address(this), amountLent + amountRepaid);
+        borrowableAsset.setBalance(address(this), amountLent);
         blue.supply(market, amountLent, address(this), hex"");
 
         vm.prank(onBehalf);
         blue.borrow(market, amountBorrowed, onBehalf, onBehalf);
 
-        blue.repay(market, amountRepaid, onBehalf, hex"");
+        uint256 borrowShareBefore = blue.borrowShare(id, onBehalf);
+        sharesRepaid = bound(sharesRepaid, 1, borrowShareBefore);
+        uint256 amountRepaid = SharesMath.toAssetsUp(sharesRepaid, blue.totalBorrow(id), blue.totalBorrowShares(id));
+        borrowableAsset.setBalance(address(this), amountRepaid);
 
-        assertApproxEqAbs(
-            blue.borrowShare(id, onBehalf),
-            (amountBorrowed - amountRepaid) * SharesMath.VIRTUAL_SHARES,
-            100,
-            "borrow share"
-        );
+        blue.repay(market, sharesRepaid, onBehalf, hex"");
+
+        assertEq(blue.borrowShare(id, onBehalf), borrowShareBefore - sharesRepaid, "borrow share");
         assertEq(borrowableAsset.balanceOf(onBehalf), amountBorrowed, "onBehalf balance");
         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), hex"");
+
+        vm.startPrank(BORROWER);
+        blue.borrow(market, amountBorrowed, BORROWER, BORROWER);
+        blue.repay(market, blue.borrowShare(id, BORROWER), BORROWER, hex"");
+        vm.stopPrank();
+
+        assertEq(blue.borrowShare(id, BORROWER), 0, "borrow share");
+        assertEq(borrowableAsset.balanceOf(BORROWER), 0, "receiver 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);
@@ -495,6 +524,19 @@ contract BlueTest is
         assertEq(collateralAsset.balanceOf(address(blue)), amountDeposited - amountWithdrawn, "blue balance");
     }
 
+    function testWithdrawCollateralAll(uint256 amountDeposited, address receiver) public {
+        vm.assume(receiver != address(blue));
+        amountDeposited = bound(amountDeposited, 1, 2 ** 64);
+
+        collateralAsset.setBalance(address(this), amountDeposited);
+        blue.supplyCollateral(market, amountDeposited, address(this), hex"");
+        blue.withdrawCollateral(market, blue.collateral(id, address(this)), address(this), receiver);
+
+        assertEq(blue.collateral(id, address(this)), 0, "this collateral");
+        assertEq(collateralAsset.balanceOf(receiver), amountDeposited, "receiver balance");
+        assertEq(collateralAsset.balanceOf(address(blue)), 0, "blue balance");
+    }
+
     function testLiquidate(uint256 amountLent) public {
         borrowableOracle.setPrice(1e18);
         amountLent = bound(amountLent, 1000, 2 ** 64);
@@ -830,7 +872,12 @@ contract BlueTest is
         );
         assertGt(blue.borrowShare(market.id(), address(this)), 0, "no borrow");
 
-        blue.repay(market, toBorrow, address(this), abi.encode(this.testFlashActions.selector, abi.encode(amount)));
+        blue.repay(
+            market,
+            blue.borrowShare(id, address(this)),
+            address(this),
+            abi.encode(this.testFlashActions.selector, abi.encode(amount))
+        );
         assertEq(blue.collateral(market.id(), address(this)), 0, "no withdraw collateral");
     }
 

test/hardhat/Blue.spec.ts

@@ -129,25 +129,28 @@ describe("Blue", () => {
 
       const user = signers[i];
 
-      let amount = BigNumber.WAD.mul(1 + Math.floor(random() * 100));
-
+      let amount = BigNumber.WAD.mul(1 + Math.floor(random() * 100));      
+      
       if (random() < 2 / 3) {
+        const totalSupply = await blue.totalSupply(id);
+        const totalSupplyShares = await blue.totalSupplyShares(id);
         Promise.all([
           blue.connect(user).supply(market, amount, user.address, []),
-          blue.connect(user).withdraw(market, amount.div(2), user.address, user.address),
+          blue.connect(user).withdraw(market, amount.mul(totalSupplyShares.add(BigNumber.WAD)).div(totalSupply.add(1)).div(2), user.address, user.address),
         ]);
       } else {
         const totalSupply = await blue.totalSupply(id);
         const totalBorrow = await blue.totalBorrow(id);
+        const totalBorrowShares = await blue.totalBorrowShares(id);
         const liquidity = BigNumber.from(totalSupply).sub(BigNumber.from(totalBorrow));
-
+        
         amount = BigNumber.min(amount, BigNumber.from(liquidity).div(2));
-
+        
         if (amount > BigNumber.from(0)) {
           Promise.all([
             blue.connect(user).supplyCollateral(market, amount, user.address, []),
             blue.connect(user).borrow(market, amount.div(2), user.address, user.address),
-            blue.connect(user).repay(market, amount.div(4), user.address, []),
+            blue.connect(user).repay(market, amount.mul(totalBorrowShares.add(BigNumber.WAD)).div(totalBorrow.add(1)).div(4), user.address, []),
             blue.connect(user).withdrawCollateral(market, amount.div(8), user.address, user.address),
           ]);
         }