Merge Swap and SwapToAnyToken into a single example

examples/swap/contracts/Swap.sol

@@ -9,30 +9,59 @@ import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
 import {RevertContext, RevertOptions} from "@zetachain/protocol-contracts/contracts/Revert.sol";
 import "@zetachain/protocol-contracts/contracts/zevm/interfaces/UniversalContract.sol";
 import "@zetachain/protocol-contracts/contracts/zevm/interfaces/IGatewayZEVM.sol";
+import "@zetachain/protocol-contracts/contracts/zevm/interfaces/IWZETA.sol";
 import {GatewayZEVM} from "@zetachain/protocol-contracts/contracts/zevm/GatewayZEVM.sol";
 
-contract Swap is UniversalContract {
-    address public immutable uniswapRouter;
+import {Initializable} from "@openzeppelin/contracts-upgradeable/proxy/utils/Initializable.sol";
+import {OwnableUpgradeable} from "@openzeppelin/contracts-upgradeable/access/OwnableUpgradeable.sol";
+import {UUPSUpgradeable} from "@openzeppelin/contracts-upgradeable/proxy/utils/UUPSUpgradeable.sol";
+
+contract Swap is
+    UniversalContract,
+    Initializable,
+    UUPSUpgradeable,
+    OwnableUpgradeable
+{
+    address public uniswapRouter;
     GatewayZEVM public gateway;
-    uint256 constant BITCOIN = 18332;
-    uint256 public immutable gasLimit;
+    uint256 constant BITCOIN = 8332;
+    uint256 constant BITCOIN_TESTNET = 18332;
+    uint256 public gasLimit;
 
     error InvalidAddress();
     error Unauthorized();
     error ApprovalFailed();
+    error TransferFailed();
+
+    event TokenSwap(
+        address sender,
+        bytes indexed recipient,
+        address indexed inputToken,
+        address indexed targetToken,
+        uint256 inputAmount,
+        uint256 outputAmount
+    );
 
     modifier onlyGateway() {
         if (msg.sender != address(gateway)) revert Unauthorized();
         _;
     }
 
-    constructor(
+    /// @custom:oz-upgrades-unsafe-allow constructor
+    constructor() {
+        _disableInitializers();
+    }
+
+    function initialize(
         address payable gatewayAddress,
         address uniswapRouterAddress,
-        uint256 gasLimitAmount
-    ) {
+        uint256 gasLimitAmount,
+        address owner
+    ) public initializer {
         if (gatewayAddress == address(0) || uniswapRouterAddress == address(0))
             revert InvalidAddress();
+        __UUPSUpgradeable_init();
+        __Ownable_init(owner);
         uniswapRouter = uniswapRouterAddress;
         gateway = GatewayZEVM(gatewayAddress);
         gasLimit = gasLimitAmount;
@@ -41,6 +70,7 @@ contract Swap is UniversalContract {
     struct Params {
         address target;
         bytes to;
+        bool withdraw;
     }
 
     function onCall(
@@ -49,29 +79,90 @@ contract Swap is UniversalContract {
         uint256 amount,
         bytes calldata message
     ) external onlyGateway {
-        Params memory params = Params({target: address(0), to: bytes("")});
-        if (context.chainID == BITCOIN) {
+        Params memory params = Params({
+            target: address(0),
+            to: bytes(""),
+            withdraw: true
+        });
+
+        if (context.chainID == BITCOIN_TESTNET || context.chainID == BITCOIN) {
             params.target = BytesHelperLib.bytesToAddress(message, 0);
             params.to = abi.encodePacked(
                 BytesHelperLib.bytesToAddress(message, 20)
             );
+            if (message.length >= 41) {
+                params.withdraw = BytesHelperLib.bytesToBool(message, 40);
+            }
         } else {
-            (address targetToken, bytes memory recipient) = abi.decode(
-                message,
-                (address, bytes)
-            );
+            (
+                address targetToken,
+                bytes memory recipient,
+                bool withdrawFlag
+            ) = abi.decode(message, (address, bytes, bool));
             params.target = targetToken;
             params.to = recipient;
+            params.withdraw = withdrawFlag;
         }
 
         (uint256 out, address gasZRC20, uint256 gasFee) = handleGasAndSwap(
             zrc20,
             amount,
             params.target
         );
+        emit TokenSwap(
+            context.sender,
+            params.to,
+            zrc20,
+            params.target,
+            amount,
+            out
+        );
         withdraw(params, context.sender, gasFee, gasZRC20, out, zrc20);
     }
 
+    function swap(
+        address inputToken,
+        uint256 amount,
+        address targetToken,
+        bytes memory recipient,
+        bool withdrawFlag
+    ) public {
+        bool success = IZRC20(inputToken).transferFrom(
+            msg.sender,
+            address(this),
+            amount
+        );
+        if (!success) {
+            revert TransferFailed();
+        }
+
+        (uint256 out, address gasZRC20, uint256 gasFee) = handleGasAndSwap(
+            inputToken,
+            amount,
+            targetToken
+        );
+        emit TokenSwap(
+            msg.sender,
+            recipient,
+            inputToken,
+            targetToken,
+            amount,
+            out
+        );
+        withdraw(
+            Params({
+                target: targetToken,
+                to: recipient,
+                withdraw: withdrawFlag
+            }),
+            msg.sender,
+            gasFee,
+            gasZRC20,
+            out,
+            inputToken
+        );
+    }
+
     function handleGasAndSwap(
         address inputToken,
         uint256 amount,
@@ -97,55 +188,58 @@ contract Swap is UniversalContract {
             swapAmount = amount - inputForGas;
         }
 
-        uint256 outputAmount = SwapHelperLib.swapExactTokensForTokens(
+        uint256 out = SwapHelperLib.swapExactTokensForTokens(
             uniswapRouter,
             inputToken,
             swapAmount,
             targetToken,
             0
         );
-        return (outputAmount, gasZRC20, gasFee);
+        return (out, gasZRC20, gasFee);
     }
 
     function withdraw(
         Params memory params,
         address sender,
         uint256 gasFee,
         address gasZRC20,
-        uint256 outputAmount,
+        uint256 out,
         address inputToken
     ) public {
-        if (gasZRC20 == params.target) {
-            if (
-                !IZRC20(gasZRC20).approve(
-                    address(gateway),
-                    outputAmount + gasFee
-                )
-            ) {
-                revert ApprovalFailed();
+        if (params.withdraw) {
+            if (gasZRC20 == params.target) {
+                if (!IZRC20(gasZRC20).approve(address(gateway), out + gasFee)) {
+                    revert ApprovalFailed();
+                }
+            } else {
+                if (!IZRC20(gasZRC20).approve(address(gateway), gasFee)) {
+                    revert ApprovalFailed();
+                }
+                if (!IZRC20(params.target).approve(address(gateway), out)) {
+                    revert ApprovalFailed();
+                }
             }
+            gateway.withdraw(
+                abi.encodePacked(params.to),
+                out,
+                params.target,
+                RevertOptions({
+                    revertAddress: address(this),
+                    callOnRevert: true,
+                    abortAddress: address(0),
+                    revertMessage: abi.encode(sender, inputToken),
+                    onRevertGasLimit: gasLimit
+                })
+            );
         } else {
-            if (!IZRC20(gasZRC20).approve(address(gateway), gasFee)) {
-                revert ApprovalFailed();
-            }
-            if (
-                !IZRC20(params.target).approve(address(gateway), outputAmount)
-            ) {
-                revert ApprovalFailed();
+            bool success = IWETH9(params.target).transfer(
+                address(uint160(bytes20(params.to))),
+                out
+            );
+            if (!success) {
+                revert TransferFailed();
             }
         }
-        gateway.withdraw(
-            params.to,
-            outputAmount,
-            params.target,
-            RevertOptions({
-                revertAddress: address(this),
-                callOnRevert: true,
-                abortAddress: address(0),
-                revertMessage: abi.encode(sender, inputToken),
-                onRevertGasLimit: gasLimit
-            })
-        );
     }
 
     function onRevert(RevertContext calldata context) external onlyGateway {
@@ -158,6 +252,7 @@ contract Swap is UniversalContract {
             context.amount,
             zrc20
         );
+
         gateway.withdraw(
             abi.encodePacked(sender),
             out,
@@ -175,4 +270,8 @@ contract Swap is UniversalContract {
     fallback() external payable {}
 
     receive() external payable {}
+
+    function _authorizeUpgrade(
+        address newImplementation
+    ) internal override onlyOwner {}
 }

examples/swap/contracts/SwapCompanion.sol

@@ -0,0 +1,55 @@
+// SPDX-License-Identifier: MIT
+pragma solidity 0.8.26;
+
+import "@zetachain/protocol-contracts/contracts/evm/GatewayEVM.sol";
+import "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
+
+contract SwapCompanion {
+    using SafeERC20 for IERC20;
+
+    error ApprovalFailed();
+
+    GatewayEVM public immutable gateway;
+
+    constructor(address payable gatewayAddress) {
+        gateway = GatewayEVM(gatewayAddress);
+    }
+
+    function swapNativeGas(
+        address universalSwapContract,
+        address targetToken,
+        bytes memory recipient,
+        bool withdraw
+    ) public payable {
+        gateway.depositAndCall{value: msg.value}(
+            universalSwapContract,
+            abi.encode(targetToken, recipient, withdraw),
+            RevertOptions(msg.sender, false, address(0), "", 0)
+        );
+    }
+
+    function swapERC20(
+        address universalSwapContract,
+        address targetToken,
+        bytes memory recipient,
+        uint256 amount,
+        address asset,
+        bool withdraw
+    ) public {
+        IERC20(asset).safeTransferFrom(msg.sender, address(this), amount);
+        if (!IERC20(asset).approve(address(gateway), amount)) {
+            revert ApprovalFailed();
+        }
+        gateway.depositAndCall(
+            universalSwapContract,
+            amount,
+            asset,
+            abi.encode(targetToken, recipient, withdraw),
+            RevertOptions(msg.sender, false, address(0), "", 0)
+        );
+    }
+
+    receive() external payable {}
+
+    fallback() external payable {}
+}