feat: perform all transfers in executors

For organization (and thus safety) purposes.

Rename to RestrictTransferFrom.sol so that we can perform multiple transfer froms (upto an allowance) in the case of split swaps (where the split is the first swap).

TODO: Fix tests.

foundry/src/OneTransferFromOnly.sol

@@ -1,88 +0,0 @@
-// SPDX-License-Identifier: BUSL-1.1
-pragma solidity ^0.8.26;
-
-import "@interfaces/IExecutor.sol";
-import "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
-import "@permit2/src/interfaces/IAllowanceTransfer.sol";
-
-error OneTransferFromOnly__AddressZero();
-error OneTransferFromOnly__MultipleTransferFrom();
-
-/**
- * @title OneTransferFromOnly - Restrict to one transferFrom on approved params per swap
- * @dev Restricts to one `transferFrom` (using `permit2` or regular `transferFrom`)
- * per swap, while ensuring that the `transferFrom` is only performed on the input
- * token and the input amount, from the msg.sender's wallet that calls the main swap
- * method. Reverts if multiple `transferFrom`s are attempted.
- */
-contract OneTransferFromOnly {
-    using SafeERC20 for IERC20;
-
-    IAllowanceTransfer public immutable permit2;
-    // keccak256("Dispatcher#TOKEN_IN_SLOT")
-    uint256 private constant _TOKEN_IN_SLOT =
-        0x66f353cfe8e3cbe0d03292348fbf0fca32e6e07fa0c2a52b4aac22193ac3b894;
-    // keccak256("Dispatcher#AMOUNT_IN_SLOT")
-    uint256 private constant _AMOUNT_IN_SLOT =
-        0x1f40aa2d23d66d03722685ce02e5d3a95545dfc8e7c56d1026790aa30be48937;
-    // keccak256("Dispatcher#IS_PERMIT2_SLOT")
-    uint256 private constant _IS_PERMIT2_SLOT =
-        0x3162c9d1175ca0ca7441f87984fdac41bbfdb13246f42c8bb4414d345da39e2a;
-    // keccak256("Dispatcher#SENDER_SLOT")
-    uint256 private constant _SENDER_SLOT =
-        0x5dcc7974be5cb30f183f878073999aaa6620995b9e052ab5a713071ff60ae9b5;
-    // keccak256("Dispatcher#IS_TRANSFER_EXECUTED_SLOT")
-    uint256 private constant _IS_TRANSFER_EXECUTED_SLOT =
-        0x1c64085c839fc2ff0f0aad20613eb6d056a1024e5990211e9eb30824dcd128c2;
-
-    constructor(address _permit2) {
-        if (_permit2 == address(0)) {
-            revert OneTransferFromOnly__AddressZero();
-        }
-        permit2 = IAllowanceTransfer(_permit2);
-    }
-
-    // slither-disable-next-line assembly
-    function _tstoreTransferFromInfo(
-        address tokenIn,
-        uint256 amountIn,
-        bool isPermit2
-    ) internal {
-        assembly {
-            tstore(_TOKEN_IN_SLOT, tokenIn)
-            tstore(_AMOUNT_IN_SLOT, amountIn)
-            tstore(_IS_PERMIT2_SLOT, isPermit2)
-            tstore(_SENDER_SLOT, caller())
-            tstore(_IS_TRANSFER_EXECUTED_SLOT, false)
-        }
-    }
-
-    // slither-disable-next-line assembly
-    function _transfer(address receiver) internal {
-        address tokenIn;
-        uint256 amount;
-        bool isPermit2;
-        address sender;
-        bool isTransferExecuted;
-        assembly {
-            tokenIn := tload(_TOKEN_IN_SLOT)
-            amount := tload(_AMOUNT_IN_SLOT)
-            isPermit2 := tload(_IS_PERMIT2_SLOT)
-            sender := tload(_SENDER_SLOT)
-            isTransferExecuted := tload(_IS_TRANSFER_EXECUTED_SLOT)
-        }
-        if (isTransferExecuted) {
-            revert OneTransferFromOnly__MultipleTransferFrom();
-        }
-        assembly {
-            tstore(_IS_TRANSFER_EXECUTED_SLOT, true)
-        }
-        if (isPermit2) {
-            // Permit2.permit is already called from the TychoRouter
-            permit2.transferFrom(sender, receiver, uint160(amount), tokenIn);
-        } else {
-            // slither-disable-next-line arbitrary-send-erc20
-            IERC20(tokenIn).safeTransferFrom(sender, receiver, amount);
-        }
-    }
-}

foundry/src/RestrictTransferFrom.sol

