[SMR-2030] Refactor deposit

src/child/ChildERC20Bridge.sol

@@ -129,6 +129,9 @@ contract ChildERC20Bridge is BridgeRoles, IChildERC20BridgeErrors, IChildERC20Br
         // Initialize
         clonedETHToken.initialize(NATIVE_ETH, "Ethereum", "ETH", 18);
         childETHToken = address(clonedETHToken);
+
+        // Map the supported tokens by default
+        rootTokenToChildToken[NATIVE_ETH] = childETHToken;
     }
 
     /**
@@ -295,81 +298,121 @@ contract ChildERC20Bridge is BridgeRoles, IChildERC20BridgeErrors, IChildERC20Br
         }
 
         address rootToken;
-        uint256 feeAmount = msg.value;
         if (childTokenAddr == NATIVE_IMX) {
-            // Native IMX.
-            if (msg.value < amount) {
-                revert InsufficientValue();
-            }
-
-            feeAmount = msg.value - amount;
-            rootToken = rootIMXToken;
+            rootToken = _withdrawIMX(amount);
         } else if (childTokenAddr == wIMXToken) {
-            // Wrapped IMX.
-            // Transfer and unwrap IMX.
-            uint256 expectedBalance = address(this).balance + amount;
+            rootToken = _withdrawWIMX(amount);
+        } else if (childTokenAddr == childETHToken) {
+            rootToken = _withdrawETH(amount);
+        } else {
+            rootToken = _withdrawERC20(childTokenAddr, amount);
+        }
 
-            IWIMX wIMX = IWIMX(wIMXToken);
-            if (!wIMX.transferFrom(msg.sender, address(this), amount)) {
-                revert TransferWIMXFailed();
-            }
-            wIMX.withdraw(amount);
+        // Encode the message payload
+        bytes memory payload = abi.encode(WITHDRAW_SIG, rootToken, msg.sender, receiver, amount);
 
-            if (address(this).balance != expectedBalance) {
-                revert BalanceInvariantCheckFailed(address(this).balance, expectedBalance);
-            }
+        // Account for fee amount
+        uint256 feeAmount = (childTokenAddr == NATIVE_IMX) ? msg.value - amount : msg.value;
 
-            rootToken = rootIMXToken;
-        } else if (childTokenAddr == childETHToken) {
-            // Wrapped ETH.
-            IChildERC20 childToken = IChildERC20(childTokenAddr);
-            rootToken = NATIVE_ETH;
+        // Send the message to the bridge adaptor and up to root chain
+        bridgeAdaptor.sendMessage{value: feeAmount}(payload, msg.sender);
 
-            if (!childToken.burn(msg.sender, amount)) {
-                revert BurnFailed();
-            }
-        } else {
-            // Other ERC20 Tokens
-            IChildERC20 childToken = IChildERC20(childTokenAddr);
+        _emitWithdrawEvent(rootToken, childTokenAddr, msg.sender, receiver, amount);
+    }
 
-            if (address(childToken).code.length == 0) {
-                revert EmptyTokenContract();
-            }
-            rootToken = childToken.rootToken();
+    /**
+     * @notice Private function to handle withdrawal of L2 native IMX.
+     */
+    function _withdrawIMX(uint256 amount) private returns (address) {
+        if (msg.value < amount) {
+            revert InsufficientValue();
+        }
+        return rootIMXToken;
+    }
 
-            if (rootTokenToChildToken[rootToken] != address(childToken)) {
-                revert NotMapped();
-            }
+    /**
+     * @notice Private function to handle withdrawal of L1 native ETH.
+     */
+    function _withdrawETH(uint256 amount) private returns (address) {
+        if (!IChildERC20(childETHToken).burn(msg.sender, amount)) {
+            revert BurnFailed();
+        }
+        return NATIVE_ETH;
+    }
 
