From ff158041a8e5932b6a2d2a52513600132a831831 Mon Sep 17 00:00:00 2001
From: "moxey.eth" <jakemoxey@gmail.com>
Date: Tue, 10 Oct 2023 20:00:14 +1100
Subject: [PATCH] refactor: account tokens

---
 src/hooks/useAccountTokens.ts | 39 +++++++++++++++----
 src/zustand/tokens.ts         | 73 ++++++++++++++++++++++++-----------
 2 files changed, 82 insertions(+), 30 deletions(-)

diff --git a/src/hooks/useAccountTokens.ts b/src/hooks/useAccountTokens.ts
index 615e50c..224cc7c 100644
--- a/src/hooks/useAccountTokens.ts
+++ b/src/hooks/useAccountTokens.ts
@@ -1,4 +1,5 @@
 import { queryOptions, useQuery } from '@tanstack/react-query'
+import { useCallback, useMemo } from 'react'
 import { type Address, stringify } from 'viem'
 
 import { getAccountTokens } from '~/actions'
@@ -6,7 +7,7 @@ import { createQueryKey } from '~/react-query'
 import type { Client } from '~/viem'
 import { useNetworkStore, useTokensStore } from '~/zustand'
 
-import { useCallback, useMemo } from 'react'
+import { getTokensKey } from '~/zustand/tokens'
 import { useClient } from './useClient'
 
 type UseAccountTokensParameters = {
@@ -35,7 +36,11 @@ export function useAccountTokensQueryOptions(args: UseAccountTokensParameters) {
         fromBlock: network.forkBlockNumber,
         toBlock: 'latest',
       })
-      syncTokens(tokens, address)
+      syncTokens({
+        accountAddress: address,
+        tokenAddresses: tokens,
+        rpcUrl: network.rpcUrl,
+      })
       return tokens
     },
   })
@@ -44,25 +49,43 @@ export function useAccountTokensQueryOptions(args: UseAccountTokensParameters) {
 export function useAccountTokens(args: UseAccountTokensParameters) {
   const queryOptions = useAccountTokensQueryOptions(args)
   const tokensStore = useTokensStore()
+  const { network } = useNetworkStore()
 
   const addToken = useCallback(
     (address: Address) =>
-      args.address ? tokensStore.addToken(address, args.address) : undefined,
-    [args.address, tokensStore.addToken],
+      args.address
+        ? tokensStore.addToken({
+            accountAddress: args.address,
+            rpcUrl: network.rpcUrl,
+            tokenAddress: address,
+          })
+        : undefined,
+    [args.address, network.rpcUrl, tokensStore.addToken],
   )
   const removeToken = useCallback(
     (address: Address) =>
-      args.address ? tokensStore.removeToken(address, args.address) : undefined,
-    [args.address, tokensStore.removeToken],
+      args.address
+        ? tokensStore.removeToken({
+            accountAddress: args.address,
+            rpcUrl: network.rpcUrl,
+            tokenAddress: address,
+          })
+        : undefined,
+    [args.address, network.rpcUrl, tokensStore.removeToken],
   )
   const tokens = useMemo(
     () =>
       args.address
-        ? tokensStore.tokens[args.address]
+        ? tokensStore.tokens[
+            getTokensKey({
+              accountAddress: args.address,
+              rpcUrl: network.rpcUrl,
+            })
+          ]
             ?.map((token) => (!token.removed ? token.address : undefined))
             .filter(Boolean)
         : [],
-    [args.address, tokensStore.tokens],
+    [args.address, network.rpcUrl, tokensStore.tokens],
   )
 
   return Object.assign(useQuery(queryOptions), {
diff --git a/src/zustand/tokens.ts b/src/zustand/tokens.ts
index 7c5ef8d..aca2bc7 100644
--- a/src/zustand/tokens.ts
+++ b/src/zustand/tokens.ts
@@ -8,26 +8,51 @@ type Token = {
   removed: boolean
 }
 
+type TokensKey = string
+
 export type TokensState = {
-  tokens: Record<Address, Token[]>
+  tokens: Record<TokensKey, Token[]>
 }
 export type TokensActions = {
-  addToken: (tokenAddress: Address, address: Address) => void
-  removeToken: (tokenAddress: Address, address: Address) => void
-  syncTokens: (tokenAddresses: Address[], address: Address) => void
+  addToken: (parameters: {
+    accountAddress: Address
+    tokenAddress: Address
+    rpcUrl: string
+  }) => void
+  removeToken: (parameters: {
+    accountAddress: Address
+    tokenAddress: Address
+    rpcUrl: string
+  }) => void
+  syncTokens: (parameters: {
+    accountAddress: Address
+    tokenAddresses: Address[]
+    rpcUrl: string
+  }) => void
 }
 export type TokensStore = TokensState & TokensActions
 
+export function getTokensKey(args: {
+  accountAddress: Address
+  rpcUrl: string
+}): TokensKey {
+  const { accountAddress, rpcUrl } = args
+  return `${rpcUrl}-${accountAddress}`.replace(/\./g, '-')
+}
+
 export const tokensStore = createStore<TokensStore>(
   (set) => ({
     tokens: {},
-    addToken(tokenAddress: Address, accountAddress: Address) {
+    addToken(args) {
+      const { accountAddress, tokenAddress, rpcUrl } = args
+      const key = getTokensKey({ accountAddress, rpcUrl })
+
       set((state) => {
         const tokens = { ...state.tokens }
-        tokens[accountAddress] = uniqBy(
+        tokens[key] = uniqBy(
           [
             { address: tokenAddress, removed: false },
-            ...(state.tokens[accountAddress] || []),
+            ...(state.tokens[key] || []),
           ],
           (x) => x.address,
         )
@@ -38,36 +63,40 @@ export const tokensStore = createStore<TokensStore>(
         }
       })
     },
-    removeToken(tokenAddress: Address, accountAddress: Address) {
+    removeToken(args) {
+      const { accountAddress, tokenAddress, rpcUrl } = args
+      const key = getTokensKey({ accountAddress, rpcUrl })
+
       set((state) => {
         const tokens = { ...state.tokens }
-        tokens[accountAddress] = (state.tokens[accountAddress] || []).map(
-          (token) => {
-            if (token.address === tokenAddress)
-              return {
-                ...token,
-                removed: true,
-              }
-            return token
-          },
-        )
+        tokens[key] = (state.tokens[key] || []).map((token) => {
+          if (token.address === tokenAddress)
+            return {
+              ...token,
+              removed: true,
+            }
+          return token
+        })
         return {
           ...state,
           tokens,
         }
       })
     },
-    syncTokens(tokenAddresses: Address[], accountAddress: Address) {
+    syncTokens(args) {
+      const { accountAddress, tokenAddresses, rpcUrl } = args
+      const key = getTokensKey({ accountAddress, rpcUrl })
+
       set((state) => {
         const tokens = { ...state.tokens }
 
         for (const tokenAddress of tokenAddresses) {
-          const exists = (tokens[accountAddress] || []).some(
+          const exists = (tokens[key] || []).some(
             (x) => x.address === tokenAddress,
           )
           if (!exists)
-            tokens[accountAddress] = [
-              ...(tokens[accountAddress] || []),
+            tokens[key] = [
+              ...(tokens[key] || []),
               {
                 address: tokenAddress,
                 removed: false,