@@ -0,0 +1,110 @@
+// SPDX-License-Identifier: BUSL-1.1
+pragma solidity ^0.8.26;
+
+import "@interfaces/IExecutor.sol";
+import "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
+import "@permit2/src/interfaces/IAllowanceTransfer.sol";
+
+error RestrictTransferFrom__AddressZero();
+error RestrictTransferFrom__ExceededTransferFromAllowance();
+error RestrictTransferFrom__UnknownTransferType();
+
+/**
+ * @title RestrictTransferFrom - Restrict transferFrom upto allowed amount of token
+ * @dev Restricts to one `transferFrom` (using `permit2` or regular `transferFrom`)
+ * per swap, while ensuring that the `transferFrom` is only performed on the input
+ * token upto input amount, from the msg.sender's wallet that calls the main swap
+ * method. Reverts if `transferFrom`s are attempted above this allowed amount.
+ */
+contract RestrictTransferFrom {
+    using SafeERC20 for IERC20;
+
+    IAllowanceTransfer public immutable permit2;
+    // keccak256("Dispatcher#TOKEN_IN_SLOT")
+    uint256 private constant _TOKEN_IN_SLOT =
+        0x66f353cfe8e3cbe0d03292348fbf0fca32e6e07fa0c2a52b4aac22193ac3b894;
+    // keccak256("Dispatcher#AMOUNT_ALLOWED_SLOT")
+    uint256 private constant _AMOUNT_ALLOWED_SLOT =
+        0xc76591aca92830b1554f3dcc7893e7519ec7c57bd4e64fec0c546d9078033291;
+    // keccak256("Dispatcher#IS_PERMIT2_SLOT")
+    uint256 private constant _IS_PERMIT2_SLOT =
+        0x3162c9d1175ca0ca7441f87984fdac41bbfdb13246f42c8bb4414d345da39e2a;
+    // keccak256("Dispatcher#SENDER_SLOT")
+    uint256 private constant _SENDER_SLOT =
+        0x5dcc7974be5cb30f183f878073999aaa6620995b9e052ab5a713071ff60ae9b5;
+    // keccak256("Dispatcher#AMOUNT_SPENT_SLOT")
+    uint256 private constant _AMOUNT_SPENT_SLOT =
+        0x56044a5eb3aa5bd3ad908b7f15d1e8cb830836bb4ad178a0bf08955c94c40d30;
+
+    constructor(address _permit2) {
+        if (_permit2 == address(0)) {
+            revert RestrictTransferFrom__AddressZero();
+        }
+        permit2 = IAllowanceTransfer(_permit2);
+    }
+
+    enum TransferType {
+        TransferFrom,
+        Transfer,
+        None
+    }
+
+    // slither-disable-next-line assembly
+    function _tstoreTransferFromInfo(
+        address tokenIn,
+        uint256 amountIn,
+        bool isPermit2
+    ) internal {
+        assembly {
+            tstore(_TOKEN_IN_SLOT, tokenIn)
+            tstore(_AMOUNT_IN_SLOT, amountIn)
+            tstore(_IS_PERMIT2_SLOT, isPermit2)
+            tstore(_SENDER_SLOT, caller())
+            tstore(_IS_TRANSFER_EXECUTED_SLOT, false)
+        }
+    }
+
+    // slither-disable-next-line assembly
+    function _transfer(
+        address receiver,
+        TransferType transferType,
+        address tokenIn,
+        uint256 amount
+    ) internal {
+        if (transferType == TransferType.TransferFrom){
+            bool isPermit2;
+            address sender;
+            bool isTransferExecuted;
+            assembly {
+                tokenIn := tload(_TOKEN_IN_SLOT)
+                amountPermitted := tload(_AMOUNT_IN_SLOT)
+                isPermit2 := tload(_IS_PERMIT2_SLOT)
+                sender := tload(_SENDER_SLOT)
+                amountSpent := tload(_IS_TRANSFER_EXECUTED_SLOT)
+            }
+            if (amount + amountSpent > amountPermitted) {
+                revert RestrictTransferFrom__ExceededTransferFromAllowance();
+            }
+            assembly {
+                tstore(_AMOUNT_SPENT_SLOT, amount)
+            }
+            if (isPermit2) {
+                // Permit2.permit is already called from the TychoRouter
+                permit2.transferFrom(sender, receiver, uint160(amount), tokenIn);
+            } else {
+                // slither-disable-next-line arbitrary-send-erc20
+                IERC20(tokenIn).safeTransferFrom(sender, receiver, amount);
+            }
+        } else if (transferType == TransferType.Transfer) {
+            if (tokenIn == address(0)) {
+                Address.sendValue(payable(receiver), amount);
+            } else {
+                IERC20(tokenIn).safeTransfer(receiver, amount);
+            }
+        } else if (transferType == TransferType.None) {
+            return;
+        } else {
+            revert RestrictTransferFrom__UnknownTransferType();
+        }
+    }
+}

