feat: update useAllTokensMultichain usage (#6493)

* feat: updates useAllTokensMultichain to return userAddedTokens

* fix: use correct types in tests

* chore: add documentation for future removal of TokenAddressMap

* fix: use doc comments

* test: add unit test for useAllTokensMultichain

* fix: check userAddedTokens for undefined

src/components/AccountDrawer/MiniPortfolio/Activity/parseLocal.test.ts

@@ -2,7 +2,7 @@ import { SupportedChainId, Token, TradeType as MockTradeType } from '@uniswap/sd
 import { PERMIT2_ADDRESS } from '@uniswap/universal-router-sdk'
 import { DAI as MockDAI, nativeOnChain, USDC_MAINNET as MockUSDC_MAINNET } from 'constants/tokens'
 import { TransactionStatus as MockTxStatus } from 'graphql/data/__generated__/types-and-hooks'
-import { TokenAddressMap } from 'state/lists/hooks'
+import { ChainTokenMap } from 'hooks/Tokens'
 import {
   ExactInputSwapTransactionInfo,
   ExactOutputSwapTransactionInfo,
@@ -89,15 +89,15 @@ function mockMultiStatus(info: TransactionInfo, id: string): [TransactionDetails
   ]
 }
 
-const mockTokenAddressMap: TokenAddressMap = {
+const mockTokenAddressMap: ChainTokenMap = {
   [mockChainId]: {
-    [MockDAI.address]: { token: MockDAI },
-    [MockUSDC_MAINNET.address]: { token: MockUSDC_MAINNET },
-  } as TokenAddressMap[number],
+    [MockDAI.address]: MockDAI,
+    [MockUSDC_MAINNET.address]: MockUSDC_MAINNET,
+  },
 }
 
-jest.mock('../../../../state/lists/hooks', () => ({
-  useCombinedActiveList: () => mockTokenAddressMap,
+jest.mock('../../../../hooks/Tokens', () => ({
+  useAllTokensMultichain: () => mockTokenAddressMap,
 }))
 
 jest.mock('../../../../state/transactions/hooks', () => {
@@ -300,7 +300,7 @@ describe('parseLocalActivity', () => {
       },
     } as TransactionDetails
     const chainId = SupportedChainId.MAINNET
-    const tokens = {} as TokenAddressMap
+    const tokens = {} as ChainTokenMap
     expect(parseLocalActivity(details, chainId, tokens)).toMatchObject({
       chainId: 1,
       currencies: [undefined, undefined],

src/components/AccountDrawer/MiniPortfolio/Activity/parseLocal.ts

@@ -4,8 +4,8 @@ import { Currency, CurrencyAmount, TradeType } from '@uniswap/sdk-core'
 import { nativeOnChain } from '@uniswap/smart-order-router'
 import { SupportedChainId } from 'constants/chains'
 import { TransactionPartsFragment, TransactionStatus } from 'graphql/data/__generated__/types-and-hooks'
+import { ChainTokenMap, useAllTokensMultichain } from 'hooks/Tokens'
 import { useMemo } from 'react'
-import { TokenAddressMap, useCombinedActiveList } from 'state/lists/hooks'
 import { useMultichainTransactions } from 'state/transactions/hooks'
 import {
   AddLiquidityV2PoolTransactionInfo,
@@ -25,8 +25,8 @@ import {
 import { getActivityTitle } from '../constants'
 import { Activity, ActivityMap } from './types'
 
-function getCurrency(currencyId: string, chainId: SupportedChainId, tokens: TokenAddressMap): Currency | undefined {
-  return currencyId === 'ETH' ? nativeOnChain(chainId) : tokens[chainId]?.[currencyId]?.token
+function getCurrency(currencyId: string, chainId: SupportedChainId, tokens: ChainTokenMap): Currency | undefined {
+  return currencyId === 'ETH' ? nativeOnChain(chainId) : tokens[chainId]?.[currencyId]
 }
 
 function buildCurrencyDescriptor(
@@ -46,7 +46,7 @@ function buildCurrencyDescriptor(
 function parseSwap(
   swap: ExactInputSwapTransactionInfo | ExactOutputSwapTransactionInfo,
   chainId: SupportedChainId,
-  tokens: TokenAddressMap
+  tokens: ChainTokenMap
 ): Partial<Activity> {
   const tokenIn = getCurrency(swap.inputCurrencyId, chainId, tokens)
   const tokenOut = getCurrency(swap.outputCurrencyId, chainId, tokens)
@@ -76,7 +76,7 @@ function parseWrap(wrap: WrapTransactionInfo, chainId: SupportedChainId, status:
 function parseApproval(
   approval: ApproveTransactionInfo,
   chainId: SupportedChainId,
-  tokens: TokenAddressMap
+  tokens: ChainTokenMap
 ): Partial<Activity> {
   // TODO: Add 'amount' approved to ApproveTransactionInfo so we can distinguish between revoke and approve
   const currency = getCurrency(approval.tokenAddress, chainId, tokens)
@@ -91,7 +91,7 @@ type GenericLPInfo = Omit<
   AddLiquidityV3PoolTransactionInfo | RemoveLiquidityV3TransactionInfo | AddLiquidityV2PoolTransactionInfo,
   'type'
 >
-function parseLP(lp: GenericLPInfo, chainId: SupportedChainId, tokens: TokenAddressMap): Partial<Activity> {
+function parseLP(lp: GenericLPInfo, chainId: SupportedChainId, tokens: ChainTokenMap): Partial<Activity> {
   const baseCurrency = getCurrency(lp.baseCurrencyId, chainId, tokens)
   const quoteCurrency = getCurrency(lp.quoteCurrencyId, chainId, tokens)
   const [baseRaw, quoteRaw] = [lp.expectedAmountBaseRaw, lp.expectedAmountQuoteRaw]
@@ -103,7 +103,7 @@ function parseLP(lp: GenericLPInfo, chainId: SupportedChainId, tokens: TokenAddr
 function parseCollectFees(
   collect: CollectFeesTransactionInfo,
   chainId: SupportedChainId,
-  tokens: TokenAddressMap
+  tokens: ChainTokenMap
 ): Partial<Activity> {
   // Adapts CollectFeesTransactionInfo to generic LP type
   const {
@@ -118,7 +118,7 @@ function parseCollectFees(
 function parseMigrateCreateV3(
   lp: MigrateV2LiquidityToV3TransactionInfo | CreateV3PoolTransactionInfo,
   chainId: SupportedChainId,
-  tokens: TokenAddressMap
+  tokens: ChainTokenMap
 ): Partial<Activity> {
   const baseCurrency = getCurrency(lp.baseCurrencyId, chainId, tokens)
   const baseSymbol = baseCurrency?.symbol ?? t`Unknown`
@@ -132,7 +132,7 @@ function parseMigrateCreateV3(
 export function parseLocalActivity(
   details: TransactionDetails,
   chainId: SupportedChainId,
-  tokens: TokenAddressMap
+  tokens: ChainTokenMap
 ): Activity | undefined {
   try {
     const status = !details.receipt
@@ -188,7 +188,7 @@ export function parseLocalActivity(
 
 export function useLocalActivities(account: string): ActivityMap {
   const allTransactions = useMultichainTransactions()
-  const tokens = useCombinedActiveList()
+  const tokens = useAllTokensMultichain()
 
   return useMemo(() => {
     const activityByHash: ActivityMap = {}

src/components/AccountDrawer/MiniPortfolio/Pools/cache.ts

@@ -126,7 +126,7 @@ export function useGetCachedTokens(chains: SupportedChainId[]): TokenGetterFn {
       const local: { [address: string]: Token | undefined } = {}
       const missing = new Set<string>()
       addresses.forEach((address) => {
-        const cached = tokenCache.get(chainId, address) ?? allTokens[chainId][address]?.token
+        const cached = tokenCache.get(chainId, address) ?? allTokens[chainId]?.[address]
         cached ? (local[address] = cached) : missing.add(address)
       })
 

src/components/Popups/TransactionPopup.tsx

@@ -3,8 +3,8 @@ import { parseLocalActivity } from 'components/AccountDrawer/MiniPortfolio/Activ
 import { PortfolioLogo } from 'components/AccountDrawer/MiniPortfolio/PortfolioLogo'
 import PortfolioRow from 'components/AccountDrawer/MiniPortfolio/PortfolioRow'
 import Column from 'components/Column'
+import { useAllTokensMultichain } from 'hooks/Tokens'
 import useENSName from 'hooks/useENSName'
-import { useCombinedActiveList } from 'state/lists/hooks'
 import { useTransaction } from 'state/transactions/hooks'
 import { TransactionDetails } from 'state/transactions/types'
 import styled from 'styled-components/macro'
@@ -19,7 +19,7 @@ const Descriptor = styled(ThemedText.BodySmall)`
 
 function TransactionPopupContent({ tx, chainId }: { tx: TransactionDetails; chainId: number }) {
   const success = tx.receipt?.status === 1
-  const tokens = useCombinedActiveList()
+  const tokens = useAllTokensMultichain()
   const activity = parseLocalActivity(tx, chainId, tokens)
   const { ENSName } = useENSName(activity?.otherAccount)
 

src/hooks/Tokens.test.ts

+import { SupportedChainId as MockSupportedChainId } from 'constants/chains'
+import {
+  DAI as MockDAI,
+  USDC_MAINNET as MockUSDC_MAINNET,
+  USDC_OPTIMISM as MockUSDC_OPTIMISM,
+  USDT as MockUSDT,
+  WETH_POLYGON as MockWETH_POLYGON,
+} from 'constants/tokens'
+import { renderHook } from 'test-utils/render'
+
+import { useAllTokensMultichain } from './Tokens'
+
+jest.mock('../state/lists/hooks.ts', () => {
+  return {
+    useCombinedTokenMapFromUrls: () => ({
+      [MockSupportedChainId.MAINNET]: {
+        [MockDAI.address]: { token: MockDAI },
+        [MockUSDC_MAINNET.address]: { token: MockUSDC_MAINNET },
+      },
+      [MockSupportedChainId.POLYGON]: {
+        [MockWETH_POLYGON.address]: { token: MockWETH_POLYGON },
+      },
+    }),
+  }
+})
+
+jest.mock('../state/hooks.ts', () => {
+  return {
+    useAppSelector: () => ({
+      [MockSupportedChainId.MAINNET]: {
+        [MockDAI.address]: MockDAI,
+        [MockUSDT.address]: MockUSDT,
+      },
+      [MockSupportedChainId.OPTIMISM]: {
+        [MockUSDC_OPTIMISM.address]: MockUSDC_OPTIMISM,
+      },
+    }),
+  }
+})
+
+describe('useAllTokensMultichain', () => {
+  it('should return multi-chain tokens from lists and userAddedTokens', () => {
+    const { result } = renderHook(() => useAllTokensMultichain())
+
+    expect(result.current).toStrictEqual({
+      [MockSupportedChainId.MAINNET]: {
+        [MockDAI.address]: MockDAI,
+        [MockUSDC_MAINNET.address]: MockUSDC_MAINNET,
+        [MockUSDT.address]: MockUSDT,
+      },
+      [MockSupportedChainId.POLYGON]: {
+        [MockWETH_POLYGON.address]: MockWETH_POLYGON,
+      },
+      [MockSupportedChainId.OPTIMISM]: {
+        [MockUSDC_OPTIMISM.address]: MockUSDC_OPTIMISM,
+      },
+    })
+  })
+})

src/hooks/Tokens.ts

@@ -5,13 +5,15 @@ import { SupportedChainId } from 'constants/chains'
 import { DEFAULT_INACTIVE_LIST_URLS, DEFAULT_LIST_OF_LISTS } from 'constants/lists'
 import { useCurrencyFromMap, useTokenFromMapOrNetwork } from 'lib/hooks/useCurrency'
 import { getTokenFilter } from 'lib/hooks/useTokenList/filtering'
+import { TokenAddressMap } from 'lib/hooks/useTokenList/utils'
 import { useMemo } from 'react'
+import { useAppSelector } from 'state/hooks'
 import { isL2ChainId } from 'utils/chains'
 
 import { useAllLists, useCombinedActiveList, useCombinedTokenMapFromUrls } from '../state/lists/hooks'
 import { WrappedTokenInfo } from '../state/lists/wrappedTokenInfo'
-import { useUserAddedTokens, useUserAddedTokensOnChain } from '../state/user/hooks'
-import { TokenAddressMap, useUnsupportedTokenList } from './../state/lists/hooks'
+import { deserializeToken, useUserAddedTokens, useUserAddedTokensOnChain } from '../state/user/hooks'
+import { useUnsupportedTokenList } from './../state/lists/hooks'
 
 type Maybe<T> = T | null | undefined
 
@@ -28,11 +30,41 @@ function useTokensFromMap(tokenMap: TokenAddressMap, chainId: Maybe<SupportedCha
   }, [chainId, tokenMap])
 }
 
-export function useAllTokensMultichain(): TokenAddressMap {
-  return useCombinedTokenMapFromUrls(DEFAULT_LIST_OF_LISTS)
+// TODO(INFRA-164): after disallowing unchecked index access, refactor ChainTokenMap to not use ?'s
+export type ChainTokenMap = { [chainId in number]?: { [address in string]?: Token } }
+/** Returns tokens from all token lists on all chains, combined with user added tokens */
+export function useAllTokensMultichain(): ChainTokenMap {
+  const allTokensFromLists = useCombinedTokenMapFromUrls(DEFAULT_LIST_OF_LISTS)
+  const userAddedTokensMap = useAppSelector(({ user: { tokens } }) => tokens)
+
+  return useMemo(() => {
+    const chainTokenMap: ChainTokenMap = {}
+
+    if (userAddedTokensMap) {
+      Object.keys(userAddedTokensMap).forEach((key) => {
+        const chainId = Number(key)
+        const tokenMap = {} as { [address in string]?: Token }
+        Object.values(userAddedTokensMap[chainId]).forEach((serializedToken) => {
+          tokenMap[serializedToken.address] = deserializeToken(serializedToken)
+        })
+        chainTokenMap[chainId] = tokenMap
+      })
+    }
+
+    Object.keys(allTokensFromLists).forEach((key) => {
+      const chainId = Number(key)
+      const tokenMap = chainTokenMap[chainId] ?? {}
+      Object.values(allTokensFromLists[chainId]).forEach(({ token }) => {
+        tokenMap[token.address] = token
+      })
+      chainTokenMap[chainId] = tokenMap
+    })
+
+    return chainTokenMap
+  }, [allTokensFromLists, userAddedTokensMap])
 }
 
-// Returns all tokens from the default list + user added tokens
+/** Returns all tokens from the default list + user added tokens */
 export function useDefaultActiveTokens(chainId: Maybe<SupportedChainId>): { [address: string]: Token } {
   const defaultListTokens = useCombinedActiveList()
   const tokensFromMap = useTokensFromMap(defaultListTokens, chainId)

src/lib/hooks/useTokenList/utils.ts

@@ -2,20 +2,21 @@ import { TokenInfo, TokenList } from '@uniswap/token-lists'
 import { WrappedTokenInfo } from 'state/lists/wrappedTokenInfo'
 
 type TokenMap = Readonly<{ [tokenAddress: string]: { token: WrappedTokenInfo; list?: TokenList } }>
-export type ChainTokenMap = Readonly<{ [chainId: number]: TokenMap }>
+// TODO(INFRA-164): replace usage of the misnomered TokenAddressMap w/ ChainTokenMap from src/hooks/Tokens.ts
+export type TokenAddressMap = Readonly<{ [chainId: number]: TokenMap }>
 
 type Mutable<T> = {
   -readonly [P in keyof T]: Mutable<T[P]>
 }
 
-const mapCache = typeof WeakMap !== 'undefined' ? new WeakMap<TokenList | TokenInfo[], ChainTokenMap>() : null
+const mapCache = typeof WeakMap !== 'undefined' ? new WeakMap<TokenList | TokenInfo[], TokenAddressMap>() : null
 
-export function tokensToChainTokenMap(tokens: TokenList | TokenInfo[]): ChainTokenMap {
+export function tokensToChainTokenMap(tokens: TokenList | TokenInfo[]): TokenAddressMap {
   const cached = mapCache?.get(tokens)
   if (cached) return cached
 
   const [list, infos] = Array.isArray(tokens) ? [undefined, tokens] : [tokens, tokens.tokens]
-  const map = infos.reduce<Mutable<ChainTokenMap>>((map, info) => {
+  const map = infos.reduce<Mutable<TokenAddressMap>>((map, info) => {
     try {
       const token = new WrappedTokenInfo(info, list)
       if (map[token.chainId]?.[token.address] !== undefined) {
@@ -30,7 +31,7 @@ export function tokensToChainTokenMap(tokens: TokenList | TokenInfo[]): ChainTok
     } catch {
       return map
     }
-  }, {}) as ChainTokenMap
+  }, {}) as TokenAddressMap
   mapCache?.set(tokens, map)
   return map
 }

src/state/lists/hooks.ts

-import { ChainTokenMap, tokensToChainTokenMap } from 'lib/hooks/useTokenList/utils'
+import { TokenAddressMap, tokensToChainTokenMap } from 'lib/hooks/useTokenList/utils'
 import { useMemo } from 'react'
 import { useAppSelector } from 'state/hooks'
 import sortByListPriority from 'utils/listSort'
@@ -7,8 +7,6 @@ import BROKEN_LIST from '../../constants/tokenLists/broken.tokenlist.json'
 import { AppState } from '../types'
 import { DEFAULT_ACTIVE_LIST_URLS, UNSUPPORTED_LIST_URLS } from './../../constants/lists'
 
-export type TokenAddressMap = ChainTokenMap
-
 type Mutable<T> = {
   -readonly [P in keyof T]: Mutable<T[P]>
 }