WIP: BalancerV3: DRY and boilerplate

src/core/BalancerV3.sol

@@ -7,6 +7,18 @@ import {ISignatureTransfer} from "@permit2/interfaces/ISignatureTransfer.sol";
 import {SafeTransferLib} from "../vendor/SafeTransferLib.sol";
 import {SettlerAbstract} from "../SettlerAbstract.sol";
 
+import {Panic} from "../utils/Panic.sol";
+import {UnsafeMath} from "../utils/UnsafeMath.sol";
+
+import {
+    TooMuchSlippage,
+    ZeroSellAmount
+} from "./SettlerErrors.sol";
+
+import {Encoder, NotesLib, StateLib, Decoder, Take} from "./FlashAccountingCommon.sol";
+
+import {FreeMemory} from "../utils/FreeMemory.sol";
+
 interface IBalancerV3Vault {
     /**
      * @notice Creates a context for a sequence of operations (i.e., "unlocks" the Vault).
@@ -131,7 +143,13 @@ interface IBalancerV3Callback {
     function balancerUnlockCallback(bytes calldata data) external returns (bytes memory);
 }
 
-abstract contract BalancerV3 is SettlerAbstract {
+abstract contract BalancerV3 is SettlerAbstract, FreeMemory {
+    using SafeTransferLib for IERC20;
+    using UnsafeMath for uint256;
+    using NotesLib for NotesLib.Note;
+    using NotesLib for NotesLib.Note[];
+    using StateLib for StateLib.State;
+
     constructor() {
         assert(BASIS == Encoder.BASIS);
         assert(BASIS == Decoder.BASIS);
@@ -197,7 +215,7 @@ abstract contract BalancerV3 is SettlerAbstract {
             amountOutMin
         );
         bytes memory encodedBuyAmount = _setOperatorAndCall(
-            address(VAULT), data, uint32(IBalancerV3Callback.unlockCallback.selector), _uniV4Callback
+            address(VAULT), data, uint32(IBalancerV3Callback.balancerUnlockCallback.selector), _balV3Callback
         );
         // buyAmount = abi.decode(abi.decode(encodedBuyAmount, (bytes)), (uint256));
         assembly ("memory-safe") {
@@ -227,27 +245,55 @@ abstract contract BalancerV3 is SettlerAbstract {
             mstore(add(0x20, swapParams), shr(0x60, calldataload(data.offset)))
             data.offset := add(0x14, data.offset)
             data.length := sub(data.length, 0x14)
+            // we don't check for array out-of-bounds here; we will check it later in `Decoder.decodeBytes`
         }
         swapParams.tokenIn = state.sell.token;
         swapParams.tokenOut = state.buy.token;
+        return data;
     }
 
-    function _decodeHookdataAndSwap(
+    function _decodeUserdataAndSwap(
         IBalancerV3Vault.VaultSwapParams memory swapParams,
         StateLib.State memory state,
         bytes calldata data
-    ) internal freeMemory returns (bytes calldata) {
+    ) internal DANGEROUS_freeMemory returns (bytes calldata) {
         (data, swapParams.userData) = Decoder.decodeBytes(data);
 
         (, uint256 amountIn, uint256 amountOut) = IBalancerV3Vault(msg.sender).swap(swapParams);
         state.sell.amount -= amountIn;
         state.buy.amount += amountOut;
 
-        swapParams.data = new bytes(0);
+        swapParams.userData = new bytes(0);
 
         return data;
     }
 
+    function _pay(
+        IERC20 sellToken,
+        address payer,
+        uint256 sellAmount,
+        ISignatureTransfer.PermitTransferFrom calldata permit,
+        bool isForwarded,
+        bytes calldata sig
+    ) private returns (uint256) {
+        if (payer == address(this)) {
+            sellToken.safeTransfer(msg.sender, sellAmount);
+        } else {
+            // assert(payer == address(0));
+            ISignatureTransfer.SignatureTransferDetails memory transferDetails =
+                ISignatureTransfer.SignatureTransferDetails({to: msg.sender, requestedAmount: sellAmount});
+            _transferFrom(permit, transferDetails, sig, isForwarded);
+        }
+        return IBalancerV3Vault(msg.sender).settle(sellToken, sellAmount);
+    }
+
+    // the mandatory fields are
+    // 2 - sell bps
+    // 1 - pool key tokens case
+    // 20 - pool
+    // 3 - user data length
+    uint256 private constant _HOP_DATA_LENGTH = 26;
+
     function balancerUnlockCallback(bytes calldata data) private returns (bytes memory) {
         address recipient;
         uint256 minBuyAmount;
@@ -283,7 +329,7 @@ abstract contract BalancerV3 is SettlerAbstract {
         IBalancerV3Vault.VaultSwapParams memory swapParams;
         swapParams.kind = IBalancerV3Vault.SwapKind.EXACT_IN;
         swapParams.limitRaw = 0; // TODO: price limits for partial filling
-        while (...) {
+        while (data.length >= _HOP_DATA_LENGTH) {
             uint16 bps;
             assembly ("memory-safe") {
                 bps := shr(0xf0, calldataload(data.offset))
@@ -298,7 +344,49 @@ abstract contract BalancerV3 is SettlerAbstract {
             data = Decoder.updateState(state, notes, data);
             data = _setSwapParams(swapParams, state, data);
             swapParams.amountGivenRaw = (state.sell.amount * bps).unsafeDiv(BASIS);
-            data = _decodeHookdataAndSwap(swapParams, state, data);
+            data = _decodeUserdataAndSwap(swapParams, state, data);
+        }
+
+
+        // `data` has been consumed. All that remains is to settle out the net result of all the
+        // swaps. Any credit in any token other than `state.buy.token` will be swept to
+        // Settler. `state.buy.token` will be sent to `recipient`.
+        {
+            (IERC20 globalSellToken, uint256 globalSellAmount) = (state.globalSell.token, state.globalSell.amount);
+            uint256 globalBuyAmount = Take.take(state, notes, uint32(IBalancerV3Vault.sendTo.selector), recipient, minBuyAmount);
+            if (feeOnTransfer) {
+                // We've already transferred the sell token to the pool manager and
+                // `settle`'d. `globalSellAmount` is the verbatim credit in that token stored by the
+                // pool manager. We only need to handle the case of incomplete filling.
+                if (globalSellAmount != 0) {
+                    Take._callSelector(uint32(IBalancerV3Vault.sendTo.selector), globalSellToken, payer == address(this) ? address(this) : _msgSender(), globalSellAmount);
+                }
+            } else {
+                // While `notes` records a credit value, the pool manager actually records a debt
+                // for the global sell token. We recover the exact amount of that debt and then pay
+                // it.
+                // `globalSellAmount` is _usually_ zero, but if it isn't it represents a partial
+                // fill. This subtraction recovers the actual debt recorded in the pool manager.
+                uint256 debt;
+                unchecked {
+                    debt = state.globalSellAmount - globalSellAmount;
+                }
+                if (debt == 0) {
+                    revert ZeroSellAmount(globalSellToken);
+                }
+                _pay(globalSellToken, payer, debt, permit, isForwarded, sig);
+            }
+
+            bytes memory returndata;
+            assembly ("memory-safe") {
+                returndata := mload(0x40)
+                mstore(returndata, 0x60)
+                mstore(add(0x20, returndata), 0x20)
+                mstore(add(0x40, returndata), 0x20)
+                mstore(add(0x60, returndata), globalBuyAmount)
+                mstore(0x40, add(0x80, returndata))
+            }
+            return returndata;
         }
     }
 }

src/core/FlashAccountingCommon.sol

@@ -8,7 +8,7 @@ import {SafeTransferLib} from "../vendor/SafeTransferLib.sol";
 import {Panic} from "../utils/Panic.sol";
 import {UnsafeMath} from "../utils/UnsafeMath.sol";
 
-import {BoughtSellToken} from "./SettlerErrors.sol";
+import {TooMuchSlippage, BoughtSellToken} from "./SettlerErrors.sol";
 
 /// This library is a highly-optimized, in-memory, enumerable mapping from tokens to amounts. It
 /// consists of 2 components that must be kept synchronized. There is a `memory` array of `Note`
@@ -584,3 +584,76 @@ library Decoder {
         newData = data;
     }
 }
+
+library Take {
+    using UnsafeMath for uint256;
+    using NotesLib for NotesLib.Note;
+    using NotesLib for NotesLib.Note[];
+
+    function _callSelector(uint256 selector, IERC20 token, address to, uint256 amount) internal {
+        assembly ("memory-safe") {
+            token := and(0xffffffffffffffffffffffffffffffffffffffff, token)
+            if iszero(amount) {
+                mstore(0x00, 0xcbf0dbf5) // selector for `ZeroBuyAmount(address)`
+                mstore(0x20, token)
+                revert(0x1c, 0x24)
+            }
+            let ptr := mload(0x40)
+            mstore(ptr, selector)
+            if eq(token, 0xeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee) { token := 0x00 }
+            mstore(add(0x20, ptr), token)
+            mstore(add(0x40, ptr), and(0xffffffffffffffffffffffffffffffffffffffff, to))
+            mstore(add(0x60, ptr), amount)
+            if iszero(call(gas(), caller(), 0x00, add(0x1c, ptr), 0x64, 0x00, 0x00)) {
+                returndatacopy(ptr, 0x00, returndatasize())
+                revert(ptr, returndatasize())
+            }
+        }
+    }
+
+    /// `_take` is responsible for removing the accumulated credit in each token from the vault. The
+    /// current `state.buy` is the global buy token. We return the settled amount of that token
+    /// (`buyAmount`), after checking it against the slippage limit (`minBuyAmount`). Each token
+    /// with credit causes a corresponding call to `msg.sender.<selector>(token, recipient,
+    /// amount)`.
+    function take(StateLib.State memory state, NotesLib.Note[] memory notes, uint32 selector, address recipient, uint256 minBuyAmount)
+        internal
+        returns (uint256 buyAmount)
+    {
+        notes.del(state.buy);
+        if (state.sell.amount == 0) {
+            notes.del(state.sell);
+        }
+
+        uint256 length = notes.length;
+        // `length` of zero implies that we fully liquidated the global sell token (there is no
+        // `amount` remaining) and that the only token in which we have credit is the global buy
+        // token. We're about to `take` that token below.
+        if (length != 0) {
+            {
+                NotesLib.Note memory firstNote = notes[0]; // out-of-bounds is impossible
+                if (!firstNote.eq(state.globalSell)) {
+                    // The global sell token being in a position other than the 1st would imply that
+                    // at some point we _bought_ that token. This is illegal and results in a revert
+                    // with reason `BoughtSellToken(address)`.
+                    _callSelector(selector, firstNote.token, address(this), firstNote.amount);
+                }
+            }
+            for (uint256 i = 1; i < length; i = i.unsafeInc()) {
+                (IERC20 token, uint256 amount) = notes.unsafeGet(i);
+                _callSelector(selector, token, address(this), amount);
+            }
+        }
+
+        // The final token to be bought is considered the global buy token. We bypass `notes` and
+        // read it directly from `state`. Check the slippage limit. Transfer to the recipient.
+        {
+            IERC20 buyToken = state.buy.token;
+            buyAmount = state.buy.amount;
+            if (buyAmount < minBuyAmount) {
+                revert TooMuchSlippage(buyToken, minBuyAmount, buyAmount);
+            }
+            _callSelector(selector, buyToken, recipient, buyAmount);
+        }
+    }
+}

src/core/UniswapV4.sol

@@ -13,17 +13,13 @@ import {
     TooMuchSlippage,
     DeltaNotPositive,
     DeltaNotNegative,
-    ZeroSellAmount,
-    ZeroBuyAmount,
-    BoughtSellToken,
-    TokenHashCollision,
-    ZeroToken
+    ZeroSellAmount
 } from "./SettlerErrors.sol";
 
 import {
     BalanceDelta, IHooks, IPoolManager, UnsafePoolManager, POOL_MANAGER, IUnlockCallback
 } from "./UniswapV4Types.sol";
-import {Encoder, NotesLib, StateLib, Decoder} from "./FlashAccountingCommon.sol";
+import {Encoder, NotesLib, StateLib, Decoder, Take} from "./FlashAccountingCommon.sol";
 
 library CreditDebt {
     using UnsafeMath for int256;
@@ -245,51 +241,6 @@ abstract contract UniswapV4 is SettlerAbstract {
         return (zeroForOne, data);
     }
 
-    /// `_take` is responsible for removing the accumulated credit in each token from the pool
-    /// manager. The current `state.buy` is the global buy token. We return the settled amount of
-    /// that token (`buyAmount`), after checking it against the slippage limit
-    /// (`minBuyAmount`). Each token with credit causes a corresponding call to `POOL_MANAGER.take`.
-    function _take(StateLib.State memory state, NotesLib.Note[] memory notes, address recipient, uint256 minBuyAmount)
-        private
-        returns (uint256 buyAmount)
-    {
-        notes.del(state.buy);
-        if (state.sell.amount == 0) {
-            notes.del(state.sell);
-        }
-
-        uint256 length = notes.length;
-        // `length` of zero implies that we fully liquidated the global sell token (there is no
-        // `amount` remaining) and that the only token in which we have credit is the global buy
-        // token. We're about to `take` that token below.
-        if (length != 0) {
-            {
-                NotesLib.Note memory firstNote = notes[0]; // out-of-bounds is impossible
-                if (!firstNote.eq(state.globalSell)) {
-                    // The global sell token being in a position other than the 1st would imply that
-                    // at some point we _bought_ that token. This is illegal and results in a revert
-                    // with reason `BoughtSellToken(address)`.
-                    IPoolManager(msg.sender).unsafeTake(firstNote.token, address(this), firstNote.amount);
-                }
-            }
-            for (uint256 i = 1; i < length; i = i.unsafeInc()) {
-                (IERC20 token, uint256 amount) = notes.unsafeGet(i);
-                IPoolManager(msg.sender).unsafeTake(token, address(this), amount);
-            }
-        }
-
-        // The final token to be bought is considered the global buy token. We bypass `notes` and
-        // read it directly from `state`. Check the slippage limit. Transfer to the recipient.
-        {
-            IERC20 buyToken = state.buy.token;
-            buyAmount = state.buy.amount;
-            if (buyAmount < minBuyAmount) {
-                revert TooMuchSlippage(buyToken, minBuyAmount, buyAmount);
-            }
-            IPoolManager(msg.sender).unsafeTake(buyToken, recipient, buyAmount);
-        }
-    }
-
     function _pay(
         IERC20 sellToken,
         address payer,
@@ -391,7 +342,7 @@ abstract contract UniswapV4 is SettlerAbstract {
         // Settler. `state.buy.token` will be sent to `recipient`.
         {
             (IERC20 globalSellToken, uint256 globalSellAmount) = (state.globalSell.token, state.globalSell.amount);
-            uint256 globalBuyAmount = _take(state, notes, recipient, minBuyAmount);
+            uint256 globalBuyAmount = Take.take(state, notes, uint32(IPoolManager.take.selector), recipient, minBuyAmount);
             if (feeOnTransfer) {
                 // We've already transferred the sell token to the pool manager and
                 // `settle`'d. `globalSellAmount` is the verbatim credit in that token stored by the

src/core/UniswapV4Types.sol

@@ -2,6 +2,7 @@
 pragma solidity ^0.8.25;
 
 import {IERC20} from "@forge-std/interfaces/IERC20.sol";
+import {Take} from "./FlashAccountingCommon.sol";
 
 type IHooks is address;
 
@@ -111,24 +112,7 @@ library UnsafePoolManager {
     }
 
     function unsafeTake(IPoolManager poolManager, IERC20 token, address to, uint256 amount) internal {
-        assembly ("memory-safe") {
-            token := and(0xffffffffffffffffffffffffffffffffffffffff, token)
-            if iszero(amount) {
-                mstore(0x00, 0xcbf0dbf5) // selector for `ZeroBuyAmount(address)`
-                mstore(0x20, token)
-                revert(0x1c, 0x24)
-            }
-            let ptr := mload(0x40)
-            mstore(ptr, 0x0b0d9c09) // selector for `take(address,address,uint256)`
-            if eq(token, 0xeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee) { token := 0x00 }
-            mstore(add(0x20, ptr), token)
-            mstore(add(0x40, ptr), and(0xffffffffffffffffffffffffffffffffffffffff, to))
-            mstore(add(0x60, ptr), amount)
-            if iszero(call(gas(), poolManager, 0x00, add(0x1c, ptr), 0x64, 0x00, 0x00)) {
-                returndatacopy(ptr, 0x00, returndatasize())
-                revert(ptr, returndatasize())
-            }
-        }
+        return Take._callSelector(uint32(IPoolManager.take.selector), token, to, amount);
     }
 
     function unsafeSwap(