From f14c8ee29ba67ca4cd8cfe8b00e8b907a0352f9a Mon Sep 17 00:00:00 2001
From: Diana Carvalho <diana@propellerheads.xyz>
Date: Fri, 9 May 2025 10:48:54 +0100
Subject: [PATCH] feat: Remove special handling of the Univ4 callback

The problem was that the pool manager was expecting an ABI encoded result to be returned and we were not returning that (we were returning just a result)
Special thanks to Max for figuring this out

Took 31 minutes
---
 foundry/src/Dispatcher.sol                  |  4 +++-
 foundry/src/TychoRouter.sol                 | 21 ++-------------------
 foundry/src/executors/UniswapV4Executor.sol |  7 +++++--
 3 files changed, 10 insertions(+), 22 deletions(-)

diff --git a/foundry/src/Dispatcher.sol b/foundry/src/Dispatcher.sol
index 1bb6209..dcbd1b4 100644
--- a/foundry/src/Dispatcher.sol
+++ b/foundry/src/Dispatcher.sol
@@ -127,7 +127,9 @@ contract Dispatcher {
             tstore(_CURRENTLY_SWAPPING_EXECUTOR_SLOT, 0)
         }
 
-        // this is necessary because the delegatecall will prepend extra bytes we don't want like the length and prefix
+        // The final callback result should not be ABI encoded. That is why we are decoding here.
+        // ABI encoding is very gas expensive and we want to avoid it if possible.
+        // The result from `handleCallback` is always ABI encoded.
         bytes memory decodedResult = abi.decode(result, (bytes));
         return decodedResult;
     }
diff --git a/foundry/src/TychoRouter.sol b/foundry/src/TychoRouter.sol
index 4d9d7b7..c237d64 100644
--- a/foundry/src/TychoRouter.sol
+++ b/foundry/src/TychoRouter.sol
@@ -657,13 +657,8 @@ contract TychoRouter is AccessControl, Dispatcher, Pausable, ReentrancyGuard {
     /**
      * @dev We use the fallback function to allow flexibility on callback.
      */
-    fallback() external {
-        bytes memory result = _callHandleCallbackOnExecutor(msg.data);
-        // slither-disable-next-line assembly
-        assembly ("memory-safe") {
-            // Propagate the result
-            return(add(result, 32), mload(result))
-        }
+    fallback(bytes calldata data) external returns (bytes memory) {
+        return _callHandleCallbackOnExecutor(data);
     }
 
     /**
@@ -778,18 +773,6 @@ contract TychoRouter is AccessControl, Dispatcher, Pausable, ReentrancyGuard {
         require(msg.sender.code.length != 0);
     }
 
-    /**
-     * @dev Called by UniswapV4 pool manager after achieving unlock state.
-     */
-    function unlockCallback(bytes calldata data)
-        external
-        returns (bytes memory)
-    {
-        if (data.length < 24) revert TychoRouter__InvalidDataLength();
-        bytes memory result = _callHandleCallbackOnExecutor(data);
-        return result;
-    }
-
     /**
      * @dev Gets balance of a token for a given address. Supports both native ETH and ERC20 tokens.
      */
diff --git a/foundry/src/executors/UniswapV4Executor.sol b/foundry/src/executors/UniswapV4Executor.sol
index f594adc..9b878ad 100644
--- a/foundry/src/executors/UniswapV4Executor.sol
+++ b/foundry/src/executors/UniswapV4Executor.sol
@@ -184,8 +184,11 @@ contract UniswapV4Executor is
         external
         returns (bytes memory)
     {
-        verifyCallback(data);
-        return _unlockCallback(data);
+        bytes calldata stripped = data[68:];
+        verifyCallback(stripped);
+        // Our general callback logic returns a not ABI encoded result.
+        // However, the pool manager expects the result to be ABI encoded. That is why we need to encode it here again.
+        return abi.encode(_unlockCallback(stripped));
     }
 
     function verifyCallback(bytes calldata) public view poolManagerOnly {}