Merge pull request #268 from sablierhq/prb/withdraw-max

feat: add `withdrawMax` function

src/SablierV2.sol

@@ -101,6 +101,9 @@ abstract contract SablierV2 is
     /// @inheritdoc ISablierV2
     function getRecipient(uint256 streamId) public view virtual override returns (address recipient);
 
+    /// @inheritdoc ISablierV2
+    function getWithdrawableAmount(uint256 streamId) public view virtual override returns (uint128 withdrawableAmount);
+
     /// @inheritdoc ISablierV2
     function isCancelable(uint256 streamId) public view virtual override returns (bool result);
 
@@ -211,7 +214,7 @@ abstract contract SablierV2 is
         uint256 streamId,
         address to,
         uint128 amount
-    ) external override streamExists(streamId) isAuthorizedForStream(streamId) {
+    ) public override streamExists(streamId) isAuthorizedForStream(streamId) {
         // Checks: the provided address is the recipient if `msg.sender` is the sender of the stream.
         if (_isCallerStreamSender(streamId) && to != getRecipient(streamId)) {
             revert Errors.SablierV2_WithdrawSenderUnauthorized(streamId, msg.sender, to);
@@ -226,6 +229,11 @@ abstract contract SablierV2 is
         _withdraw(streamId, to, amount);
     }
 
+    /// @inheritdoc ISablierV2
+    function withdrawMax(uint256 streamId, address to) external override {
+        withdraw(streamId, to, getWithdrawableAmount(streamId));
+    }
+
     /// @inheritdoc ISablierV2
     function withdrawMultiple(uint256[] calldata streamIds, address to, uint128[] calldata amounts) external override {
         // Checks: the provided address to withdraw to is not zero.

src/SablierV2Linear.sol

@@ -157,7 +157,9 @@ contract SablierV2Linear is
     }
 
     /// @inheritdoc ISablierV2
-    function getWithdrawableAmount(uint256 streamId) public view override returns (uint128 withdrawableAmount) {
+    function getWithdrawableAmount(
+        uint256 streamId
+    ) public view override(ISablierV2, SablierV2) returns (uint128 withdrawableAmount) {
         unchecked {
             withdrawableAmount = getStreamedAmount(streamId) - _streams[streamId].amounts.withdrawn;
         }

src/SablierV2Pro.sol

@@ -154,7 +154,9 @@ contract SablierV2Pro is
     }
 
     /// @inheritdoc ISablierV2
-    function getWithdrawableAmount(uint256 streamId) public view override returns (uint128 withdrawableAmount) {
+    function getWithdrawableAmount(
+        uint256 streamId
+    ) public view override(ISablierV2, SablierV2) returns (uint128 withdrawableAmount) {
         unchecked {
             withdrawableAmount = getStreamedAmount(streamId) - _streams[streamId].amounts.withdrawn;
         }

src/interfaces/ISablierV2.sol

@@ -173,7 +173,7 @@ interface ISablierV2 is
     /// @param newComptroller The address of the new SablierV2Comptroller contract.
     function setComptroller(ISablierV2Comptroller newComptroller) external;
 
-    /// @notice Withdraws tokens from the stream to the recipient's account.
+    /// @notice Withdraws the provided amount of tokens from the stream to the provide address `to`.
     ///
     /// @dev Emits a {Withdraw} and a {Transfer} event.
     ///
@@ -188,10 +188,24 @@ interface ISablierV2 is
     /// - `amount` must not be zero and must not exceed the withdrawable amount.
     ///
     /// @param streamId The id of the stream to withdraw.
-    /// @param to The address that receives the withdrawn tokens, if the `msg.sender` is not the stream sender.
+    /// @param to The address that receives the withdrawn tokens.
     /// @param amount The amount to withdraw, in units of the token's decimals.
     function withdraw(uint256 streamId, address to, uint128 amount) external;
 
+    /// @notice Withdraws the maximum withdrawable amount from the stream to the provided address `to`.
+    ///
+    /// @dev Emits a {Withdraw} and a {Transfer} event.
+    ///
+    /// Notes:
+    /// - All from `withdraw`.
+    ///
+    /// Requirements:
+    /// - All from `withdraw`.
+    ///
+    /// @param streamId The id of the stream to withdraw.
+    /// @param to The address that receives the withdrawn tokens.
+    function withdrawMax(uint256 streamId, address to) external;
+
     /// @notice Withdraws tokens from multiple streams to the provided address `to`.
     ///
     /// @dev Emits multiple {Withdraw} and {Transfer} events.

test/helpers/mocks/SablierV2Mock.t.sol

@@ -68,7 +68,7 @@ contract SablierV2Mock is SablierV2 {
         return 0;
     }
 
-    function getWithdrawableAmount(uint256 streamId) external pure override returns (uint128) {
+    function getWithdrawableAmount(uint256 streamId) public pure override returns (uint128) {
         streamId;
         return 0;
     }

test/unit/sablier-v2/linear/withdraw-max/withdrawMax.t.sol

@@ -0,0 +1,14 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity >=0.8.13 <0.9.0;
+
+import { ISablierV2 } from "src/interfaces/ISablierV2.sol";
+
+import { LinearTest } from "test/unit/sablier-v2/linear/LinearTest.t.sol";
+import { WithdrawMax_Test } from "test/unit/sablier-v2/shared/withdraw-max/withdrawMax.t.sol";
+
+contract WithdrawMax_LinearTest is LinearTest, WithdrawMax_Test {
+    function setUp() public virtual override(LinearTest, WithdrawMax_Test) {
+        WithdrawMax_Test.setUp();
+        sablierV2 = ISablierV2(linear);
+    }
+}

test/unit/sablier-v2/pro/withdraw-max/withdrawMax.t.sol

@@ -0,0 +1,14 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity >=0.8.13 <0.9.0;
+
+import { ISablierV2 } from "src/interfaces/ISablierV2.sol";
+
+import { ProTest } from "test/unit/sablier-v2/pro/ProTest.t.sol";
+import { WithdrawMax_Test } from "test/unit/sablier-v2/shared/withdraw-max/withdrawMax.t.sol";
+
+contract WithdrawMax_ProTest is ProTest, WithdrawMax_Test {
+    function setUp() public virtual override(ProTest, WithdrawMax_Test) {
+        WithdrawMax_Test.setUp();
+        sablierV2 = ISablierV2(pro);
+    }
+}

test/unit/sablier-v2/shared/get-withdrawn-amount/getWithdrawnAmount.t.sol

@@ -51,8 +51,8 @@ abstract contract GetWithdrawnAmount_Test is SharedTest {
         vm.warp({ timestamp: DEFAULT_START_TIME + timeWarp });
 
         // Bound the withdraw amount.
-        uint128 withdrawableAmount = sablierV2.getWithdrawableAmount(defaultStreamId);
-        withdrawAmount = boundUint128(withdrawAmount, 1, withdrawableAmount);
+        uint128 streamedAmount = sablierV2.getStreamedAmount(defaultStreamId);
+        withdrawAmount = boundUint128(withdrawAmount, 1, streamedAmount);
 
         // Make the withdrawal.
         sablierV2.withdraw({ streamId: defaultStreamId, to: users.recipient, amount: withdrawAmount });

test/unit/sablier-v2/shared/withdraw-max/withdrawMax.t.sol

@@ -0,0 +1,69 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity >=0.8.13 <0.9.0;
+
+import { IERC20 } from "@prb/contracts/token/erc20/IERC20.sol";
+
+import { Events } from "src/libraries/Events.sol";
+
+import { SharedTest } from "../SharedTest.t.sol";
+
+abstract contract WithdrawMax_Test is SharedTest {
+    uint256 internal defaultStreamId;
+
+    function setUp() public virtual override {
+        super.setUp();
+
+        // Create the default stream.
+        defaultStreamId = createDefaultStream();
+
+        // Make the recipient the caller in this test suite.
+        changePrank(users.recipient);
+    }
+
+    /// @dev it should make the withdrawal and delete the stream.
+    function test_WithdrawMax_CurrentTimeEqualToStopTime() external {
+        // Warp to the end of the stream.
+        vm.warp({ timestamp: DEFAULT_STOP_TIME });
+
+        // Make the withdrawal.
+        sablierV2.withdrawMax({ streamId: defaultStreamId, to: users.recipient });
+
+        // Assert that the stream was deleted.
+        assertDeleted(defaultStreamId);
+
+        // Assert that the NFT was not burned.
+        address actualNFTowner = sablierV2.ownerOf({ tokenId: defaultStreamId });
+        address expectedNFTOwner = users.recipient;
+        assertEq(actualNFTowner, expectedNFTOwner);
+    }
+
+    modifier currentTimeLessThanStopTime() {
+        _;
+    }
+
+    /// @dev it should make the max withdrawal, emit a Withdraw event, and update the withdrawn amount
+    function testFuzz_WithdrawMax(uint256 timeWarp) external currentTimeLessThanStopTime {
+        timeWarp = bound(timeWarp, DEFAULT_CLIFF_DURATION, DEFAULT_TOTAL_DURATION - 1);
+
+        // Warp into the future.
+        vm.warp({ timestamp: DEFAULT_START_TIME + timeWarp });
+
+        // Bound the withdraw amount.
+        uint128 withdrawAmount = sablierV2.getWithdrawableAmount(defaultStreamId);
+
+        // Expect the withdrawal to be made to the recipient.
+        vm.expectCall(address(dai), abi.encodeCall(IERC20.transfer, (users.recipient, withdrawAmount)));
+
+        // Expect an event to be emitted.
+        vm.expectEmit({ checkTopic1: true, checkTopic2: true, checkTopic3: false, checkData: true });
+        emit Events.Withdraw({ streamId: defaultStreamId, to: users.recipient, amount: withdrawAmount });
+
+        // Make the withdrawal.
+        sablierV2.withdrawMax(defaultStreamId, users.recipient);
+
+        // Assert that the withdrawn amount was updated.
+        uint128 actualWithdrawnAmount = sablierV2.getWithdrawnAmount(defaultStreamId);
+        uint128 expectedWithdrawnAmount = withdrawAmount;
+        assertEq(actualWithdrawnAmount, expectedWithdrawnAmount);
+    }
+}

test/unit/sablier-v2/shared/withdraw-max/withdrawMax.tree

@@ -0,0 +1,5 @@
+withdrawMax.t.sol
+├── when the current time is greater than or equal to the stop time
+│  └── it should make the max withdrawal and delete the stream
+└── when the current time is less than the stop time
+   └── it should make the max withdrawal, emit a Withdraw event, and update the withdrawn amount

test/unit/sablier-v2/shared/withdraw/withdraw.t.sol

@@ -181,7 +181,7 @@ abstract contract Withdraw_Test is SharedTest {
     }
 
     /// @dev it should make the withdrawal and delete the stream.
-    function test_Withdraw_StreamEnded()
+    function test_Withdraw_CurrentTimeEqualToStopTime()
         external
         streamExistent
         callerAuthorized
@@ -205,7 +205,7 @@ abstract contract Withdraw_Test is SharedTest {
         assertEq(actualNFTowner, expectedNFTOwner);
     }
 
-    modifier streamOngoing() {
+    modifier currentTimeLessThanStopTime() {
         // Warp to 2,600 seconds after the start time (26% of the default stream duration).
         vm.warp({ timestamp: DEFAULT_START_TIME + DEFAULT_TIME_WARP });
         _;
@@ -224,7 +224,7 @@ abstract contract Withdraw_Test is SharedTest {
         withdrawAmountNotZero
         withdrawAmountLessThanOrEqualToWithdrawableAmount
         callerSender
-        streamOngoing
+        currentTimeLessThanStopTime
     {
         timeWarp = bound(timeWarp, DEFAULT_CLIFF_DURATION, DEFAULT_TOTAL_DURATION - 1);
         vm.assume(to != address(0) && to.code.length == 0);
@@ -268,7 +268,7 @@ abstract contract Withdraw_Test is SharedTest {
         withdrawAmountNotZero
         withdrawAmountLessThanOrEqualToWithdrawableAmount
         callerSender
-        streamOngoing
+        currentTimeLessThanStopTime
         recipientContract
     {
         // Create the stream with the recipient as a contract.
@@ -296,7 +296,7 @@ abstract contract Withdraw_Test is SharedTest {
         withdrawAmountNotZero
         withdrawAmountLessThanOrEqualToWithdrawableAmount
         callerSender
-        streamOngoing
+        currentTimeLessThanStopTime
         recipientContract
         recipientImplementsHook
     {
@@ -325,7 +325,7 @@ abstract contract Withdraw_Test is SharedTest {
         withdrawAmountNotZero
         withdrawAmountLessThanOrEqualToWithdrawableAmount
         callerSender
-        streamOngoing
+        currentTimeLessThanStopTime
         recipientContract
         recipientImplementsHook
         recipientDoesNotRevert
@@ -361,7 +361,7 @@ abstract contract Withdraw_Test is SharedTest {
         withdrawAmountNotZero
         withdrawAmountLessThanOrEqualToWithdrawableAmount
         callerSender
-        streamOngoing
+        currentTimeLessThanStopTime
         recipientContract
         recipientImplementsHook
         recipientDoesNotRevert

test/unit/sablier-v2/shared/withdraw/withdraw.tree

@@ -24,9 +24,9 @@ withdraw.t.sol
                ├── when the caller is an approved operator
                │  └── it should make the withdrawal and update the withdrawn amount
                └── when the caller is the sender
-                  ├── when the stream ended
+                  ├── when the current time is equal to the stop time
                   │  └── it should make the withdrawal and delete the stream
-                  └── when the stream did not end
+                  └── when the current time is less than the stop time
                      ├── when the recipient is not a contract
                      │  └── it should make the withdrawal and update the withdrawn amount
                      └── when the recipient is a contract