feat: add Treble dapp allocateValue logic

Author: BenSparksCode

src/contracts/examples/trebleswap/TrebleSwapDAppControl.sol

@@ -9,10 +9,6 @@ import { CallConfig } from "src/contracts/types/ConfigTypes.sol";
 import "src/contracts/types/UserOperation.sol";
 import "src/contracts/types/SolverOperation.sol";
 
-// NOTES:
-// Support for native in/out?
-// Support for bid token (TREB) in/out?
-
 // ODOS v2 Router on Base: https://basescan.org/address/0x19ceead7105607cd444f5ad10dd51356436095a1
 // Main function: swapCompact()
 
@@ -26,13 +22,16 @@ struct SwapTokenInfo {
 contract TrebleSwapDAppControl is DAppControl {
     address public constant ODOS_ROUTER = 0x19cEeAd7105607Cd444F5ad10dd51356436095a1;
 
-    // TODO WETH on Base for now, change to TREB when token deployed
-    address public constant TREB = address(0x4200000000000000000000000000000000000006);
+    // TODO TREB token not available yet - replace when it is. DEGEN address for now.
+    address public constant TREB = 0x4ed4E862860beD51a9570b96d89aF5E1B0Efefed;
     address internal constant _ETH = address(0);
+    address internal constant _BURN = address(0xdead);
 
     error InvalidUserOpData();
     error UserOpDappNotOdosRouter();
     error InsufficientUserOpValue();
+    error InsufficientTrebBalance();
+    error InsufficientOutputBalance();
 
     constructor(
         address atlas
@@ -73,30 +72,56 @@ contract TrebleSwapDAppControl is DAppControl {
     function _preOpsCall(UserOperation calldata userOp) internal virtual override returns (bytes memory) {
         if (userOp.dapp != ODOS_ROUTER) revert UserOpDappNotOdosRouter();
 
-        (bool success, bytes memory data) =
+        (bool success, bytes memory swapData) =
             CONTROL.staticcall(abi.encodePacked(this.decodeUserOpData.selector, userOp.data));
 
         if (!success) revert InvalidUserOpData();
 
-        SwapTokenInfo memory swapTokenInfo = abi.decode(data, (SwapTokenInfo));
+        SwapTokenInfo memory _swapInfo = abi.decode(swapData, (SwapTokenInfo));
 
         // If inputToken is ERC20, transfer tokens from user to EE, and approve Odos router for swap
-        if (swapTokenInfo.inputToken != _ETH) {
-            _transferUserERC20(swapTokenInfo.inputToken, address(this), swapTokenInfo.inputAmount);
-            SafeTransferLib.safeApprove(swapTokenInfo.inputToken, ODOS_ROUTER, swapTokenInfo.inputAmount);
+        if (_swapInfo.inputToken != _ETH) {
+            _transferUserERC20(_swapInfo.inputToken, address(this), _swapInfo.inputAmount);
+            SafeTransferLib.safeApprove(_swapInfo.inputToken, ODOS_ROUTER, _swapInfo.inputAmount);
         } else {
-            if (userOp.value < swapTokenInfo.inputAmount) revert InsufficientUserOpValue();
+            if (userOp.value < _swapInfo.inputAmount) revert InsufficientUserOpValue();
         }
 
-        return data; // return SwapTokenInfo in bytes format, to be used in allocateValue.
+        return swapData; // return SwapTokenInfo in bytes format, to be used in allocateValue.
     }
 
-    // UserOperation happens here. EE calls router (userOp.dapp) with userOp.data as calldata.
+    function _allocateValueCall(address, uint256 bidAmount, bytes calldata data) internal virtual override {
+        SwapTokenInfo memory _swapInfo = abi.decode(data, (SwapTokenInfo));
+        uint256 _outputTokenBalance = _balanceOf(_swapInfo.outputToken);
+
+        // Check enough output token was received after swap
+        if (_swapInfo.outputToken == TREB) {
+            // If outputToken is TREB, check there is enough to burn the solver bid and fulfill user swap
+            if (_outputTokenBalance < bidAmount + _swapInfo.outputMin) revert InsufficientTrebBalance();
+            _outputTokenBalance -= bidAmount; // deduct solver bid to get final user output amount to send
+        } else {
+            if (_outputTokenBalance < _swapInfo.outputMin) revert InsufficientOutputBalance();
+        }
+
+        // Transfer output token to user
+        if (_swapInfo.outputToken == _ETH) {
+            SafeTransferLib.safeTransferETH(_user(), _outputTokenBalance);
+        } else {
+            SafeTransferLib.safeTransfer(_swapInfo.outputToken, _user(), _outputTokenBalance);
+        }
+
+        // If solver won, burn TREB bid
+        if (bidAmount > 0) SafeTransferLib.safeTransfer(TREB, _BURN, bidAmount);
 
-    function _allocateValueCall(address bidToken, uint256 bidAmount, bytes calldata data) internal virtual override {
-        // Solver bid (in TREB) gets burnt here
-        // Send user back their tokenOut from the swap, and any leftover tokenIn
-        // Revert here if swap fails
+        // If any leftover input token, transfer back to user
+        uint256 _inputTokenBalance = _balanceOf(_swapInfo.inputToken);
+        if (_inputTokenBalance > 0) {
+            if (_swapInfo.inputToken == _ETH) {
+                SafeTransferLib.safeTransferETH(_user(), _inputTokenBalance);
+            } else {
+                SafeTransferLib.safeTransfer(_swapInfo.inputToken, _user(), _inputTokenBalance);
+            }
+        }
     }
 
     // ---------------------------------------------------- //
@@ -184,4 +209,12 @@ contract TrebleSwapDAppControl is DAppControl {
             }
         }
     }
+
+    function _balanceOf(address token) internal view returns (uint256) {
+        if (token == _ETH) {
+            return address(this).balance;
+        } else {
+            return SafeTransferLib.balanceOf(token, address(this));
+        }
+    }
 }

test/TrebleSwap.t.sol

@@ -48,6 +48,7 @@ contract TrebleSwapTest is BaseTest {
     IERC20 bWETH = IERC20(0x4200000000000000000000000000000000000006);
     IERC20 USDC = IERC20(0x833589fCD6eDb6E08f4c7C32D4f71b54bdA02913);
     IERC20 WUF = IERC20(0x4da78059D97f155E18B37765e2e042270f4E0fC4);
+    IERC20 TREB = IERC20(0x4ed4E862860beD51a9570b96d89aF5E1B0Efefed);
 
     address executionEnvironment;
 
@@ -90,8 +91,9 @@ contract TrebleSwapTest is BaseTest {
         bytes memory swapCompactCalldata =
             hex"83bd37f9000400014da78059d97f155e18b37765e2e042270f4e0fc4040bc108800601d1d9f50a5a028f5c0001f73f77f9466da712590ae432a80f07fd50a7de600001616535324976f8dbcef19df0705b95ace86ebb480001736F6980876FDa51A610AB79E2856528a62Bf80e0000000006020207003401000001020180000005020a0004040500000301010003060119ff0000000000000000000000000000000000000000000000000000000000000000616535324976f8dbcef19df0705b95ace86ebb48833589fcd6edb6e08f4c7c32d4f71b54bda02913569d81c17b5b4ac08929dc1769b8e39668d3ae29f6c0a374a483101e04ef5f7ac9bd15d9142bac95d9aaec86b65d86f6a7b5b1b0c42ffa531710b6ca42000000000000000000000000000000000000060000000000000000";
 
-        bytes memory encodedCall = abi.encodePacked(this.decodeSwapCompactCalldata.selector, swapCompactCalldata);
-        (bool res, bytes memory returnData) = address(this).staticcall(encodedCall);
+        bytes memory encodedCall =
+            abi.encodePacked(TrebleSwapDAppControl.decodeUserOpData.selector, swapCompactCalldata);
+        (bool res, bytes memory returnData) = address(trebleSwapControl).staticcall(encodedCall);
 
         console.log("res", res);
 
@@ -121,86 +123,6 @@ contract TrebleSwapTest is BaseTest {
             signature: new bytes(0)
         });
     }
-
-    // struct swapTokenInfo {
-    //     address inputToken;
-    //     uint256 inputAmount;
-    //     address outputToken;
-    //     uint256 outputMin;
-    // }
-
-    // TODO once this works, move to TrebleSwapDAppControl
-    function decodeSwapCompactCalldata() public view returns (SwapTokenInfo memory swapTokenInfo) {
-        assembly {
-            // helper function to get address either from storage or calldata
-            function getAddress(currPos) -> result, newPos {
-                let inputPos := shr(240, calldataload(currPos))
-
-                switch inputPos
-                // Reserve the null address as a special case that can be specified with 2 null bytes
-                case 0x0000 { newPos := add(currPos, 2) }
-                // This case means that the address is encoded in the calldata directly following the code
-                case 0x0001 {
-                    result := and(shr(80, calldataload(currPos)), 0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF)
-                    newPos := add(currPos, 22)
-                }
-                default {
-                    // 0000 and 0001 are reserved for cases above, so offset by 2 for addressList index
-                    let arg := sub(inputPos, 2)
-                    let selector := 0xb810fb43 // function selector for "addressList(uint256)"
-                    let ptr := mload(0x40) // get the free memory pointer
-                    mstore(ptr, shl(224, selector)) // shift selector to left of slot and store
-                    mstore(add(ptr, 4), arg) // store the uint256 argument after the selector
-
-                    // Perform the external call
-                    let success :=
-                        staticcall(
-                            gas(), // gas remaining
-                            ODOS_ROUTER,
-                            ptr, // input location
-                            0x24, // input size (4 byte selector + uint256 arg)
-                            ptr, // output location
-                            0x20 // output size (32 bytes for the address)
-                        )
-
-                    if eq(success, 0) { revert(0, 0) }
-
-                    result := mload(ptr)
-                    newPos := add(currPos, 2)
-                }
-            }
-
-            let result := 0
-            let pos := 8 // starts at 4 to skip the selector
-
-            // swapTokenInfo.inputToken (slot 0)
-            result, pos := getAddress(pos)
-            mstore(swapTokenInfo, result)
-
-            // swapTokenInfo.outputToken (slot 2)
-            result, pos := getAddress(pos)
-            mstore(add(swapTokenInfo, 0x40), result)
-
-            // swapTokenInfo.inputAmount (slot 1)
-            let inputAmountLength := shr(248, calldataload(pos))
-            pos := add(pos, 1)
-            if inputAmountLength {
-                mstore(add(swapTokenInfo, 0x20), shr(mul(sub(32, inputAmountLength), 8), calldataload(pos)))
-                pos := add(pos, inputAmountLength)
-            }
-
-            // swapTokenInfo.outputMin (slot 3)
-            // get outputQuote and slippageTolerance from calldata, then calculate outputMin
-            let quoteAmountLength := shr(248, calldataload(pos))
-            pos := add(pos, 1)
-            let outputQuote := shr(mul(sub(32, quoteAmountLength), 8), calldataload(pos))
-            pos := add(pos, quoteAmountLength)
-            {
-                let slippageTolerance := shr(232, calldataload(pos))
-                mstore(add(swapTokenInfo, 0x60), div(mul(outputQuote, sub(0xFFFFFF, slippageTolerance)), 0xFFFFFF))
-            }
-        }
-    }
 }
 
 interface IOdosRouterV2 {