From 45fee2a07a6d11d34c06ab238dc2c05e7a9aff20 Mon Sep 17 00:00:00 2001
From: Alistair Singh <alistair.singh7@gmail.com>
Date: Mon, 21 Oct 2024 14:04:46 +0200
Subject: [PATCH] convert to address lookup

---
 contracts/src/Gateway.sol             |  5 ++---
 contracts/src/interfaces/IGateway.sol |  4 ++--
 contracts/test/ForkUpgrade.t.sol      | 15 ++++++---------
 3 files changed, 10 insertions(+), 14 deletions(-)

diff --git a/contracts/src/Gateway.sol b/contracts/src/Gateway.sol
index df013c0714..316def08dc 100644
--- a/contracts/src/Gateway.sol
+++ b/contracts/src/Gateway.sol
@@ -18,7 +18,6 @@ import {
     MultiAddress,
     Ticket,
     Costs,
-    TokenInfo,
     AgentExecuteCommand
 } from "./Types.sol";
 import {Upgrade} from "./Upgrade.sol";
@@ -439,8 +438,8 @@ contract Gateway is IGateway, IInitializable, IUpgradable {
         return Assets.isTokenRegistered(token);
     }
 
-    function tokenInfo(address token) external view returns (TokenInfo memory) {
-        return AssetsStorage.layout().tokenRegistry[token];
+    function tokenForeignIDOf(address token) external view returns (bytes32) {
+        return AssetsStorage.layout().tokenRegistry[token].foreignID;
     }
 
     // Total fee for registering a token
diff --git a/contracts/src/interfaces/IGateway.sol b/contracts/src/interfaces/IGateway.sol
index b502876242..086fa2daa0 100644
--- a/contracts/src/interfaces/IGateway.sol
+++ b/contracts/src/interfaces/IGateway.sol
@@ -86,8 +86,8 @@ interface IGateway {
     /// @dev Check whether a token is registered
     function isTokenRegistered(address token) external view returns (bool);
 
-    /// @dev Get information about a token.
-    function tokenInfo(address token) external view returns (TokenInfo memory);
+    /// @dev Get token id an ERC20 contract address.
+    function tokenForeignIDOf(address token) external view returns (bytes32);
 
     /// @dev Quote a fee in Ether for registering a token, covering
     /// 1. Delivery costs to BridgeHub
diff --git a/contracts/test/ForkUpgrade.t.sol b/contracts/test/ForkUpgrade.t.sol
index b0a325951d..f10a8db7da 100644
--- a/contracts/test/ForkUpgrade.t.sol
+++ b/contracts/test/ForkUpgrade.t.sol
@@ -40,12 +40,10 @@ contract ForkUpgradeTest is Test {
     }
 
     function checkLegacyToken() public {
-        TokenInfo memory weth = GatewayPNA(GatewayProxy).tokenInfo(0xC02aaA39b223FE8D0A0e5C4F27eAD9083C756Cc2);
-        assertEq(weth.isRegistered, true);
-        assertEq(weth.foreignID, bytes32(""));
-        TokenInfo memory myth = GatewayPNA(GatewayProxy).tokenInfo(0xBA41Ddf06B7fFD89D1267b5A93BFeF2424eb2003);
-        assertEq(myth.isRegistered, true);
-        assertEq(myth.foreignID, bytes32(""));
+        assert(IGateway(GatewayProxy).isTokenRegistered(0xC02aaA39b223FE8D0A0e5C4F27eAD9083C756Cc2));
+        assertEq(IGateway(GatewayProxy).tokenForeignIDOf(0xC02aaA39b223FE8D0A0e5C4F27eAD9083C756Cc2), bytes32(""));
+        assert(IGateway(GatewayProxy).isTokenRegistered(0xBA41Ddf06B7fFD89D1267b5A93BFeF2424eb2003));
+        assertEq(IGateway(GatewayProxy).tokenForeignIDOf(0xBA41Ddf06B7fFD89D1267b5A93BFeF2424eb2003), bytes32(""));
     }
 
     function registerForeignToken() public {
@@ -57,9 +55,8 @@ contract ForkUpgradeTest is Test {
         emit IGateway.ForeignTokenRegistered(dotId, address(0x0));
 
         GatewayPNA(GatewayProxy).registerForeignToken(abi.encode(params));
-        TokenInfo memory dot = GatewayPNA(GatewayProxy).tokenInfo(0x70D9d338A6b17957B16836a90192BD8CDAe0b53d);
-        assertEq(dot.isRegistered, true);
-        assertEq(dot.foreignID, dotId);
+        assert(IGateway(GatewayProxy).isTokenRegistered(0x70D9d338A6b17957B16836a90192BD8CDAe0b53d));
+        assertEq(IGateway(GatewayProxy).tokenForeignIDOf(0x70D9d338A6b17957B16836a90192BD8CDAe0b53d), dotId);
     }
 
     function testSanityCheck() public {