Add "wrap/unwrap all" functions; always wrap/unwrap all in QuarkBuilder

src/WrapperScripts.sol

@@ -17,6 +17,13 @@ contract WrapperActions {
         }
     }
 
+    function wrapAllETH(address weth) external payable {
+        uint256 ethBalance = address(this).balance;
+        if (ethBalance > 0) {
+            IWETH(weth).deposit{value: ethBalance}();
+        }
+    }
+
     function unwrapWETH(address weth, uint256 amount) external {
         IWETH(weth).withdraw(amount);
     }
@@ -28,12 +35,34 @@ contract WrapperActions {
         }
     }
 
+    function unwrapAllWETH(address weth) external payable {
+        uint256 wethBalance = IERC20(weth).balanceOf(address(this));
+        if (wethBalance > 0) {
+            IWETH(weth).withdraw(wethBalance);
+        }
+    }
+
     function wrapLidoStETH(address wstETH, address stETH, uint256 amount) external {
         IERC20(stETH).approve(wstETH, amount);
         IWstETH(wstETH).wrap(amount);
     }
 
+    function wrapAllLidoStETH(address wstETH, address stETH) external payable {
+        uint256 stETHBalance = IERC20(stETH).balanceOf(address(this));
+        if (stETHBalance > 0) {
+            IERC20(stETH).approve(wstETH, stETHBalance);
+            IWstETH(wstETH).wrap(stETHBalance);
+        }
+    }
+
     function unwrapLidoWstETH(address wstETH, uint256 amount) external {
         IWstETH(wstETH).unwrap(amount);
     }
+
+    function unwrapAllLidoWstETH(address wstETH) external {
+        uint256 wstETHBalance = IERC20(wstETH).balanceOf(address(this));
+        if (wstETHBalance > 0) {
+            IWstETH(wstETH).unwrap(wstETHBalance);
+        }
+    }
 }

src/builder/BridgeRoutes.sol

@@ -4,6 +4,7 @@ pragma solidity ^0.8.27;
 import {AcrossActions} from "src/AcrossScripts.sol";
 import {CCTPBridgeActions} from "src/BridgeScripts.sol";
 import {Errors} from "src/builder/Errors.sol";
+import {HashMap} from "src/builder/HashMap.sol";
 import {QuarkBuilder} from "src/builder/QuarkBuilder.sol";
 
 import "src/builder/Strings.sol";
@@ -196,4 +197,18 @@ library Across {
             )
         );
     }
+
+    // Returns whether or not an asset is bridged non-deterministically. This applies to WETH/ETH, where Across will send either ETH or WETH
+    // to the target address depending on if address is an EOA or contract.
+    function isNonDeterministicBridgeAction(HashMap.Map memory assetsBridged, string memory assetSymbol)
+        internal
+        pure
+        returns (bool)
+    {
+        uint256 bridgedAmount = HashMap.contains(assetsBridged, abi.encode(assetSymbol))
+            ? HashMap.getUint256(assetsBridged, abi.encode(assetSymbol))
+            : 0;
+        return bridgedAmount > 0
+            && (Strings.stringEqIgnoreCase(assetSymbol, "ETH") || Strings.stringEqIgnoreCase(assetSymbol, "WETH"));
+    }
 }

src/builder/QuarkBuilderBase.sol

@@ -5,7 +5,7 @@ import {IQuarkWallet} from "quark-core/src/interfaces/IQuarkWallet.sol";
 
 import {Actions} from "src/builder/actions/Actions.sol";
 import {Accounts} from "src/builder/Accounts.sol";
-import {BridgeRoutes} from "src/builder/BridgeRoutes.sol";
+import {Across, BridgeRoutes} from "src/builder/BridgeRoutes.sol";
 import {EIP712Helper} from "src/builder/EIP712Helper.sol";
 import {Math} from "src/lib/Math.sol";
 import {MorphoInfo} from "src/builder/MorphoInfo.sol";
