Add settleAndRefund (#2)

* feat: Add settleAndRefund

* feat: Add to in parameter for settleAndRefund

* feat: update to Snoopy implementation

* feat: add negative balance delta test

* bug: move transfer only at the end to prevent reentrancy

.forge-snapshots/VaultTest#Vault.snap

@@ -1 +1 @@
-6762
+7075

.forge-snapshots/VaultTest#lockSettledWhenAddLiquidity.snap

@@ -0,0 +1 @@
+122540

.forge-snapshots/VaultTest#registerPoolManager.snap

@@ -1 +1 @@
-24484
+24506

.forge-snapshots/VaultTest#testLock_NoOp.snap

@@ -1 +1 @@
-11502
+11692

.forge-snapshots/VaultTest#testSettleAndRefund_WithErc20Transfer.snap

@@ -0,0 +1 @@
+56176

.forge-snapshots/VaultTest#testSettleAndRefund_WithoutErc20Transfer.snap

@@ -0,0 +1 @@
+36514

src/Vault.sol

@@ -126,6 +126,28 @@ contract Vault is IVault, VaultToken, Ownable {
         SettlementGuard.accountDelta(msg.sender, currency, -(paid.toInt128()));
     }
 
+    function settleAndRefund(Currency currency, address to)
+        external
+        payable
+        override
+        isLocked
+        returns (uint256 paid, uint256 refund)
+    {
+        paid = currency.balanceOfSelf() - reservesOfVault[currency];
+        int256 currentDelta = SettlementGuard.getCurrencyDelta(msg.sender, currency);
+
+        if (currentDelta >= 0 && paid > currentDelta.toUint256()) {
+            // msg.sender owes vault but paid more than than whats owed
+            refund = paid - currentDelta.toUint256();
+            paid = currentDelta.toUint256();
+        }
+
+        reservesOfVault[currency] += paid;
+        SettlementGuard.accountDelta(msg.sender, currency, -(paid.toInt128()));
+
+        if (refund > 0) currency.transfer(to, refund);
+    }
+
     /// @inheritdoc IVault
     function settleFor(Currency currency, address target, uint256 amount) external isLocked {
         /// @notice settle all outstanding debt if amount is 0

src/interfaces/IVault.sol

@@ -67,6 +67,13 @@ interface IVault is IVaultToken {
     /// @notice Called by the user to pay what is owed
     function settle(Currency token) external payable returns (uint256 paid);
 
+    /// @notice Called by the user to pay what is owed. If the payment is more than the debt, the surplus is refunded
+    /// @param currency The currency to settle
+    /// @param to The address to refund the surplus to
+    /// @return paid The amount paid
+    /// @return refund The amount refunded
+    function settleAndRefund(Currency currency, address to) external payable returns (uint256 paid, uint256 refund);
+
     /// @notice move the delta from target to the msg.sender, only payment delta can be moved
     /// @param currency The currency to settle
     /// @param target The address whose delta will be updated

test/vault/FakePoolManagerRouter.sol

@@ -128,6 +128,18 @@ contract FakePoolManagerRouter {
             // settle ETH
             vault.settle{value: 5 ether}(CurrencyLibrary.NATIVE);
             vault.take(CurrencyLibrary.NATIVE, address(this), 5 ether);
+        } else if (data[0] == 0x18) {
+            // call this method via vault.lock(abi.encodePacked(hex"18", alice));
+            address to = address(uint160(uint256(bytes32(data[1:0x15]) >> 96)));
+            vault.settleAndRefund(poolKey.currency0, to);
+            vault.settleAndRefund(poolKey.currency1, to);
+        } else if (data[0] == 0x19) {
+            poolManager.mockAccounting(poolKey, 3 ether, -3 ether);
+            vault.settle(poolKey.currency0);
+
+            /// try to call settleAndRefund should not revert
+            vault.settleAndRefund(poolKey.currency1, address(this));
+            vault.take(poolKey.currency1, address(this), 3 ether);
         }
 
         return "";

test/vault/Vault.t.sol

@@ -173,6 +173,55 @@ contract VaultTest is Test, GasSnapshot {
         vault.lock(hex"02");
     }
 
+    function testSettleAndRefund_WithErc20Transfer() public {
+        address alice = makeAddr("alice");
+
+        // simulate someone transferred token to vault
+        currency0.transfer(address(vault), 10 ether);
+        assertEq(IERC20(Currency.unwrap(currency0)).balanceOf(address(fakePoolManagerRouter)), 0 ether);
+        assertEq(IERC20(Currency.unwrap(currency0)).balanceOf(address(alice)), 0 ether);
+
+        // settle and refund
+        vm.prank(address(fakePoolManagerRouter));
+        snapStart("VaultTest#testSettleAndRefund_WithErc20Transfer");
+        vault.lock(abi.encodePacked(hex"18", alice));
+        snapEnd();
+
+        // verify
+        assertEq(IERC20(Currency.unwrap(currency0)).balanceOf(address(fakePoolManagerRouter)), 0 ether);
+        assertEq(IERC20(Currency.unwrap(currency0)).balanceOf(address(alice)), 10 ether);
+    }
+
+    function testSettleAndRefund_WithoutErc20Transfer() public {
+        address alice = makeAddr("alice");
+
+        assertEq(IERC20(Currency.unwrap(currency0)).balanceOf(address(fakePoolManagerRouter)), 0 ether);
+        assertEq(IERC20(Currency.unwrap(currency0)).balanceOf(address(alice)), 0 ether);
+
+        // settleAndRefund works even if there's no excess currency
+        vm.prank(address(fakePoolManagerRouter));
+        snapStart("VaultTest#testSettleAndRefund_WithoutErc20Transfer");
+        vault.lock(abi.encodePacked(hex"18", alice));
+        snapEnd();
+
+        // verify
+        assertEq(IERC20(Currency.unwrap(currency0)).balanceOf(address(fakePoolManagerRouter)), 0 ether);
+        assertEq(IERC20(Currency.unwrap(currency0)).balanceOf(address(alice)), 0 ether);
+    }
+
+    function testSettleAndRefund_NegativeBalanceDelta() public {
+        // pre-req: ensure vault has some value in reserveOfVault[] before
+        currency0.transfer(address(vault), 10 ether);
+        currency1.transfer(address(vault), 10 ether);
+        vm.prank(address(fakePoolManagerRouter));
+        vault.lock(hex"02");
+
+        // settleAndRefund should not revert even if negative balanceDelta
+        currency0.transfer(address(vault), 3 ether);
+        vm.prank(address(fakePoolManagerRouter));
+        vault.lock(hex"19");
+    }
+
     function testNotCorrectPoolManager() public {
         // router => vault.lock
         // vault.lock => periphery.lockAcquired
@@ -208,7 +257,7 @@ contract VaultTest is Test, GasSnapshot {
         vm.prank(address(fakePoolManagerRouter));
         snapStart("VaultTest#lockSettledWhenAddLiquidity");
         vault.lock(hex"02");
-        snapStart("end");
+        snapEnd();
 
         assertEq(IERC20(Currency.unwrap(currency0)).balanceOf(address(vault)), 10 ether);
         assertEq(IERC20(Currency.unwrap(currency1)).balanceOf(address(vault)), 10 ether);

test/vault/VaultInvariant.t.sol

@@ -31,6 +31,7 @@ contract VaultPoolManager is Test {
     enum ActionType {
         Take,
         Settle,
+        SettleAndRefund,
         SettleFor,
         Mint,
         Burn
@@ -84,6 +85,21 @@ contract VaultPoolManager is Test {
         vault.lock(abi.encode(Action(ActionType.Settle, uint128(amt0), uint128(amt1))));
     }
 
+    /// @dev In settleAndRefund case, assume user add liquidity and paying to the vault
+    ///      but theres another folk who minted extra token to the vault
+    function settleAndRefund(uint256 amt0, uint256 amt1, bool sendToVault) public {
+        amt0 = bound(amt0, 0, MAX_TOKEN_BALANCE - 1 ether);
+        amt1 = bound(amt1, 0, MAX_TOKEN_BALANCE - 1 ether);
+
+        // someone send some token directly to vault
+        if (sendToVault) token0.mint(address(vault), 1 ether);
+
+        // mint token to VaultPoolManager, so VaultPoolManager can pay to the vault
+        token0.mint(address(this), amt0);
+        token1.mint(address(this), amt1);
+        vault.lock(abi.encode(Action(ActionType.SettleAndRefund, uint128(amt0), uint128(amt1))));
+    }
+
     /// @dev In settleFor case, assume user is paying for hook
     function settleFor(uint256 amt0, uint256 amt1) public {
         amt0 = bound(amt0, 0, MAX_TOKEN_BALANCE);
@@ -165,6 +181,15 @@ contract VaultPoolManager is Test {
 
             vault.settle(currency0);
             vault.settle(currency1);
+        } else if (action.actionType == ActionType.SettleAndRefund) {
+            BalanceDelta delta = toBalanceDelta(int128(action.amt0), int128(action.amt1));
+            vault.accountPoolBalanceDelta(poolKey, delta, address(this));
+
+            token0.transfer(address(vault), action.amt0);
+            token1.transfer(address(vault), action.amt1);
+
+            vault.settleAndRefund(currency0, address(this));
+            vault.settleAndRefund(currency1, address(this));
         } else if (action.actionType == ActionType.SettleFor) {
             // hook cash out the fee ahead
             BalanceDelta delta = toBalanceDelta(int128(action.amt0), int128(action.amt1));
@@ -212,13 +237,14 @@ contract VaultInvariant is Test, GasSnapshot {
         // Only call vaultPoolManager, otherwise all other contracts deployed in setUp will be called
         targetContract(address(vaultPoolManager));
 
-        bytes4[] memory selectors = new bytes4[](6);
+        bytes4[] memory selectors = new bytes4[](7);
         selectors[0] = VaultPoolManager.take.selector;
         selectors[1] = VaultPoolManager.mint.selector;
         selectors[2] = VaultPoolManager.settle.selector;
         selectors[3] = VaultPoolManager.burn.selector;
         selectors[4] = VaultPoolManager.settleFor.selector;
         selectors[5] = VaultPoolManager.collectFee.selector;
+        selectors[6] = VaultPoolManager.settleAndRefund.selector;
         targetSelector(FuzzSelector({addr: address(vaultPoolManager), selectors: selectors}));
     }