foundry/src/TychoRouter.sol

@@ -14,7 +14,7 @@ import "@permit2/src/interfaces/IAllowanceTransfer.sol";
 import "./Dispatcher.sol";
 import {LibSwap} from "../lib/LibSwap.sol";
 import {IPoolManager} from "@uniswap/v4-core/src/interfaces/IPoolManager.sol";
-import {OneTransferFromOnly} from "./OneTransferFromOnly.sol";
+import {RestrictTransferFrom} from "./RestrictTransferFrom.sol";
 
 //                                         ✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷
 //                                   ✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷
@@ -71,7 +71,7 @@ contract TychoRouter is
     Dispatcher,
     Pausable,
     ReentrancyGuard,
-    OneTransferFromOnly
+RestrictTransferFrom
 {
     IWETH private immutable _weth;
 
@@ -93,7 +93,7 @@ contract TychoRouter is
         address indexed token, uint256 amount, address indexed receiver
     );
 
-    constructor(address _permit2, address weth) OneTransferFromOnly(_permit2) {
+    constructor(address _permit2, address weth) RestrictTransferFrom(_permit2) {
         if (_permit2 == address(0) || weth == address(0)) {
             revert TychoRouter__AddressZero();
         }

foundry/src/executors/BalancerV2Executor.sol

@@ -10,15 +10,16 @@ import {
 import {IAsset} from "@balancer-labs/v2-interfaces/contracts/vault/IAsset.sol";
 // slither-disable-next-line solc-version
 import {IVault} from "@balancer-labs/v2-interfaces/contracts/vault/IVault.sol";
+import {RestrictTransferFrom} from "../RestrictTransferFrom.sol";
 
-error BalancerV2Executor__InvalidDataLength();
+    error BalancerV2Executor__InvalidDataLength();
 
-contract BalancerV2Executor is IExecutor {
+contract BalancerV2Executor is IExecutor, RestrictTransferFrom {
     using SafeERC20 for IERC20;
 
     address private constant VAULT = 0xBA12222222228d8Ba445958a75a0704d566BF2C8;
 
-    constructor(address _permit2) {}
+    constructor(address _permit2) RestrictTransferFrom(_permit2) {}
 
     // slither-disable-next-line locked-ether
     function swap(uint256 givenAmount, bytes calldata data)
@@ -31,9 +32,12 @@ contract BalancerV2Executor is IExecutor {
             IERC20 tokenOut,
             bytes32 poolId,
             address receiver,
-            bool approvalNeeded
+            bool approvalNeeded,
+            TransferType transferType
         ) = _decodeData(data);
 
+        _transfer(address(this), transferType, tokenIn, givenAmount);
+
         if (approvalNeeded) {
             // slither-disable-next-line unused-return
             tokenIn.forceApprove(VAULT, type(uint256).max);

foundry/src/executors/CurveExecutor.sol

@@ -4,8 +4,9 @@ pragma solidity ^0.8.26;
 import "@interfaces/IExecutor.sol";
 import "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
 import "@openzeppelin/contracts/utils/Address.sol";
+import {RestrictTransferFrom} from "../RestrictTransferFrom.sol";
 
-error CurveExecutor__AddressZero();
+    error CurveExecutor__AddressZero();
 error CurveExecutor__InvalidDataLength();
 
 interface CryptoPool {
@@ -34,12 +35,12 @@ interface CryptoPoolETH {
     // slither-disable-end naming-convention
 }
 
-contract CurveExecutor is IExecutor {
+contract CurveExecutor is IExecutor, RestrictTransferFrom {
     using SafeERC20 for IERC20;
 
     address public immutable nativeToken;
 
-    constructor(address _nativeToken, address _permit2) {
+    constructor(address _nativeToken, address _permit2) RestrictTransferFrom(_permit2) {
         if (_nativeToken == address(0)) {
             revert CurveExecutor__AddressZero();
         }
@@ -52,7 +53,7 @@ contract CurveExecutor is IExecutor {
         payable
         returns (uint256)
     {
-        if (data.length != 84) revert CurveExecutor__InvalidDataLength();
+        if (data.length != 85) revert CurveExecutor__InvalidDataLength();
 
         (
             address tokenIn,
@@ -62,6 +63,7 @@ contract CurveExecutor is IExecutor {
             int128 i,
             int128 j,
             bool approvalNeeded,
+            TransferType transferType,
             address receiver
         ) = _decodeData(data);
 
@@ -69,6 +71,7 @@ contract CurveExecutor is IExecutor {
             // slither-disable-next-line unused-return
             IERC20(tokenIn).forceApprove(address(pool), type(uint256).max);
         }
+        _transfer(address(this), transferType, tokenIn, amountIn);
 
         /// Inspired by Curve's router contract: https://github.com/curvefi/curve-router-ng/blob/9ab006ca848fc7f1995b6fbbecfecc1e0eb29e2a/contracts/Router.vy#L44
         uint256 balanceBefore = _balanceOf(tokenOut);
@@ -120,6 +123,7 @@ contract CurveExecutor is IExecutor {
             int128 i,
             int128 j,
             bool approvalNeeded,
+            TransferType transferType,
             address receiver
         )
     {
@@ -130,7 +134,8 @@ contract CurveExecutor is IExecutor {
         i = int128(uint128(uint8(data[61])));
         j = int128(uint128(uint8(data[62])));
         approvalNeeded = data[63] != 0;
-        receiver = address(bytes20(data[64:84]));
+        transferType = TransferType(uint8(data[64]));
+        receiver = address(bytes20(data[65:85]));
     }
 
     /**

foundry/src/executors/EkuboExecutor.sol

@@ -11,7 +11,7 @@ import {SafeTransferLib} from "@solady/utils/SafeTransferLib.sol";
 import {LibBytes} from "@solady/utils/LibBytes.sol";
 import {Config, EkuboPoolKey} from "@ekubo/types/poolKey.sol";
 import {MAX_SQRT_RATIO, MIN_SQRT_RATIO} from "@ekubo/types/sqrtRatio.sol";
-import "../OneTransferFromOnly.sol";
+import "../RestrictTransferFrom.sol";
 import "@openzeppelin/contracts/utils/Address.sol";
 
 contract EkuboExecutor is
@@ -19,7 +19,7 @@ contract EkuboExecutor is
     ILocker,
     IPayer,
     ICallback,
-    OneTransferFromOnly
+RestrictTransferFrom
 {
     error EkuboExecutor__InvalidDataLength();
     error EkuboExecutor__CoreOnly();
@@ -36,7 +36,7 @@ contract EkuboExecutor is
     using SafeERC20 for IERC20;
 
     constructor(address _core, address _permit2)
-        OneTransferFromOnly(_permit2)
+        RestrictTransferFrom(_permit2)
     {
         core = ICore(_core);
     }
@@ -46,7 +46,7 @@ contract EkuboExecutor is
         payable
         returns (uint256 calculatedAmount)
     {
-        if (data.length < 93) revert EkuboExecutor__InvalidDataLength();
+        if (data.length < 92) revert EkuboExecutor__InvalidDataLength();
 
         // amountIn must be at most type(int128).MAX
         calculatedAmount =
@@ -125,10 +125,9 @@ contract EkuboExecutor is
     function _locked(bytes calldata swapData) internal returns (int128) {
         int128 nextAmountIn = int128(uint128(bytes16(swapData[0:16])));
         uint128 tokenInDebtAmount = uint128(nextAmountIn);
-        bool transferFromNeeded = (swapData[16] != 0);
+        TransferType transferType = TransferType(uint8(swapData[16]));
-        bool transferNeeded = (swapData[17] != 0);
+        address receiver = address(bytes20(swapData[17:37]));
-        address receiver = address(bytes20(swapData[18:38]));
+        address tokenIn = address(bytes20(swapData[37:57]));
-        address tokenIn = address(bytes20(swapData[38:58]));
 
         address nextTokenIn = tokenIn;
 
@@ -162,7 +161,7 @@ contract EkuboExecutor is
             offset += HOP_BYTE_LEN;
         }
 
-        _pay(tokenIn, tokenInDebtAmount, transferFromNeeded, transferNeeded);
+        _pay(tokenIn, tokenInDebtAmount, transferType);
         core.withdraw(nextTokenIn, receiver, uint128(nextAmountIn));
         return nextAmountIn;
     }
@@ -170,8 +169,7 @@ contract EkuboExecutor is
     function _pay(
         address token,
         uint128 amount,
-        bool transferFromNeeded,
+        TransferType transferType
-        bool transferNeeded
     ) internal {
         address target = address(core);
 
@@ -185,11 +183,10 @@ contract EkuboExecutor is
                 mstore(free, shl(224, 0x0c11dedd))
                 mstore(add(free, 4), token)
                 mstore(add(free, 36), shl(128, amount))
-                mstore(add(free, 52), shl(248, transferFromNeeded))
+                mstore(add(free, 52), shl(248, transferType))
-                mstore(add(free, 53), shl(248, transferNeeded))
 
-                // 4 (selector) + 32 (token) + 16 (amount) + 1 (transferFromNeeded) + 1 (transferNeeded) = 54
+                // 4 (selector) + 32 (token) + 16 (amount) + 1 (transferType) = 53
-                if iszero(call(gas(), target, 0, free, 54, 0, 0)) {
+                if iszero(call(gas(), target, 0, free, 53, 0, 0)) {
                     returndatacopy(0, 0, returndatasize())
                     revert(0, returndatasize())
                 }
@@ -200,13 +197,8 @@ contract EkuboExecutor is
     function _payCallback(bytes calldata payData) internal {
         address token = address(bytes20(payData[12:32])); // This arg is abi-encoded
         uint128 amount = uint128(bytes16(payData[32:48]));
-        bool transferFromNeeded = (payData[48] != 0);
+        TransferType transferType = TransferType(uint8(payData[48]));
-        bool transferNeeded = (payData[49] != 0);
+        _transfer(core, transferType, token, amount);
-        if (transferFromNeeded) {
-            _transfer(msg.sender);
-        } else if (transferNeeded) {
-            IERC20(token).safeTransfer(msg.sender, amount);
-        }
     }
 
     // To receive withdrawals from Core

foundry/src/executors/MaverickV2Executor.sol

@@ -4,17 +4,18 @@ pragma solidity ^0.8.26;
 import "@interfaces/IExecutor.sol";
 import "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
 import "@openzeppelin/contracts/utils/Address.sol";
+import {RestrictTransferFrom} from "../RestrictTransferFrom.sol";
 
-error MaverickV2Executor__InvalidDataLength();
+    error MaverickV2Executor__InvalidDataLength();
 error MaverickV2Executor__InvalidTarget();
 error MaverickV2Executor__InvalidFactory();
 
-contract MaverickV2Executor is IExecutor {
+contract MaverickV2Executor is IExecutor, RestrictTransferFrom {
     using SafeERC20 for IERC20;
 
     address public immutable factory;
 
-    constructor(address _factory, address _permit2) {
+    constructor(address _factory, address _permit2) RestrictTransferFrom(_permit2) {
         if (_factory == address(0)) {
             revert MaverickV2Executor__InvalidFactory();
         }
@@ -30,9 +31,9 @@ contract MaverickV2Executor is IExecutor {
         address target;
         address receiver;
         IERC20 tokenIn;
-        bool transferNeeded;
+        TransferType transferType;
 
-        (tokenIn, target, receiver, transferNeeded) = _decodeData(data);
+        (tokenIn, target, receiver, transferType) = _decodeData(data);
 
         _verifyPairAddress(target);
         IMaverickV2Pool pool = IMaverickV2Pool(target);
@@ -47,14 +48,7 @@ contract MaverickV2Executor is IExecutor {
             tickLimit: tickLimit
         });
 
-        if (transferNeeded) {
+        _transfer(target, transferType, tokenIn, givenAmount);
-            if (address(tokenIn) == address(0)) {
-                Address.sendValue(payable(target), givenAmount);
-            } else {
-                // slither-disable-next-line arbitrary-send-erc20
-                tokenIn.safeTransfer(target, givenAmount);
-            }
-        }
 
         // slither-disable-next-line unused-return
         (, calculatedAmount) = pool.swap(receiver, swapParams, "");
@@ -67,7 +61,7 @@ contract MaverickV2Executor is IExecutor {
             IERC20 inToken,
             address target,
             address receiver,
-            bool transferNeeded
+            TransferType transferType
         )
     {
         if (data.length != 61) {
@@ -76,7 +70,7 @@ contract MaverickV2Executor is IExecutor {
         inToken = IERC20(address(bytes20(data[0:20])));
         target = address(bytes20(data[20:40]));
         receiver = address(bytes20(data[40:60]));
-        transferNeeded = (data[60] != 0);
+        transferType = TransferType(uint8(data[60]));
     }
 
     function _verifyPairAddress(address target) internal view {

foundry/src/executors/UniswapV2Executor.sol

@@ -4,14 +4,15 @@ pragma solidity ^0.8.26;
 import "@interfaces/IExecutor.sol";
 import "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
 import "@uniswap-v2/contracts/interfaces/IUniswapV2Pair.sol";
+import {RestrictTransferFrom} from "../RestrictTransferFrom.sol";
 
-error UniswapV2Executor__InvalidDataLength();
+    error UniswapV2Executor__InvalidDataLength();
 error UniswapV2Executor__InvalidTarget();
 error UniswapV2Executor__InvalidFactory();
 error UniswapV2Executor__InvalidInitCode();
 error UniswapV2Executor__InvalidFee();
 
-contract UniswapV2Executor is IExecutor {
+contract UniswapV2Executor is IExecutor, RestrictTransferFrom {
     using SafeERC20 for IERC20;
 
     address public immutable factory;
@@ -24,7 +25,7 @@ contract UniswapV2Executor is IExecutor {
         bytes32 _initCode,
         address _permit2,
         uint256 _feeBps
-    ) {
+    ) RestrictTransferFrom(_permit2) {
         if (_factory == address(0)) {
             revert UniswapV2Executor__InvalidFactory();
         }
@@ -50,19 +51,16 @@ contract UniswapV2Executor is IExecutor {
         address target;
         address receiver;
         bool zeroForOne;
-        bool transferNeeded;
+        TransferType transferType;
 
-        (tokenIn, target, receiver, zeroForOne, transferNeeded) =
+        (tokenIn, target, receiver, zeroForOne, transferType) =
             _decodeData(data);
 
         _verifyPairAddress(target);
 
         calculatedAmount = _getAmountOut(target, givenAmount, zeroForOne);
 
-        if (transferNeeded) {
+        _transfer(target, transferType, address(tokenIn), givenAmount);
-            // slither-disable-next-line arbitrary-send-erc20
-            tokenIn.safeTransfer(target, givenAmount);
-        }
 
         IUniswapV2Pair pool = IUniswapV2Pair(target);
         if (zeroForOne) {
@@ -80,7 +78,7 @@ contract UniswapV2Executor is IExecutor {
             address target,
             address receiver,
             bool zeroForOne,
-            bool transferNeeded
+            TransferType transferType,
         )
     {
         if (data.length != 62) {
@@ -90,7 +88,7 @@ contract UniswapV2Executor is IExecutor {
         target = address(bytes20(data[20:40]));
         receiver = address(bytes20(data[40:60]));
         zeroForOne = data[60] != 0;
-        transferNeeded = data[61] != 0;
+        transferType = TransferType(uint8(data[61]));
     }
 
     function _getAmountOut(address target, uint256 amountIn, bool zeroForOne)

foundry/src/executors/UniswapV3Executor.sol

@@ -5,7 +5,7 @@ import "@interfaces/IExecutor.sol";
 import "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
 import "@uniswap/v3-core/contracts/interfaces/IUniswapV3Pool.sol";
 import "@interfaces/ICallback.sol";
-import {OneTransferFromOnly} from "../OneTransferFromOnly.sol";
+import {RestrictTransferFrom} from "../RestrictTransferFrom.sol";
 import "@openzeppelin/contracts/utils/Address.sol";
 
 error UniswapV3Executor__InvalidDataLength();
@@ -13,7 +13,7 @@ error UniswapV3Executor__InvalidFactory();
 error UniswapV3Executor__InvalidTarget();
 error UniswapV3Executor__InvalidInitCode();
 
-contract UniswapV3Executor is IExecutor, ICallback, OneTransferFromOnly {
+contract UniswapV3Executor is IExecutor, ICallback, RestrictTransferFrom {
     using SafeERC20 for IERC20;
 
     uint160 private constant MIN_SQRT_RATIO = 4295128739;
@@ -25,7 +25,7 @@ contract UniswapV3Executor is IExecutor, ICallback, OneTransferFromOnly {
     address private immutable self;
 
     constructor(address _factory, bytes32 _initCode, address _permit2)
-        OneTransferFromOnly(_permit2)
+        RestrictTransferFrom(_permit2)
     {
         if (_factory == address(0)) {
             revert UniswapV3Executor__InvalidFactory();
@@ -51,8 +51,7 @@ contract UniswapV3Executor is IExecutor, ICallback, OneTransferFromOnly {
             address receiver,
             address target,
             bool zeroForOne,
-            bool transferFromNeeded,
+            uint8 transferType
-            bool transferNeeded
         ) = _decodeData(data);
 
         _verifyPairAddress(tokenIn, tokenOut, fee, target);
@@ -62,7 +61,7 @@ contract UniswapV3Executor is IExecutor, ICallback, OneTransferFromOnly {
         IUniswapV3Pool pool = IUniswapV3Pool(target);
 
         bytes memory callbackData = _makeV3CallbackData(
-            tokenIn, tokenOut, fee, transferFromNeeded, transferNeeded
+            tokenIn, tokenOut, fee, transferType
         );
 
         {
@@ -99,24 +98,14 @@ contract UniswapV3Executor is IExecutor, ICallback, OneTransferFromOnly {
             abi.decode(msgData[4:68], (int256, int256));
 
         address tokenIn = address(bytes20(msgData[132:152]));
-
+        bool transferType = TransferType(uint8(msgData[175]));
-        bool transferFromNeeded = msgData[175] != 0;
-        bool transferNeeded = msgData[176] != 0;
 
         verifyCallback(msgData[132:]);
 
         uint256 amountOwed =
             amount0Delta > 0 ? uint256(amount0Delta) : uint256(amount1Delta);
 
-        if (transferFromNeeded) {
+        _transfer(msg.sender, transferType, tokenIn, amountOwed);
-            _transfer(msg.sender);
-        } else if (transferNeeded) {
-            if (tokenIn == address(0)) {
-                Address.sendValue(payable(msg.sender), amountOwed);
-            } else {
-                IERC20(tokenIn).safeTransfer(msg.sender, amountOwed);
-            }
-        }
 
         return abi.encode(amountOwed, tokenIn);
     }
@@ -147,11 +136,10 @@ contract UniswapV3Executor is IExecutor, ICallback, OneTransferFromOnly {
             address receiver,
             address target,
             bool zeroForOne,
-            bool transferFromNeeded,
+            uint8 transferType
-            bool transferNeeded
         )
     {
-        if (data.length != 86) {
+        if (data.length != 85) {
             revert UniswapV3Executor__InvalidDataLength();
         }
         tokenIn = address(bytes20(data[0:20]));
@@ -160,19 +148,17 @@ contract UniswapV3Executor is IExecutor, ICallback, OneTransferFromOnly {
         receiver = address(bytes20(data[43:63]));
         target = address(bytes20(data[63:83]));
         zeroForOne = uint8(data[83]) > 0;
-        transferFromNeeded = data[84] != 0;
+        transferType = uint8(data[84]);
-        transferNeeded = data[85] != 0;
     }
 
     function _makeV3CallbackData(
         address tokenIn,
         address tokenOut,
         uint24 fee,
-        bool transferFromNeeded,
+        uint8 transferType
-        bool transferNeeded
     ) internal pure returns (bytes memory) {
         return abi.encodePacked(
-            tokenIn, tokenOut, fee, transferFromNeeded, transferNeeded
+            tokenIn, tokenOut, fee, transferType
         );
     }
 

foundry/src/executors/UniswapV4Executor.sol

@@ -22,7 +22,7 @@ import {IUnlockCallback} from
 import {SafeCast} from "@uniswap/v4-core/src/libraries/SafeCast.sol";
 import {TransientStateLibrary} from
     "@uniswap/v4-core/src/libraries/TransientStateLibrary.sol";
-import "../OneTransferFromOnly.sol";
+import "../RestrictTransferFrom.sol";
 import "@openzeppelin/contracts/utils/Address.sol";
 
 error UniswapV4Executor__InvalidDataLength();
@@ -38,7 +38,7 @@ contract UniswapV4Executor is
     IExecutor,
     IUnlockCallback,
     ICallback,
-    OneTransferFromOnly
+RestrictTransferFrom
 {
     using SafeERC20 for IERC20;
     using CurrencyLibrary for Currency;
@@ -58,7 +58,7 @@ contract UniswapV4Executor is
     }
 
     constructor(IPoolManager _poolManager, address _permit2)
-        OneTransferFromOnly(_permit2)
+        RestrictTransferFrom(_permit2)
     {
         poolManager = _poolManager;
         _self = address(this);
@@ -83,8 +83,7 @@ contract UniswapV4Executor is
             address tokenIn,
             address tokenOut,
             bool zeroForOne,
-            bool transferFromNeeded,
+            TransferType transferType,
-            bool transferNeeded,
             address receiver,
             UniswapV4Executor.UniswapV4Pool[] memory pools
         ) = _decodeData(data);
@@ -102,8 +101,7 @@ contract UniswapV4Executor is
                 key,
                 zeroForOne,
                 amountIn,
-                transferFromNeeded,
+                transferType,
-                transferNeeded,
                 receiver,
                 bytes("")
             );
@@ -125,8 +123,7 @@ contract UniswapV4Executor is
                 currencyIn,
                 path,
                 amountIn,
-                transferFromNeeded,
+                transferType,
-                transferNeeded,
                 receiver
             );
         }
@@ -144,26 +141,24 @@ contract UniswapV4Executor is
             address tokenIn,
             address tokenOut,
             bool zeroForOne,
-            bool transferFromNeeded,
+            TransferType transferType,
-            bool transferNeeded,
             address receiver,
             UniswapV4Pool[] memory pools
         )
     {
-        if (data.length < 89) {
+        if (data.length < 88) {
             revert UniswapV4Executor__InvalidDataLength();
         }
 
         tokenIn = address(bytes20(data[0:20]));
         tokenOut = address(bytes20(data[20:40]));
         zeroForOne = data[40] != 0;
-        transferFromNeeded = data[41] != 0;
+        transferType = TransferType(uint8(data[41]));
-        transferNeeded = data[42] != 0;
+        receiver = address(bytes20(data[42:62]));
-        receiver = address(bytes20(data[43:63]));
 
-        uint256 poolsLength = (data.length - 63) / 26; // 26 bytes per pool object
+        uint256 poolsLength = (data.length - 62) / 26; // 26 bytes per pool object
         pools = new UniswapV4Pool[](poolsLength);
-        bytes memory poolsData = data[63:];
+        bytes memory poolsData = data[62:];
         uint256 offset = 0;
         for (uint256 i = 0; i < poolsLength; i++) {
             address intermediaryToken;
@@ -243,8 +238,7 @@ contract UniswapV4Executor is
      * @param poolKey The key of the pool to swap in.
      * @param zeroForOne Whether the swap is from token0 to token1 (true) or vice versa (false).
      * @param amountIn The amount of tokens to swap in.
-     * @param transferFromNeeded Whether to transferFrom input tokens into the core contract from the swapper's wallet .
+     * @param transferType The type of action necessary to pay back the pool.
-     * @param transferNeeded Whether to transfer input tokens into the core contract from the router contract
      * @param receiver The address of the receiver.
      * @param hookData Additional data for hook contracts.
      */
@@ -252,8 +246,7 @@ contract UniswapV4Executor is
         PoolKey memory poolKey,
         bool zeroForOne,
         uint128 amountIn,
-        bool transferFromNeeded,
+        TransferType transferType,
-        bool transferNeeded,
         address receiver,
         bytes calldata hookData
     ) external returns (uint128) {
@@ -266,7 +259,7 @@ contract UniswapV4Executor is
         if (amount > amountIn) {
             revert UniswapV4Executor__V4TooMuchRequested(amountIn, amount);
         }
-        _settle(currencyIn, amount, transferFromNeeded, transferNeeded);
+        _settle(currencyIn, amount, transferType);
 
         Currency currencyOut =
             zeroForOne ? poolKey.currency1 : poolKey.currency0;
@@ -279,16 +272,14 @@ contract UniswapV4Executor is
      * @param currencyIn The currency of the input token.
      * @param path The path to swap along.
      * @param amountIn The amount of tokens to swap in.
-     * @param transferFromNeeded Whether to transferFrom input tokens into the core contract from the swapper's wallet .
+     * @param transferType The type of action necessary to pay back the pool.
-     * @param transferNeeded Whether to transfer input tokens into the core contract from the router contract
      * @param receiver The address of the receiver.
      */
     function swapExactInput(
         Currency currencyIn,
         PathKey[] calldata path,
         uint128 amountIn,
-        bool transferFromNeeded,
+        TransferType transferType,
-        bool transferNeeded,
         address receiver
     ) external returns (uint128) {
         uint128 amountOut = 0;
@@ -319,7 +310,7 @@ contract UniswapV4Executor is
         if (amount > amountIn) {
             revert UniswapV4Executor__V4TooMuchRequested(amountIn, amount);
         }
-        _settle(currencyIn, amount, transferFromNeeded, transferNeeded);
+        _settle(currencyIn, amount, transferType);
 
         _take(
             swapCurrencyIn, // at the end of the loop this is actually currency out
@@ -391,17 +382,13 @@ contract UniswapV4Executor is
      * @dev The implementing contract must ensure that the `payer` is a secure address.
      * @param currency The currency to settle.
      * @param amount The amount to send.
-     * @param transferFromNeeded Whether to manually transferFrom input tokens into the
+     * @param transferType The type of action necessary to pay back the pool.
-     *  core contract from the swapper.
-     * @param transferNeeded Whether to manually transfer input tokens into the
-     *  core contract from the router.
      * @dev Returns early if the amount is 0.
      */
     function _settle(
         Currency currency,
         uint256 amount,
-        bool transferFromNeeded,
+        TransferType transferType
-        bool transferNeeded
     ) internal {
         if (amount == 0) return;
         poolManager.sync(currency);
@@ -409,18 +396,12 @@ contract UniswapV4Executor is
             // slither-disable-next-line unused-return
             poolManager.settle{value: amount}();
         } else {
-            if (transferFromNeeded) {
+            _transfer(
-                // transferFrom swapper's wallet into the core contract
+                address(poolManager),
-                _transfer(address(poolManager));
+                transferType,
-            } else if (transferNeeded) {
+                address(currency),
-                address tokenIn = Currency.unwrap(currency);
+                amount
-                // transfer from router contract into the core contract
+            );
-                if (tokenIn == address(0)) {
-                    Address.sendValue(payable(address(poolManager)), amount);
-                } else {
-                    IERC20(tokenIn).safeTransfer(address(poolManager), amount);
-                }
-            }
             // slither-disable-next-line unused-return
             poolManager.settle();
         }