Remove balanceOf checks in EkuboExecutor & remove ICallback impl

foundry/src/TychoRouter.sol

@@ -557,9 +557,7 @@ contract TychoRouter is AccessControl, Dispatcher, Pausable, ReentrancyGuard {
         address executor = address(0xA612f60d3C49E5f13f0e067b14E0eD6656F3F279);
 
         // slither-disable-next-line controlled-delegatecall,low-level-calls
-        (bool success, bytes memory result) = executor.delegatecall(
+        (bool success, bytes memory result) = executor.delegatecall(msg.data);
-            abi.encodeWithSelector(ICallback.handleCallback.selector, msg.data)
-        );
 
         if (!success) {
             revert(
@@ -570,15 +568,19 @@ contract TychoRouter is AccessControl, Dispatcher, Pausable, ReentrancyGuard {
                 )
             );
         }
+
+        // slither-disable-next-line assembly
+        assembly ("memory-safe") {
+            // Propagate the swappedAmount
+            return(add(result, 32), 16)
+        }
     }
 
     function payCallback(uint256, address /*token*/ ) external {
         address executor = address(0xA612f60d3C49E5f13f0e067b14E0eD6656F3F279);
 
         // slither-disable-next-line controlled-delegatecall,low-level-calls
-        (bool success, bytes memory result) = executor.delegatecall(
+        (bool success, bytes memory result) = executor.delegatecall(msg.data);
-            abi.encodeWithSelector(ICallback.handleCallback.selector, msg.data)
-        );
 
         if (!success) {
             revert(

foundry/src/executors/EkuboExecutor.sol

@@ -3,7 +3,6 @@ pragma solidity ^0.8.26;
 
 import {IERC20} from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
 import {IExecutor} from "@interfaces/IExecutor.sol";
-import {ICallback} from "@interfaces/ICallback.sol";
 import {ICore} from "@ekubo/interfaces/ICore.sol";
 import {ILocker, IPayer} from "@ekubo/interfaces/IFlashAccountant.sol";
 import {NATIVE_TOKEN_ADDRESS} from "@ekubo/math/constants.sol";
@@ -12,17 +11,14 @@ 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";
 
-contract EkuboExecutor is IExecutor, ICallback, ILocker, IPayer {
+contract EkuboExecutor is IExecutor, ILocker, IPayer {
     error EkuboExecutor__InvalidDataLength();
     error EkuboExecutor__CoreOnly();
     error EkuboExecutor__UnknownCallback();
 
     ICore immutable core;
 
-    bytes4 constant LOCKED_SELECTOR = 0xb45a3c0e; // locked(uint256)
+    uint256 constant POOL_DATA_OFFSET = 92;
-    bytes4 constant PAY_CALLBACK_SELECTOR = 0x599d0714; // payCallback(uint256,address)
-
-    uint256 constant POOL_DATA_OFFSET = 56;
     uint256 constant HOP_BYTE_LEN = 52;
 
     constructor(address _core) {
@@ -36,68 +32,67 @@ contract EkuboExecutor is IExecutor, ICallback, ILocker, IPayer {
     {
         if (data.length < 92) revert EkuboExecutor__InvalidDataLength();
 
-        uint256 tokenOutOffset = data.length - HOP_BYTE_LEN;
-        address tokenOut =
-            address(bytes20(LibBytes.loadCalldata(data, tokenOutOffset)));
-
-        uint256 tokenOutBalanceBefore = _balanceOf(tokenOut);
-
         // amountIn must be at most type(int128).MAX
-        _lock(bytes.concat(bytes16(uint128(amountIn)), data));
+        calculatedAmount = uint256(_lock(bytes.concat(bytes16(uint128(amountIn)), data)));
-
-        uint256 tokenOutBalanceAfter = _balanceOf(tokenOut);
-
-        // It would be better if we could somehow pass back the swapped amount from the lock but the interface doesn't offer that capability.
-        // Note that the current approach also prevents arbs that return less than their input because of arithmetic underflow.
-        calculatedAmount = tokenOutBalanceAfter - tokenOutBalanceBefore;
     }
 
-    // We can't use the return value here since it won't get propagated (see Dispatcher.sol:_handleCallback)
-    function handleCallback(bytes calldata raw)
-        external
-        returns (bytes memory)
-    {
-        verifyCallback(raw);
-
-        // Without selector and locker id
-        bytes calldata stripped = raw[36:];
-
-        bytes4 selector = bytes4(raw[:4]);
-
-        if (selector == LOCKED_SELECTOR) {
-            _locked(stripped);
-        } else if (selector == PAY_CALLBACK_SELECTOR) {
-            _payCallback(stripped);
-        } else {
-            revert EkuboExecutor__UnknownCallback();
-        }
-
-        return "";
-    }
-
-    function verifyCallback(bytes calldata) public view coreOnly {}
-
     function locked(uint256) external coreOnly {
-        // Without selector and locker id
+        int128 nextAmountIn = int128(uint128(bytes16(msg.data[36:52])));
-        _locked(msg.data[36:]);
+        uint128 tokenInDebtAmount = uint128(nextAmountIn);
+
+        address receiver = address(bytes20(msg.data[52:72]));
+        address tokenIn = address(bytes20(msg.data[72:POOL_DATA_OFFSET]));
+
+        address nextTokenIn = tokenIn;
+
+        uint256 hopsLength = (msg.data.length - POOL_DATA_OFFSET) / HOP_BYTE_LEN;
+
+        uint256 offset = POOL_DATA_OFFSET;
+
+        for (uint256 i = 0; i < hopsLength; i++) {
+            address nextTokenOut =
+                address(bytes20(LibBytes.loadCalldata(msg.data, offset)));
+            Config poolConfig =
+                Config.wrap(LibBytes.loadCalldata(msg.data, offset + 20));
+
+            (address token0, address token1, bool isToken1) = nextTokenIn
+                > nextTokenOut
+                ? (nextTokenOut, nextTokenIn, true)
+                : (nextTokenIn, nextTokenOut, false);
+
+            // slither-disable-next-line calls-loop
+            (int128 delta0, int128 delta1) = core.swap_611415377(
+                EkuboPoolKey(token0, token1, poolConfig),
+                nextAmountIn,
+                isToken1,
+                isToken1 ? MAX_SQRT_RATIO : MIN_SQRT_RATIO,
+                0
+            );
+
+            nextTokenIn = nextTokenOut;
+            nextAmountIn = -(isToken1 ? delta0 : delta1);
+
+            offset += HOP_BYTE_LEN;
+        }
+
+        _pay(tokenIn, tokenInDebtAmount);
+
+        core.withdraw(nextTokenIn, receiver, uint128(nextAmountIn));
+
+        // slither-disable-next-line assembly
+        assembly ("memory-safe") {
+            mstore(0, nextAmountIn)
+            return(0x10, 16)
+        }
     }
 
-    function payCallback(uint256, address /*token*/ ) external coreOnly {
+    function payCallback(uint256, address token) external coreOnly {
-        // Without selector and locker id
+        uint128 amount = uint128(bytes16(msg.data[68:84]));
-        _payCallback(msg.data[36:]);
+
+        SafeTransferLib.safeTransfer(token, address(core), amount);
     }
 
-    function _balanceOf(address token)
+    function _lock(bytes memory data) internal returns (uint128 swappedAmount) {
-        internal
-        view
-        returns (uint256 balance)
-    {
-        balance = token == NATIVE_TOKEN_ADDRESS
-            ? address(this).balance
-            : IERC20(token).balanceOf(address(this));
-    }
-
-    function _lock(bytes memory data) internal {
         address target = address(core);
 
         // slither-disable-next-line assembly
@@ -116,52 +111,12 @@ contract EkuboExecutor is IExecutor, ICallback, ILocker, IPayer {
                 returndatacopy(0, 0, returndatasize())
                 revert(0, returndatasize())
             }
+
+            returndatacopy(0, 0, 16)
+            swappedAmount := shr(128, mload(0))
         }
     }
 
-    function _locked(bytes calldata swapData) internal {
-        int128 nextAmountIn = int128(uint128(bytes16(swapData[0:16])));
-        uint128 tokenInDebtAmount = uint128(nextAmountIn);
-
-        address receiver = address(bytes20(swapData[16:36]));
-        address tokenIn = address(bytes20(swapData[36:POOL_DATA_OFFSET]));
-
-        address nextTokenIn = tokenIn;
-
-        uint256 hopsLength = (swapData.length - POOL_DATA_OFFSET) / HOP_BYTE_LEN;
-
-        uint256 offset = POOL_DATA_OFFSET;
-
-        for (uint256 i = 0; i < hopsLength; i++) {
-            address nextTokenOut =
-                address(bytes20(LibBytes.loadCalldata(swapData, offset)));
-            Config poolConfig =
-                Config.wrap(LibBytes.loadCalldata(swapData, offset + 20));
-
-            (address token0, address token1, bool isToken1) = nextTokenIn
-                > nextTokenOut
-                ? (nextTokenOut, nextTokenIn, true)
-                : (nextTokenIn, nextTokenOut, false);
-
-            (int128 delta0, int128 delta1) = core.swap_611415377(
-                EkuboPoolKey(token0, token1, poolConfig),
-                nextAmountIn,
-                isToken1,
-                isToken1 ? MAX_SQRT_RATIO : MIN_SQRT_RATIO,
-                0
-            );
-
-            nextTokenIn = nextTokenOut;
-            nextAmountIn = -(isToken1 ? delta0 : delta1);
-
-            offset += HOP_BYTE_LEN;
-        }
-
-        _pay(tokenIn, tokenInDebtAmount);
-
-        core.withdraw(nextTokenIn, receiver, uint128(nextAmountIn));
-    }
-
     function _pay(address token, uint128 amount) internal {
         address target = address(core);
 
@@ -185,13 +140,6 @@ contract EkuboExecutor is IExecutor, ICallback, ILocker, IPayer {
         }
     }
 
-    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]));
-
-        SafeTransferLib.safeTransfer(token, address(core), amount);
-    }
-
     // To receive withdrawals from Core
     receive() external payable {}