@@ -37,7 +37,7 @@ contract QuarkBuilderBase {
 
     /* ===== Constants ===== */
 
-    string constant VERSION = "0.2.0";
+    string constant VERSION = "0.3.0";
 
     /* ===== Custom Errors ===== */
 
@@ -232,13 +232,20 @@ contract QuarkBuilderBase {
                 uint256 supplementalBalance = HashMap.contains(assetsBridged, abi.encode(assetSymbolOut))
                     ? HashMap.getUint256(assetsBridged, abi.encode(assetSymbolOut))
                     : 0;
+                // Note: Right now, ETH/WETH is only bridged via Across. Across has a weird quirk where it will send ETH to EOAs and
+                // WETH to contracts. Since the QuarkBuilder cannot know if a QuarkWallet is deployed before the operation is actually
+                // executed on-chain, we need to assume there is no supplemental balance that arrives on the destination chain.
+                if (Across.isNonDeterministicBridgeAction(assetsBridged, assetSymbolOut)) {
+                    supplementalBalance = 0;
+                }
+
                 checkAndInsertWrapOrUnwrapAction({
                     actions: actions,
                     quarkOperations: quarkOperations,
                     chainAccountsList: chainAccountsList,
                     payment: payment,
                     assetSymbol: assetSymbolOut,
-                    amount: actionIntent.amountOuts[i],
+                    amountNeeded: actionIntent.amountOuts[i],
                     supplementalBalance: supplementalBalance,
                     chainId: actionIntent.chainId,
                     account: actionIntent.actor,
@@ -405,7 +412,7 @@ contract QuarkBuilderBase {
         Accounts.ChainAccounts[] memory chainAccountsList,
         PaymentInfo.Payment memory payment,
         string memory assetSymbol,
-        uint256 amount,
+        uint256 amountNeeded,
         uint256 supplementalBalance,
         uint256 chainId,
         address account,
@@ -415,8 +422,8 @@ contract QuarkBuilderBase {
         // Check if inserting wrapOrUnwrap action is necessary
         uint256 assetBalanceOnChain =
             Accounts.getBalanceOnChain(assetSymbol, chainId, chainAccountsList) + supplementalBalance;
-        if (assetBalanceOnChain < amount && TokenWrapper.hasWrapperContract(chainId, assetSymbol)) {
-            // If the asset has a wrapper counterpart, wrap/unwrap the token to cover the transferIntent amount
+        if (assetBalanceOnChain < amountNeeded && TokenWrapper.hasWrapperContract(chainId, assetSymbol)) {
+            // If the asset has a wrapper counterpart, wrap/unwrap the token to cover the amount needed for the intent
             string memory counterpartSymbol = TokenWrapper.getWrapperCounterpartSymbol(chainId, assetSymbol);
 
             // Wrap/unwrap the token to cover the amount
@@ -425,8 +432,8 @@ contract QuarkBuilderBase {
                 Actions.WrapOrUnwrapAsset({
                     chainAccountsList: chainAccountsList,
                     assetSymbol: counterpartSymbol,
-                    // NOTE: Wrap/unwrap the amount needed to cover the amount
-                    amount: amount - assetBalanceOnChain,
+                    // Note: The wrapper logic should only "wrap all" or "wrap up to" the amount needed
+                    amount: amountNeeded,
                     chainId: chainId,
                     sender: account,
                     blockTimestamp: blockTimestamp

src/builder/TokenWrapper.sol

@@ -108,6 +108,9 @@ library TokenWrapper {
         return Strings.stringEqIgnoreCase(tokenSymbol, getKnownWrapperTokenPair(chainId, tokenSymbol).wrappedSymbol);
     }
 
+    /// Note: We "wrap/unwrap all" for every asset except for ETH/WETH. For ETH/WETH, we will "wrap all ETH" but
+    ///       "unwrap up to X WETH". This is an intentional choice to prefer WETH over ETH since it is much more
+    ///       usable across protocols.
     function encodeActionToWrapOrUnwrap(uint256 chainId, string memory tokenSymbol, uint256 amount)
         internal
         pure
@@ -117,25 +120,24 @@ library TokenWrapper {
         if (isWrappedToken(chainId, tokenSymbol)) {
             return encodeActionToUnwrapToken(chainId, tokenSymbol, amount);
         } else {
-            return encodeActionToWrapToken(chainId, tokenSymbol, pair.underlyingToken, amount);
+            return encodeActionToWrapToken(chainId, tokenSymbol, pair.underlyingToken);
         }
     }
 
-    function encodeActionToWrapToken(uint256 chainId, string memory tokenSymbol, address tokenAddress, uint256 amount)
+    function encodeActionToWrapToken(uint256 chainId, string memory tokenSymbol, address tokenAddress)
         internal
         pure
         returns (bytes memory)
     {
         if (Strings.stringEqIgnoreCase(tokenSymbol, "ETH")) {
             return abi.encodeWithSelector(
-                WrapperActions.wrapETH.selector, getKnownWrapperTokenPair(chainId, tokenSymbol).wrapper, amount
+                WrapperActions.wrapAllETH.selector, getKnownWrapperTokenPair(chainId, tokenSymbol).wrapper
             );
         } else if (Strings.stringEqIgnoreCase(tokenSymbol, "stETH")) {
             return abi.encodeWithSelector(
-                WrapperActions.wrapLidoStETH.selector,
+                WrapperActions.wrapAllLidoStETH.selector,
                 getKnownWrapperTokenPair(chainId, tokenSymbol).wrapper,
-                tokenAddress,
-                amount
+                tokenAddress
             );
         }
         revert NotWrappable();
@@ -148,11 +150,11 @@ library TokenWrapper {
     {
         if (Strings.stringEqIgnoreCase(tokenSymbol, "WETH")) {
             return abi.encodeWithSelector(
-                WrapperActions.unwrapWETH.selector, getKnownWrapperTokenPair(chainId, tokenSymbol).wrapper, amount
+                WrapperActions.unwrapWETHUpTo.selector, getKnownWrapperTokenPair(chainId, tokenSymbol).wrapper, amount
             );
         } else if (Strings.stringEqIgnoreCase(tokenSymbol, "wstETH")) {
             return abi.encodeWithSelector(
-                WrapperActions.unwrapLidoWstETH.selector, getKnownWrapperTokenPair(chainId, tokenSymbol).wrapper, amount
+                WrapperActions.unwrapAllLidoWstETH.selector, getKnownWrapperTokenPair(chainId, tokenSymbol).wrapper
             );
         }
         revert NotUnwrappable();

src/interfaces/IWstETH.sol

@@ -2,6 +2,6 @@
 pragma solidity 0.8.27;
 
 interface IWstETH {
-    function wrap(uint256 amount) external returns (uint256);
-    function unwrap(uint256 amount) external returns (uint256);
+    function wrap(uint256 stETHAmount) external returns (uint256);
+    function unwrap(uint256 wstETHAmount) external returns (uint256);
 }

test/WrapperScripts.t.sol

@@ -115,6 +115,31 @@ contract WrapperScriptsTest is Test {
         assertEq(address(wallet).balance, 10 ether);
     }
 
+    function testWrapAllETH() public {
+        vm.pauseGasMetering();
+        QuarkWallet wallet = QuarkWallet(factory.create(alice, address(0)));
+
+        deal(address(wallet), 10 ether);
+        deal(WETH, address(wallet), 7 ether);
+
+        QuarkWallet.QuarkOperation memory op = new QuarkOperationHelper().newBasicOpWithCalldata(
+            wallet,
+            wrapperScript,
+            abi.encodeWithSelector(WrapperActions.wrapAllETH.selector, WETH),
+            ScriptType.ScriptSource
+        );
+        bytes memory signature = new SignatureHelper().signOp(alicePrivateKey, wallet, op);
+
+        assertEq(IERC20(WETH).balanceOf(address(wallet)), 7 ether);
+        assertEq(address(wallet).balance, 10 ether);
+
+        vm.resumeGasMetering();
+        wallet.executeQuarkOperation(op, signature);
+
+        assertEq(IERC20(WETH).balanceOf(address(wallet)), 17 ether);
+        assertEq(address(wallet).balance, 0 ether);
+    }
+
     function testUnwrapWETH() public {
         vm.pauseGasMetering();
         QuarkWallet wallet = QuarkWallet(factory.create(alice, address(0)));
@@ -187,6 +212,31 @@ contract WrapperScriptsTest is Test {
         assertEq(address(wallet).balance, 10 ether);
     }
 
+    function testUnwrapAllWETH() public {
+        vm.pauseGasMetering();
+        QuarkWallet wallet = QuarkWallet(factory.create(alice, address(0)));
+
+        deal(WETH, address(wallet), 10 ether);
+        deal(address(wallet), 7 ether);
+
+        QuarkWallet.QuarkOperation memory op = new QuarkOperationHelper().newBasicOpWithCalldata(
+            wallet,
+            wrapperScript,
+            abi.encodeWithSelector(WrapperActions.unwrapAllWETH.selector, WETH),
+            ScriptType.ScriptSource
+        );
+        bytes memory signature = new SignatureHelper().signOp(alicePrivateKey, wallet, op);
+
+        assertEq(IERC20(WETH).balanceOf(address(wallet)), 10 ether);
+        assertEq(address(wallet).balance, 7 ether);
+
+        vm.resumeGasMetering();
+        wallet.executeQuarkOperation(op, signature);
+
+        assertEq(IERC20(WETH).balanceOf(address(wallet)), 0 ether);
+        assertEq(address(wallet).balance, 17 ether);
+    }
+
     function testWrapStETH() public {
         vm.pauseGasMetering();
         QuarkWallet wallet = QuarkWallet(factory.create(alice, address(0)));
@@ -208,12 +258,44 @@ contract WrapperScriptsTest is Test {
 
         assertEq(IERC20(wstETH).balanceOf(address(wallet)), 0 ether);
         assertApproxEqAbs(IERC20(stETH).balanceOf(address(wallet)), 10 ether, 0.01 ether);
+
         vm.resumeGasMetering();
         wallet.executeQuarkOperation(op, signature);
+
         assertEq(IERC20(stETH).balanceOf(address(wallet)), 0 ether);
         assertApproxEqAbs(IERC20(wstETH).balanceOf(address(wallet)), 8.74 ether, 0.01 ether);
     }
 
+    function testWrapAllStETH() public {
+        vm.pauseGasMetering();
+        QuarkWallet wallet = QuarkWallet(factory.create(alice, address(0)));
+
+        // Special balance computation in Lido, have to do regular staking action to not mess up the Lido's contract
+        deal(address(wallet), 10 ether);
+        vm.startPrank(address(wallet));
+        // Call Lido's submit() to stake
+        IStETH(stETH).submit{value: 10 ether}(address(0));
+        vm.stopPrank();
+
+        QuarkWallet.QuarkOperation memory op = new QuarkOperationHelper().newBasicOpWithCalldata(
+            wallet,
+            wrapperScript,
+            abi.encodeWithSelector(WrapperActions.wrapAllLidoStETH.selector, wstETH, stETH),
+            ScriptType.ScriptSource
+        );
+        bytes memory signature = new SignatureHelper().signOp(alicePrivateKey, wallet, op);
+
+        assertEq(IERC20(wstETH).balanceOf(address(wallet)), 0 ether);
+        assertApproxEqAbs(IERC20(stETH).balanceOf(address(wallet)), 10 ether, 0.01 ether);
+
+        vm.resumeGasMetering();
+        wallet.executeQuarkOperation(op, signature);
+
+        // Due to shares math in wstETH, the user can be left with a tiny amount of stETH
+        assertApproxEqAbs(IERC20(stETH).balanceOf(address(wallet)), 0 ether, 1);
+        assertApproxEqAbs(IERC20(wstETH).balanceOf(address(wallet)), 8.74 ether, 0.01 ether);
+    }
+
     function testUnwrapWstETH() public {
         vm.pauseGasMetering();
         QuarkWallet wallet = QuarkWallet(factory.create(alice, address(0)));
@@ -230,8 +312,34 @@ contract WrapperScriptsTest is Test {
 
         assertEq(IERC20(stETH).balanceOf(address(wallet)), 0 ether);
         assertEq(IERC20(wstETH).balanceOf(address(wallet)), 10 ether);
+
+        vm.resumeGasMetering();
+        wallet.executeQuarkOperation(op, signature);
+
+        assertApproxEqAbs(IERC20(stETH).balanceOf(address(wallet)), 11.44 ether, 0.01 ether);
+        assertEq(IERC20(wstETH).balanceOf(address(wallet)), 0 ether);
+    }
+
+    function testUnwrapAllWstETH() public {
+        vm.pauseGasMetering();
+        QuarkWallet wallet = QuarkWallet(factory.create(alice, address(0)));
+
+        deal(wstETH, address(wallet), 10 ether);
+
+        QuarkWallet.QuarkOperation memory op = new QuarkOperationHelper().newBasicOpWithCalldata(
+            wallet,
+            wrapperScript,
+            abi.encodeWithSelector(WrapperActions.unwrapAllLidoWstETH.selector, wstETH),
+            ScriptType.ScriptSource
+        );
+        bytes memory signature = new SignatureHelper().signOp(alicePrivateKey, wallet, op);
+
+        assertEq(IERC20(stETH).balanceOf(address(wallet)), 0 ether);
+        assertEq(IERC20(wstETH).balanceOf(address(wallet)), 10 ether);
+
         vm.resumeGasMetering();
         wallet.executeQuarkOperation(op, signature);
+
         assertApproxEqAbs(IERC20(stETH).balanceOf(address(wallet)), 11.44 ether, 0.01 ether);
         assertEq(IERC20(wstETH).balanceOf(address(wallet)), 0 ether);
     }