Improve solidity code for PNA

contracts/src/AgentExecutor.sol

@@ -7,7 +7,6 @@ import {SubstrateTypes} from "./SubstrateTypes.sol";
 
 import {IERC20} from "./interfaces/IERC20.sol";
 import {SafeTokenTransfer, SafeNativeTransfer} from "./utils/SafeTransfer.sol";
-import {ERC20} from "./ERC20.sol";
 import {Gateway} from "./Gateway.sol";
 
 /// @title Code which will run within an `Agent` using `delegatecall`.
@@ -16,40 +15,20 @@ contract AgentExecutor {
     using SafeTokenTransfer for IERC20;
     using SafeNativeTransfer for address payable;
 
-    /// @dev Execute a message which originated from the Polkadot side of the bridge. In other terms,
-    /// the `data` parameter is constructed by the BridgeHub parachain.
-    ///
-    function execute(bytes memory data) external {
-        (AgentExecuteCommand command, bytes memory params) = abi.decode(data, (AgentExecuteCommand, bytes));
-        if (command == AgentExecuteCommand.TransferToken) {
-            (address token, address recipient, uint128 amount) = abi.decode(params, (address, address, uint128));
-            _transferToken(token, recipient, amount);
-        }
-    }
-
     /// @dev Transfer ether to `recipient`. Unlike `_transferToken` This logic is not nested within `execute`,
     /// as the gateway needs to control an agent's ether balance directly.
     ///
     function transferNative(address payable recipient, uint256 amount) external {
         recipient.safeNativeTransfer(amount);
     }
 
-    /// @dev Transfer ERC20 to `recipient`. Only callable via `execute`.
-    function _transferToken(address token, address recipient, uint128 amount) internal {
-        IERC20(token).safeTransfer(recipient, amount);
-    }
-
     /// @dev Transfer ERC20 to `recipient`. Only callable via `execute`.
     function transferToken(address token, address recipient, uint128 amount) external {
         _transferToken(token, recipient, amount);
     }
 
-    /// @dev Mint ERC20 token to `recipient`.
-    function mintToken(address token, address recipient, uint256 amount) external {
-        ERC20(token).mint(recipient, amount);
-    }
-
-    function burnToken(address token, address sender, uint256 amount) external {
-        ERC20(token).burn(sender, amount);
+    /// @dev Transfer ERC20 to `recipient`. Only callable via `execute`.
+    function _transferToken(address token, address recipient, uint128 amount) internal {
+        IERC20(token).safeTransfer(recipient, amount);
     }
 }

contracts/src/Assets.sol

@@ -16,7 +16,7 @@ import {Address} from "./utils/Address.sol";
 import {AgentExecutor} from "./AgentExecutor.sol";
 import {Agent} from "./Agent.sol";
 import {Call} from "./utils/Call.sol";
-import {ERC20} from "./ERC20.sol";
+import {Token} from "./Token.sol";
 
 /// @title Library for implementing Ethereum->Polkadot ERC20 transfers.
 library Assets {
@@ -108,6 +108,34 @@ library Assets {
     ) external returns (Ticket memory ticket) {
         AssetsStorage.Layout storage $ = AssetsStorage.layout();
 
+        TokenInfo storage info = $.tokenRegistry[token];
+
+        if (!info.isRegistered) {
+            revert TokenNotRegistered();
+        }
+
+        if (info.foreignID == bytes32(0)) {
+            return _sendNativeToken(
+                token, sender, destinationChain, destinationAddress, destinationChainFee, maxDestinationChainFee, amount
+            );
+        } else {
+            return _sendForeignToken(
+                info.foreignID, token, sender, destinationChain, destinationAddress, destinationChainFee, amount
+            );
+        }
+    }
+
+    function _sendNativeToken(
+        address token,
+        address sender,
+        ParaID destinationChain,
+        MultiAddress calldata destinationAddress,
+        uint128 destinationChainFee,
+        uint128 maxDestinationChainFee,
+        uint128 amount
+    ) internal returns (Ticket memory ticket) {
+        AssetsStorage.Layout storage $ = AssetsStorage.layout();
+
         // Lock the funds into AssetHub's agent contract
         _transferToAgent($.assetHubAgent, token, sender, amount);
 
@@ -159,6 +187,42 @@ library Assets {
         emit IGateway.TokenSent(token, sender, destinationChain, destinationAddress, amount);
     }
 
+    // @dev Transfer Polkadot-native tokens back to Polkadot
+    function _sendForeignToken(
+        bytes32 foreignID,
+        address token,
+        address sender,
+        ParaID destinationChain,
+        MultiAddress calldata destinationAddress,
+        uint128 destinationChainFee,
+        uint128 amount
+    ) internal returns (Ticket memory ticket) {
+        if (destinationChainFee == 0) {
+            revert InvalidDestinationFee();
+        }
+
+        Token(token).burn(sender, amount);
+
+        ticket.dest = destinationChain;
+        ticket.costs = _sendForeignTokenCosts(destinationChainFee);
+
+        if (destinationAddress.isAddress32()) {
+            // The receiver has a 32-byte account ID
+            ticket.payload = SubstrateTypes.SendForeignTokenToAddress32(
+                foreignID, destinationChain, destinationAddress.asAddress32(), destinationChainFee, amount
+            );
+        } else if (destinationAddress.isAddress20()) {
+            // The receiver has a 20-byte account ID
+            ticket.payload = SubstrateTypes.SendForeignTokenToAddress20(
+                foreignID, destinationChain, destinationAddress.asAddress20(), destinationChainFee, amount
+            );
+        } else {
+            revert Unsupported();
+        }
+
+        emit IGateway.TokenSent(token, sender, destinationChain, destinationAddress, amount);
+    }
+
     function registerTokenCosts() external view returns (Costs memory costs) {
         return _registerTokenCosts();
     }
@@ -195,90 +259,36 @@ library Assets {
         emit IGateway.TokenRegistrationSent(token);
     }
 
-    // @dev Transfer polkadot native tokens back
-    function sendForeignToken(
-        address agent,
-        address executor,
-        TokenInfo storage info,
-        address sender,
-        ParaID destinationChain,
-        MultiAddress calldata destinationAddress,
-        uint128 destinationChainFee,
-        uint128 amount
-    ) external returns (Ticket memory ticket) {
-        if (destinationChainFee == 0) {
-            revert InvalidDestinationFee();
-        }
-        // Polkadot-native token: burn wrapped token
-        _burnToken(executor, agent, info.token, sender, amount);
-
-        ticket.dest = destinationChain;
-        ticket.costs = _sendForeignTokenCosts(destinationChainFee);
-
-        if (destinationAddress.isAddress32()) {
-            // The receiver has a 32-byte account ID
-            ticket.payload = SubstrateTypes.SendForeignTokenToAddress32(
-                info.tokenID, destinationChain, destinationAddress.asAddress32(), destinationChainFee, amount
-            );
-        } else if (destinationAddress.isAddress20()) {
-            // The receiver has a 20-byte account ID
-            ticket.payload = SubstrateTypes.SendForeignTokenToAddress20(
-                info.tokenID, destinationChain, destinationAddress.asAddress20(), destinationChainFee, amount
-            );
-        } else {
-            revert Unsupported();
-        }
-
-        emit IGateway.TokenSent(info.token, sender, destinationChain, destinationAddress, amount);
-    }
-
-    function _burnToken(address agentExecutor, address agent, address token, address sender, uint256 amount) internal {
-        bytes memory call = abi.encodeCall(AgentExecutor.burnToken, (token, sender, amount));
-        (bool success, bytes memory returndata) = (Agent(payable(agent)).invoke(agentExecutor, call));
-        Call.verifyResult(success, returndata);
-    }
-
     function _sendForeignTokenCosts(uint128 destinationChainFee) internal pure returns (Costs memory costs) {
         costs.foreign = destinationChainFee;
         costs.native = 0;
     }
 
     // @dev Register a new fungible Polkadot token for an agent
-    function registerForeignToken(
-        bytes32 agentID,
-        address agent,
-        bytes32 tokenID,
-        string memory name,
-        string memory symbol,
-        uint8 decimals
-    ) external {
+    function registerForeignToken(bytes32 foreignTokenID, string memory name, string memory symbol, uint8 decimals)
+        external
+    {
         AssetsStorage.Layout storage $ = AssetsStorage.layout();
-        if ($.tokenRegistryByID[tokenID].isRegistered == true) {
+        if ($.tokenAddressOf[foreignTokenID] != address(0)) {
             revert TokenAlreadyRegistered();
         }
-        ERC20 foreignToken = new ERC20(agent, name, symbol, decimals);
-        address token = address(foreignToken);
-        TokenInfo memory info =
-            TokenInfo({isRegistered: true, isForeign: true, tokenID: tokenID, agentID: agentID, token: token});
-        $.tokenRegistry[token] = info;
-        $.tokenRegistryByID[tokenID] = info;
-        emit IGateway.ForeignTokenRegistered(tokenID, agentID, token);
+        Token token = new Token(name, symbol, decimals);
+        TokenInfo memory info = TokenInfo({isRegistered: true, foreignID: foreignTokenID});
+
+        $.tokenAddressOf[foreignTokenID] = address(token);
+        $.tokenRegistry[address(token)] = info;
+
+        emit IGateway.ForeignTokenRegistered(foreignTokenID, address(token));
     }
 
     // @dev Mint foreign token from Polkadot
-    function mintForeignToken(address executor, address agent, bytes32 tokenID, address recipient, uint256 amount)
-        external
-    {
-        address token = _tokenAddressOf(tokenID);
-        bytes memory call = abi.encodeCall(AgentExecutor.mintToken, (token, recipient, amount));
-        (bool success,) = Agent(payable(agent)).invoke(executor, call);
-        if (!success) {
-            revert TokenMintFailed();
-        }
+    function mintForeignToken(bytes32 foreignTokenID, address recipient, uint256 amount) external {
+        address token = _ensureTokenAddressOf(foreignTokenID);
+        Token(token).mint(recipient, amount);
     }
 
     // @dev Transfer ERC20 to `recipient`
-    function transferToken(address executor, address agent, address token, address recipient, uint128 amount)
+    function transferNativeToken(address executor, address agent, address token, address recipient, uint128 amount)
         external
     {
         bytes memory call = abi.encodeCall(AgentExecutor.transferToken, (token, recipient, amount));
@@ -290,15 +300,21 @@ library Assets {
 
     // @dev Get token address by tokenID
     function tokenAddressOf(bytes32 tokenID) external view returns (address) {
-        return _tokenAddressOf(tokenID);
+        AssetsStorage.Layout storage $ = AssetsStorage.layout();
+        return $.tokenAddressOf[tokenID];
     }
 
     // @dev Get token address by tokenID
-    function _tokenAddressOf(bytes32 tokenID) internal view returns (address) {
+    function _ensureTokenAddressOf(bytes32 tokenID) internal view returns (address) {
         AssetsStorage.Layout storage $ = AssetsStorage.layout();
-        if ($.tokenRegistryByID[tokenID].isRegistered == false) {
+        if ($.tokenAddressOf[tokenID] == address(0)) {
             revert TokenNotRegistered();
         }
-        return $.tokenRegistryByID[tokenID].token;
+        return $.tokenAddressOf[tokenID];
+    }
+
+    function _isTokenRegistered(address token) internal view returns (bool) {
+        AssetsStorage.Layout storage $ = AssetsStorage.layout();
+        return $.tokenRegistry[token].isRegistered;
     }
 }