Refactor

zeta-chain · Nov 22, 2024 · 2d82f6c · 2d82f6c
1 parent 581f201
commit 2d82f6c
Show file tree

Hide file tree

Showing 2 changed files with 58 additions and 40 deletions.
diff --git a/examples/swap/contracts/Swap.sol b/examples/swap/contracts/Swap.sol
@@ -57,26 +57,19 @@ contract Swap is UniversalContract {
             params.to = recipient;
         }
 
-        uint256 out = handleGasAndSwap(zrc20, amount, params.target);
-        gateway.withdraw(
-            params.to,
-            out,
-            params.target,
-            RevertOptions({
-                revertAddress: address(this),
-                callOnRevert: true,
-                abortAddress: address(0),
-                revertMessage: abi.encode(context.sender, zrc20),
-                onRevertGasLimit: 100000
-            })
+        (uint256 out, address gasZRC20, uint256 gasFee) = handleGasAndSwap(
+            zrc20,
+            amount,
+            params.target
         );
+        withdraw(params, context.sender, gasFee, gasZRC20, out, zrc20);
     }
 
     function handleGasAndSwap(
         address inputToken,
         uint256 amount,
         address targetToken
-    ) internal returns (uint256) {
+    ) internal returns (uint256, address, uint256) {
         uint256 inputForGas;
         address gasZRC20;
         uint256 gasFee;
@@ -104,22 +97,50 @@ contract Swap is UniversalContract {
             targetToken,
             0
         );
-        return outputAmount;
+        return (outputAmount, gasZRC20, gasFee);
+    }
+
+    function withdraw(
+        Params memory params,
+        address sender,
+        uint256 gasFee,
+        address gasZRC20,
+        uint256 outputAmount,
+        address inputToken
+    ) public {
+        if (gasZRC20 == params.target) {
+            IZRC20(gasZRC20).approve(address(gateway), outputAmount + gasFee);
+        } else {
+            IZRC20(gasZRC20).approve(address(gateway), gasFee);
+            IZRC20(params.target).approve(address(gateway), outputAmount);
+        }
+        gateway.withdraw(
+            params.to,
+            outputAmount,
+            params.target,
+            RevertOptions({
+                revertAddress: address(this),
+                callOnRevert: true,
+                abortAddress: address(0),
+                revertMessage: abi.encode(sender, inputToken),
+                onRevertGasLimit: 100000
+            })
+        );
     }
 
     function onRevert(RevertContext calldata context) external onlyGateway {
         (address sender, address zrc20) = abi.decode(
             context.revertMessage,
             (address, address)
         );
-        uint256 outputAmount = handleGasAndSwap(
+        (uint256 out, , ) = handleGasAndSwap(
             context.asset,
             context.amount,
             zrc20
         );
         gateway.withdraw(
             abi.encodePacked(sender),
-            outputAmount,
+            out,
             zrc20,
             RevertOptions({
                 revertAddress: sender,

diff --git a/examples/swap/contracts/SwapToAnyToken.sol b/examples/swap/contracts/SwapToAnyToken.sol
@@ -69,12 +69,12 @@ contract SwapToAnyToken is UniversalContract {
             params.withdraw = withdrawFlag;
         }
 
-        (
-            uint256 outputAmount,
-            address gasZRC20,
-            uint256 gasFee
-        ) = handleGasAndSwap(zrc20, amount, params.target);
-        withdraw(params, context.sender, gasFee, gasZRC20, outputAmount, zrc20);
+        (uint256 out, address gasZRC20, uint256 gasFee) = handleGasAndSwap(
+            zrc20,
+            amount,
+            params.target
+        );
+        withdraw(params, context.sender, gasFee, gasZRC20, out, zrc20);
     }
 
     function swap(
@@ -85,11 +85,11 @@ contract SwapToAnyToken is UniversalContract {
         bool withdrawFlag
     ) public {
         IZRC20(inputToken).transferFrom(msg.sender, address(this), amount);
-        (
-            uint256 outputAmount,
-            address gasZRC20,
-            uint256 gasFee
-        ) = handleGasAndSwap(inputToken, amount, targetToken);
+        (uint256 out, address gasZRC20, uint256 gasFee) = handleGasAndSwap(
+            inputToken,
+            amount,
+            targetToken
+        );
         withdraw(
             Params({
                 target: targetToken,
@@ -99,7 +99,7 @@ contract SwapToAnyToken is UniversalContract {
             msg.sender,
             gasFee,
             gasZRC20,
-            outputAmount,
+            out,
             inputToken
         );
     }
@@ -129,37 +129,34 @@ contract SwapToAnyToken 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 (params.withdraw) {
             if (gasZRC20 == params.target) {
-                IZRC20(gasZRC20).approve(
-                    address(gateway),
-                    outputAmount + gasFee
-                );
+                IZRC20(gasZRC20).approve(address(gateway), out + gasFee);
             } else {
                 IZRC20(gasZRC20).approve(address(gateway), gasFee);
-                IZRC20(params.target).approve(address(gateway), outputAmount);
+                IZRC20(params.target).approve(address(gateway), out);
             }
             gateway.withdraw(
                 abi.encodePacked(params.to),
-                outputAmount,
+                out,
                 params.target,
                 RevertOptions({
                     revertAddress: address(this),
@@ -172,7 +169,7 @@ contract SwapToAnyToken is UniversalContract {
         } else {
             IWETH9(params.target).transfer(
                 address(uint160(bytes20(params.to))),
-                outputAmount
+                out
             );
         }
     }
@@ -182,15 +179,15 @@ contract SwapToAnyToken is UniversalContract {
             context.revertMessage,
             (address, address)
         );
-        (uint256 outputAmount, , ) = handleGasAndSwap(
+        (uint256 out, , ) = handleGasAndSwap(
             context.asset,
             context.amount,
             zrc20
         );
 
         gateway.withdraw(
             abi.encodePacked(sender),
-            outputAmount,
+            out,
             zrc20,
             RevertOptions({
                 revertAddress: sender,