Refactor/ Portfolio lookups and reduce formatting

src/libs/portfolio/getOnchainBalances.ts

@@ -218,19 +218,27 @@ export async function getNFTs(
   // simulation was performed if the nonce is changed
   const hasSimulation = afterNonce !== beforeNonce
 
-  const simulationTokens: (CollectionResult & { addr: any })[] | null = hasSimulation
-    ? after.collections.map((simulationToken: any, tokenIndex: number) => ({
-        ...mapNft(simulationToken, deltaAddressesMapping[tokenIndex]),
-        addr: deltaAddressesMapping[tokenIndex]
-      }))
-    : null
+  // Index all to prevent nested loops
+  const simulationTokensByAddr = new Map<string, any>()
+
+  if (hasSimulation) {
+    after.collections.forEach((simulationToken: any, tokenIndex: number) => {
+      const addr = deltaAddressesMapping[tokenIndex]
+
+      if (addr === undefined) return
+
+      const key = addr.toLowerCase()
+
+      if (simulationTokensByAddr.has(key)) return
+
+      simulationTokensByAddr.set(key, { ...mapNft(simulationToken, addr), addr })
+    })
+  }
 
   return [
     before.collections.map((beforeToken: any, i: number) => {
-      const simulationToken = simulationTokens
-        ? simulationTokens.find(
-            (token: any) => token.addr.toLowerCase() === tokenAddrs[i]![0].toLowerCase()
-          )
+      const simulationToken = hasSimulation
+        ? simulationTokensByAddr.get(tokenAddrs[i]![0].toLowerCase())
         : null
 
       const token = mapNft(beforeToken, tokenAddrs[i]![0])
@@ -350,18 +358,22 @@ export async function getTokens(
   // simulation was performed if the nonce is changed
   const hasSimulation = afterNonce !== beforeNonce
 
-  const simulationTokens = hasSimulation
-    ? after.balances.map((simulationToken: any, tokenIndex: number) => ({
-        ...simulationToken,
-        amount: simulationToken.amount,
-        addr: deltaAddressesMapping[tokenIndex]
-      }))
-    : null
+  // Index all to prevent nested loops
+  const simulationTokensByAddr = new Map<string, any>()
+
+  if (hasSimulation) {
+    after.balances.forEach((simulationToken: any, tokenIndex: number) => {
+      const addr = deltaAddressesMapping[tokenIndex]
+
+      if (addr === undefined || simulationTokensByAddr.has(addr)) return
+
+      simulationTokensByAddr.set(addr, { ...simulationToken, addr })
+    })
+  }
+
   return [
     before.balances.map((token: any, i: number) => {
-      const simulation = simulationTokens
-        ? simulationTokens.find((simulationToken: any) => simulationToken.addr === tokenAddrs[i])
-        : null
+      const simulation = hasSimulation ? (simulationTokensByAddr.get(tokenAddrs[i]!) ?? null) : null
 
       const simulationAmount = simulation ? simulation.amount - token.amount : undefined
       const amountPostSimulation = simulation ? simulation.amount : token.amount

src/libs/portfolio/helpers.ts

@@ -521,12 +521,23 @@ export const erc721CollectionToLearnedAssetKeys = (collection: [string, bigint[]
  */
 export const learnedErc721sToHints = (keys: string[]): ERC721s => {
   const hints: ERC721s = {}
+  // Split once and collect the enumerable collections up front. Checking for an
+  // enumerable key while building the hints would mean scanning every key for
+  // every key, and an account with many collections brings thousands of them.
+  const parsedKeys: [string, string | undefined][] = []
+  const enumerableCollections = new Set<string>()
 
   keys.forEach((key) => {
     const [collectionAddress, tokenId] = key.split(':')
 
     if (!collectionAddress) return
 
+    parsedKeys.push([collectionAddress, tokenId])
+
+    if (tokenId === 'enumerable') enumerableCollections.add(collectionAddress)
+  })
+
+  parsedKeys.forEach(([collectionAddress, tokenId]) => {
     if (tokenId === 'enumerable') {
       hints[collectionAddress] = []
 
@@ -535,7 +546,7 @@ export const learnedErc721sToHints = (keys: string[]): ERC721s => {
     // The key already exists as an enumerable hint. Example:
     // collectionA:enumerable exists and collectionB:id is attempted to be added
     // (it shouldn't be)
-    if (keys.includes(`${collectionAddress}:enumerable`)) {
+    if (enumerableCollections.has(collectionAddress)) {
       return
     }
 

src/libs/portfolio/portfolio.ts

@@ -1,6 +1,8 @@
 import { ZeroAddress } from 'ethers'
 import { getAddress } from 'viem'
 
+import { getFeeToken } from '@/libs/portfolio/tokenProcessing'
+
 import BalanceGetter from '../../../contracts/compiled/BalanceGetter.json'
 import NFTGetter from '../../../contracts/compiled/NFTGetter.json'
 import gasTankFeeTokens from '../../consts/gasTankFeeTokens'
@@ -245,17 +247,26 @@ export class Portfolio {
       ...Object.values(specialErc721Hints || {})
     ])
 
-    const checksummedErc20Hints = hints.erc20s
-      .map((address) => {
-        try {
-          // getAddress may throw an error. This will break the portfolio
-          // if the error isn't caught
-          return getAddress(address)
-        } catch {
-          return null
-        }
-      })
-      .filter(Boolean) as string[]
+    // Deduped before checksumming for performance
+    const seenErc20Hints = new Set<string>()
+    const checksummedErc20Hints: string[] = []
+
+    hints.erc20s.forEach((address) => {
+      try {
+        const lowercasedAddress = address.toLowerCase()
+
+        if (seenErc20Hints.has(lowercasedAddress)) return
+
+        // getAddress may throw an error. This will break the portfolio
+        // if the error isn't caught
+        const checksummedAddress = getAddress(address)
+
+        seenErc20Hints.add(lowercasedAddress)
+        checksummedErc20Hints.push(checksummedAddress)
+      } catch {
+        // Not an address, so it can't be a token
+      }
+    })
 
     // Merge static and dynamic blacklisted addresses for this chain
     const chainIdStr = this.network.chainId.toString()
@@ -575,11 +586,7 @@ export class Portfolio {
         // return the native token
         if (t.address === ZeroAddress && t.chainId === this.network.chainId) return true
 
-        return gasTankFeeTokens.find(
-          (gasTankT) =>
-            gasTankT.address.toLowerCase() === t.address.toLowerCase() &&
-            gasTankT.chainId === t.chainId
-        )
+        return getFeeToken(t.address, t.chainId)
       }),
       beforeNonce,
       afterNonce,

src/libs/portfolio/tokenProcessing.test.ts

@@ -0,0 +1,92 @@
+import { ZeroAddress } from 'ethers'
+
+import { describe, expect, it } from '@jest/globals'
+
+import gasTankFeeTokens from '../../consts/gasTankFeeTokens'
+import { getFeeToken, getFlags } from './tokenProcessing'
+
+const USDT_ETHEREUM = '0xdAC17F958D2ee523a2206206994597C13D831ec7'
+const WETH_OPTIMISM = '0x4200000000000000000000000000000000000006'
+const DUPLICATED_ON_AVALANCHE = '0xB97EF9Ef8734C71904D8002F8b6Bc66Dd9c48a6E'
+const NOT_A_FEE_TOKEN = '0xdeadbeefdeadbeefdeadbeefdeadbeefdeadbeef'
+
+describe('getFeeToken', () => {
+  it('returns the first of two entries sharing an address and a chain', () => {
+    const duplicates = gasTankFeeTokens.filter(
+      (t) =>
+        t.address.toLowerCase() === DUPLICATED_ON_AVALANCHE.toLowerCase() && t.chainId === 43114n
+    )
+
+    expect(duplicates.length).toBeGreaterThan(1)
+    expect(getFeeToken(DUPLICATED_ON_AVALANCHE, 43114n)).toBe(duplicates[0])
+  })
+
+  it('is case-insensitive on the given address', () => {
+    const usdt = gasTankFeeTokens.find(
+      (t) => t.address.toLowerCase() === USDT_ETHEREUM.toLowerCase() && t.chainId === 1n
+    )
+
+    expect(usdt).toBeDefined()
+    expect(getFeeToken(USDT_ETHEREUM.toLowerCase(), 1n)).toBe(usdt)
+    expect(getFeeToken(USDT_ETHEREUM.toUpperCase(), 1n)).toBe(usdt)
+  })
+
+  it('returns undefined for an address that is not a fee token', () => {
+    expect(getFeeToken(NOT_A_FEE_TOKEN, 1n)).toBeUndefined()
+    expect(getFeeToken(NOT_A_FEE_TOKEN, 1n)).toBeUndefined()
+  })
+
+  it('returns undefined when the address is a fee token but on another chain', () => {
+    const wethOnOptimism = gasTankFeeTokens.find(
+      (t) => t.address.toLowerCase() === WETH_OPTIMISM.toLowerCase() && t.chainId === 10n
+    )
+
+    expect(wethOnOptimism).toBeDefined()
+    expect(getFeeToken(WETH_OPTIMISM, 1n)).toBeUndefined()
+    expect(getFeeToken(WETH_OPTIMISM, 1n)).toBeUndefined()
+  })
+
+  it('reuses the index across calls instead of rebuilding it', () => {
+    expect(getFeeToken(USDT_ETHEREUM, 1n)).toBe(getFeeToken(USDT_ETHEREUM, 1n))
+  })
+})
+
+describe('getFlags fee token flags', () => {
+  it('marks a gas tank fee token as topped up and usable as a fee', () => {
+    const usdt = gasTankFeeTokens.find(
+      (t) => t.address.toLowerCase() === USDT_ETHEREUM.toLowerCase() && t.chainId === 1n
+    )!
+
+    expect(usdt.disableGasTankDeposit).toBeFalsy()
+    expect(usdt.disableAsFeeToken).toBeFalsy()
+
+    const flags = getFlags({}, '1', 1n, USDT_ETHEREUM, 'Tether USD', 'USDT')
+
+    expect(flags.canTopUpGasTank).toBe(true)
+    expect(flags.isFeeToken).toBe(true)
+    expect(flags.onGasTank).toBe(false)
+  })
+
+  it('does not mark an unknown token as a fee token', () => {
+    const flags = getFlags({}, '1', 1n, NOT_A_FEE_TOKEN, 'Random', 'RND')
+
+    expect(flags.canTopUpGasTank).toBe(false)
+    expect(flags.isFeeToken).toBeFalsy()
+  })
+
+  it('treats the native token as a fee token even without a gas tank entry', () => {
+    const flags = getFlags({}, '31337', 31337n, ZeroAddress, 'Ether', 'ETH')
+
+    expect(getFeeToken(ZeroAddress, 31337n)).toBeUndefined()
+    expect(flags.isFeeToken).toBe(true)
+    expect(flags.canTopUpGasTank).toBe(false)
+  })
+
+  it('resolves fee tokens on the gasTank pseudo chain by the token chain id', () => {
+    const flags = getFlags({}, 'gasTank', 1n, USDT_ETHEREUM, 'Tether USD', 'USDT')
+
+    expect(flags.onGasTank).toBe(true)
+    expect(flags.canTopUpGasTank).toBe(true)
+    expect(flags.isFeeToken).toBe(true)
+  })
+})

src/libs/portfolio/tokenProcessing.ts

@@ -83,6 +83,34 @@ export const isSuspectedToken = (
   return null
 }
 
+let feeTokenIndex: Map<string, (typeof gasTankFeeTokens)[number]> | null = null
+
+const feeTokenKey = (address: string, chainId: string) => `${address.toLowerCase()}|${chainId}`
+
+const getFeeTokenIndex = () => {
+  if (feeTokenIndex) return feeTokenIndex
+
+  feeTokenIndex = new Map<string, (typeof gasTankFeeTokens)[number]>()
+
+  gasTankFeeTokens.forEach((feeToken) => {
+    const key = feeTokenKey(feeToken.address, feeToken.chainId.toString())
+
+    if (!feeTokenIndex!.has(key)) feeTokenIndex!.set(key, feeToken)
+  })
+
+  return feeTokenIndex
+}
+
+/**
+ * Look up a gas-tank fee token by address and chain in O(1)
+ */
+export function getFeeToken(
+  address: string,
+  chainid: bigint
+): (typeof gasTankFeeTokens)[number] | undefined {
+  return getFeeTokenIndex().get(feeTokenKey(address, chainid.toString()))
+}
+
 export function getFlags(
   networkData: any,
   chainId: string,
@@ -101,11 +129,7 @@ export function getFlags(
   if (networkData?.walletClaimableBalance?.address.toLowerCase() === address.toLowerCase())
     rewardsType = 'wallet-vesting'
 
-  const foundFeeToken = gasTankFeeTokens.find(
-    (t) =>
-      t.address.toLowerCase() === address.toLowerCase() &&
-      (isRewardsOrGasTank ? t.chainId === tokenChainId : t.chainId.toString() === chainId)
-  )
+  const foundFeeToken = getFeeToken(address, tokenChainId)
 
   const canTopUpGasTank = !!foundFeeToken && !foundFeeToken?.disableGasTankDeposit && !rewardsType
   const isFeeToken =