-            // A mapped token should never have root token unset
-            if (rootToken == address(0)) {
-                revert ZeroAddressRootToken();
-            }
+    /**
+     * @notice Private function to handle withdrawal of native IMX.
+     */
+    function _withdrawWIMX(uint256 amount) private returns (address) {
+        // Calculate expected balance
+        uint256 expectedBalance = address(this).balance + amount;
 
-            // A mapped token should never have the bridge unset
-            if (childToken.bridge() != address(this)) {
-                revert IncorrectBridgeAddress();
-            }
+        IWIMX wIMX = IWIMX(wIMXToken);
 
-            if (!childToken.burn(msg.sender, amount)) {
-                revert BurnFailed();
-            }
+        // Transfer to contract
+        if (!wIMX.transferFrom(msg.sender, address(this), amount)) {
+            revert TransferWIMXFailed();
         }
 
-        // Encode the message payload
-        bytes memory payload = abi.encode(WITHDRAW_SIG, rootToken, msg.sender, receiver, amount);
+        // Withdraw
+        wIMX.withdraw(amount);
 
-        // Send the message to the bridge adaptor and up to root chain
-        bridgeAdaptor.sendMessage{value: feeAmount}(payload, msg.sender);
+        // Assert balance
+        if (address(this).balance != expectedBalance) {
+            revert BalanceInvariantCheckFailed(address(this).balance, expectedBalance);
+        }
 
-        if (childTokenAddr == NATIVE_IMX) {
-            emit ChildChainNativeIMXWithdraw(rootToken, msg.sender, receiver, amount);
-        } else if (childTokenAddr == wIMXToken) {
-            emit ChildChainWrappedIMXWithdraw(rootToken, msg.sender, receiver, amount);
-        } else if (childTokenAddr == childETHToken) {
-            emit ChildChainEthWithdraw(msg.sender, receiver, amount);
+        return rootIMXToken;
+    }
+
+    /**
+     * @notice Private function to handle withdrawal of ERC20 tokens.
+     */
+    function _withdrawERC20(address childToken, uint256 amount) private returns (address) {
+        // Validate code existence
+        if (address(childToken).code.length == 0) {
+            revert EmptyTokenContract();
+        }
+
+        address rootToken = IChildERC20(childToken).rootToken();
+
+        // A mapped token should never have root token unset
+        if (rootToken == address(0)) {
+            revert ZeroAddressRootToken();
+        }
+
+        // Assert mapping
+        if (rootTokenToChildToken[rootToken] != address(childToken)) {
+            revert NotMapped();
+        }
+
+        // A mapped token should never have the bridge unset
+        if (IChildERC20(childToken).bridge() != address(this)) {
+            revert IncorrectBridgeAddress();
+        }
+
+        // Burn tokens
+        if (!IChildERC20(childToken).burn(msg.sender, amount)) {
+            revert BurnFailed();
+        }
+
+        return rootToken;
+    }
+
+    /**
+     * @notice Private function to emit withdraw events.
+     */
+    function _emitWithdrawEvent(address rootToken, address childToken, address sender, address receiver, uint256 amount)
+        private
+    {
+        if (childToken == NATIVE_IMX) {
+            emit ChildChainNativeIMXWithdraw(rootToken, sender, receiver, amount);
+        } else if (childToken == wIMXToken) {
+            emit ChildChainWrappedIMXWithdraw(rootToken, sender, receiver, amount);
+        } else if (childToken == childETHToken) {
+            emit ChildChainEthWithdraw(sender, receiver, amount);
         } else {
-            emit ChildChainERC20Withdraw(rootToken, childTokenAddr, msg.sender, receiver, amount);
+            emit ChildChainERC20Withdraw(rootToken, childToken, sender, receiver, amount);
         }
     }
 
@@ -436,33 +479,41 @@ contract ChildERC20Bridge is BridgeRoles, IChildERC20BridgeErrors, IChildERC20Br
             revert ZeroAddress();
         }
 
-        IChildERC20 childToken;
-        if (rootToken != rootIMXToken) {
-            if (rootToken == NATIVE_ETH) {
-                childToken = IChildERC20(childETHToken);
-            } else {
-                childToken = IChildERC20(rootTokenToChildToken[rootToken]);
-                if (address(childToken) == address(0)) {
-                    revert NotMapped();
-                }
-            }
+        transferTokensAndEmitEvent(rootToken, rootTokenToChildToken[rootToken], sender, receiver, amount);
+    }
 
-            if (address(childToken).code.length == 0) {
-                revert EmptyTokenContract();
+    function transferTokensAndEmitEvent(
+        address rootToken,
+        address childToken,
+        address sender,
+        address receiver,
+        uint256 amount
+    ) private {
+        if (rootToken == rootIMXToken) {
+            if (address(this).balance < amount) {
+                revert InsufficientIMX();
             }
+            Address.sendValue(payable(receiver), amount);
+            emit IMXDeposit(rootToken, sender, receiver, amount);
+            return;
+        }
 
-            if (!childToken.mint(receiver, amount)) {
-                revert MintFailed();
-            }
+        if (childToken == address(0)) {
+            revert NotMapped();
+        }
 
-            if (rootToken == NATIVE_ETH) {
-                emit NativeEthDeposit(rootToken, address(childToken), sender, receiver, amount);
-            } else {
-                emit ChildChainERC20Deposit(rootToken, address(childToken), sender, receiver, amount);
-            }
+        if (childToken.code.length == 0) {
+            revert EmptyTokenContract();
+        }
+
+        if (!IChildERC20(childToken).mint(receiver, amount)) {
+            revert MintFailed();
+        }
+
+        if (rootToken == NATIVE_ETH) {
+            emit NativeEthDeposit(rootToken, address(childToken), sender, receiver, amount);
         } else {
-            Address.sendValue(payable(receiver), amount);
-            emit IMXDeposit(rootToken, sender, receiver, amount);
+            emit ChildChainERC20Deposit(rootToken, address(childToken), sender, receiver, amount);
         }
     }
 

src/interfaces/child/IChildERC20Bridge.sol

@@ -218,4 +218,6 @@ interface IChildERC20BridgeErrors {
     error BalanceInvariantCheckFailed(uint256 actualBalance, uint256 expectedBalance);
     /// @notice Error when native transfer is sent to contract from non wrapped-token address.
     error NonWrappedNativeTransfer();
+    /// @notice Error when the bridge doesn't have enough native IMX to support the deposit.
+    error InsufficientIMX();
 }

src/root/RootERC20Bridge.sol

@@ -176,6 +176,11 @@ contract RootERC20Bridge is BridgeRoles, IRootERC20Bridge, IRootERC20BridgeEvent
         childBridgeAdaptor = newChildBridgeAdaptor;
         childChain = newChildChain;
         imxCumulativeDepositLimit = newImxCumulativeDepositLimit;
+
+        // Map the supported tokens by default
+        rootTokenToChildToken[rootIMXToken] = rootIMXToken;
+        rootTokenToChildToken[NATIVE_ETH] = NATIVE_ETH;
+        rootTokenToChildToken[rootWETHToken] = NATIVE_ETH;
     }
 
     /**
@@ -404,7 +409,11 @@ contract RootERC20Bridge is BridgeRoles, IRootERC20Bridge, IRootERC20BridgeEvent
         return childToken;
     }
 
-    function _deposit(IERC20Metadata rootToken, address receiver, uint256 amount) private whenNotPaused {
+    function _deposit(IERC20Metadata rootToken, address receiver, uint256 amount)
+        private
+        whenNotPaused
+        wontIMXOverflow(address(rootToken), amount)
+    {
         if (receiver == address(0) || address(rootToken) == address(0)) {
             revert ZeroAddress();
         }
@@ -414,52 +423,46 @@ contract RootERC20Bridge is BridgeRoles, IRootERC20Bridge, IRootERC20BridgeEvent
         if (msg.value == 0) {
             revert NoGas();
         }
-        if (
-            address(rootToken) == rootIMXToken && imxCumulativeDepositLimit != UNLIMITED_DEPOSIT
-                && IERC20Metadata(rootIMXToken).balanceOf(address(this)) + amount > imxCumulativeDepositLimit
-        ) {
-            revert ImxDepositLimitExceeded();
-        }
-
         // ETH, WETH and IMX do not need to be mapped since it should have been mapped on initialization
-        // ETH also cannot be transferred since it was received in the payable function call
-        // WETH is also not transferred here since it was earlier unwrapped to ETH
+        if (rootTokenToChildToken[address(rootToken)] == address(0)) {
+            revert NotMapped();
+        }
 
-        // TODO We can call _mapToken here, but ordering in the GMP is not guaranteed.
+        // We can call _mapToken here, but ordering in the GMP is not guaranteed.
         // Therefore, we need to decide how to handle this and it may be a UI decision to wait until map token message is executed on child chain.
         // Discuss this, and add this decision to the design doc.
 
-        address childToken;
-        uint256 feeAmount = msg.value;
-        address payloadToken = address(rootToken);
-
-        if (address(rootToken) == NATIVE_ETH) {
-            feeAmount = msg.value - amount;
-        } else if (address(rootToken) == rootWETHToken) {
-            payloadToken = NATIVE_ETH;
-        } else {
-            if (address(rootToken) != rootIMXToken) {
-                childToken = rootTokenToChildToken[address(rootToken)];
-                if (childToken == address(0)) {
-                    revert NotMapped();
-                }
-            }
-            rootToken.safeTransferFrom(msg.sender, address(this), amount);
-        }
+        address payloadToken = (address(rootToken) == rootWETHToken) ? NATIVE_ETH : address(rootToken);
 
         // Deposit sig, root token address, depositor, receiver, amount
         bytes memory payload = abi.encode(DEPOSIT_SIG, payloadToken, msg.sender, receiver, amount);
 
+        // Adjust for fee amount on native transfers
+        uint256 feeAmount = (address(rootToken) == NATIVE_ETH) ? msg.value - amount : msg.value;
+
+        // Send message to child chain
         rootBridgeAdaptor.sendMessage{value: feeAmount}(payload, msg.sender);
 
-        if (address(rootToken) == NATIVE_ETH) {
-            emit NativeEthDeposit(address(rootToken), childETHToken, msg.sender, receiver, amount);
-        } else if (address(rootToken) == rootWETHToken) {
-            emit WETHDeposit(address(rootToken), childETHToken, msg.sender, receiver, amount);
-        } else if (address(rootToken) == rootIMXToken) {
-            emit IMXDeposit(address(rootToken), msg.sender, receiver, amount);
+        // Emit the appropriate deposit event
+        transferTokensAndEmitEvent(address(rootToken), receiver, amount);
+    }
+
+    /**
+     * @notice Private helper function to emit the appropriate deposit event and execute transfer if rootIMX or rootERC20
+     */
+    function transferTokensAndEmitEvent(address rootToken, address receiver, uint256 amount) private {
+        // ETH also cannot be transferred since it was received in the payable function call
+        if (rootToken == NATIVE_ETH) {
+            emit NativeEthDeposit(rootToken, childETHToken, msg.sender, receiver, amount);
+            // WETH is also not transferred here since it was earlier unwrapped to ETH
+        } else if (rootToken == rootWETHToken) {
+            emit WETHDeposit(rootToken, childETHToken, msg.sender, receiver, amount);
+        } else if (rootToken == rootIMXToken) {
+            emit IMXDeposit(rootToken, msg.sender, receiver, amount);
+            IERC20Metadata(rootToken).safeTransferFrom(msg.sender, address(this), amount);
         } else {
-            emit ChildChainERC20Deposit(address(rootToken), childToken, msg.sender, receiver, amount);
+            emit ChildChainERC20Deposit(rootToken, rootTokenToChildToken[rootToken], msg.sender, receiver, amount);
+            IERC20Metadata(rootToken).safeTransferFrom(msg.sender, address(this), amount);
         }
     }
 
@@ -508,6 +511,17 @@ contract RootERC20Bridge is BridgeRoles, IRootERC20Bridge, IRootERC20BridgeEvent
         }
     }
 
+    modifier wontIMXOverflow(address rootToken, uint256 amount) {
+        // Assert whether the deposit is root IMX
+        if (address(rootToken) == rootIMXToken && imxCumulativeDepositLimit != UNLIMITED_DEPOSIT) {
+            // Based on the balance of this contract, check if the deposit will exceed the cumulative limit
+            if (IERC20Metadata(rootIMXToken).balanceOf(address(this)) + amount > imxCumulativeDepositLimit) {
+                revert ImxDepositLimitExceeded();
+            }
+        }
+        _;
+    }
+
     // slither-disable-next-line unused-state,naming-convention
     uint256[50] private __gapRootERC20Bridge;
 }