From b4737756fd90f38ab9b8137a90161e269ddb6c45 Mon Sep 17 00:00:00 2001
From: Spacebean <wndgud00@gmail.com>
Date: Wed, 26 Jun 2024 11:04:44 -0600
Subject: [PATCH] feat: fix too many token balances to load error

---
 .../dex-ui/src/tokens/useAllTokenBalance.tsx  | 67 ++++++++++++-------
 projects/dex-ui/src/utils/query/queryKeys.ts  |  6 +-
 2 files changed, 47 insertions(+), 26 deletions(-)

diff --git a/projects/dex-ui/src/tokens/useAllTokenBalance.tsx b/projects/dex-ui/src/tokens/useAllTokenBalance.tsx
index cbdc43e21c..91624a98e0 100644
--- a/projects/dex-ui/src/tokens/useAllTokenBalance.tsx
+++ b/projects/dex-ui/src/tokens/useAllTokenBalance.tsx
@@ -7,6 +7,8 @@ import { useAccount } from "wagmi";
 import { useTokens } from "./TokenProvider";
 import { Log } from "src/utils/logger";
 import { config } from "src/utils/wagmi/config";
+import { ContractFunctionParameters } from "viem";
+import { queryKeys } from "src/utils/query/queryKeys";
 
 const TokenBalanceABI = [
   {
@@ -20,63 +22,78 @@ const TokenBalanceABI = [
   }
 ] as const;
 
+const MAX_PER_CALL = 20;
+
 export const useAllTokensBalance = () => {
   const tokens = useTokens();
   const { address } = useAccount();
   const queryClient = useQueryClient();
 
-  // Remove ETH from this list, manually get the balance below
   const tokensToLoad = Object.values(tokens).filter((t) => t.symbol !== "ETH");
-  if (tokensToLoad.length > 20) throw new Error("Too many tokens to load balances. Fix me");
 
   const calls = useMemo(() => {
-    const contractCalls: any[] = [];
-    Log.module("app").debug(`Fetching token balances for ${tokensToLoad.length} tokens, for address ${address}`);
-    for (const t of tokensToLoad) {
-      contractCalls.push({
-        address: t.address as `0x{string}`,
+    const contractCalls: ContractFunctionParameters[][] = [];
+    Log.module("app").debug(
+      `Fetching token balances for ${tokensToLoad.length} tokens, for address ${address}`
+    );
+
+    let callBucket: ContractFunctionParameters[] = [];
+
+    tokensToLoad.forEach((token, i) => {
+      callBucket.push({
+        address: token.address as `0x{string}`,
         abi: TokenBalanceABI,
         functionName: "balanceOf",
         args: [address]
       });
-    }
+
+      if (i % MAX_PER_CALL === MAX_PER_CALL - 1) {
+        contractCalls.push([...callBucket]);
+        callBucket = [];
+      }
+    });
     return contractCalls;
 
     // eslint-disable-next-line react-hooks/exhaustive-deps -- doing just tokensToLoad doesn't work and causes multiple calls
   }, [address, tokensToLoad.map((t) => t.symbol).join()]);
 
   const { data, isLoading, error, refetch, isFetching } = useQuery({
-    queryKey: ["token", "balance"],
-
+    queryKey: queryKeys.tokenBalancesAll,
     queryFn: async () => {
       if (!address) return {};
-      const res = (await multicall(config, {
-        contracts: calls,
-        allowFailure: false
-      })) as unknown as BigNumber[];
+
+      const ETH = tokens.ETH;
+
+      const [ethBalance, ...results] = await Promise.all([
+        ETH.getBalance(address),
+        ...(calls.map((calls) =>
+          multicall(config, { contracts: calls, allowFailure: false })
+        ) as unknown as BigNumber[])
+      ]);
+
+      const res = results.flat();
       const balances: Record<string, TokenValue> = {};
 
+      if (ethBalance) {
+        Log.module("app").debug(`ETH balance: `, ethBalance.toHuman());
+        queryClient.setQueryData(queryKeys.tokenBalance(ETH.symbol), { ETH: ethBalance });
+        balances.ETH = ethBalance;
+      }
+
       for (let i = 0; i < res.length; i++) {
         const value = res[i];
         const token = tokensToLoad[i];
         balances[token.symbol] = token.fromBlockchain(value);
 
         // set the balance in the query cache too
-        queryClient.setQueryData(["token", "balance", token.symbol], { [token.symbol]: balances[token.symbol] });
-
-      }
-
-      const ETH = tokens.ETH;
-      if (ETH) {
-        const ethBalance = await ETH.getBalance(address);
-        Log.module("app").debug(`ETH balance: `, ethBalance.toHuman());
-        queryClient.setQueryData(["token", "balance", "ETH"], { ETH: ethBalance });
-        balances.ETH = ethBalance;
+        queryClient.setQueryData(queryKeys.tokenBalance(token.symbol), {
+          [token.symbol]: balances[token.symbol]
+        });
       }
 
       return balances;
     },
-
+    enabled: !!address && !!tokensToLoad.length,
     staleTime: 1000 * 30,
     refetchInterval: 1000 * 30
   });
diff --git a/projects/dex-ui/src/utils/query/queryKeys.ts b/projects/dex-ui/src/utils/query/queryKeys.ts
index 2f7a4527b9..6d5d856869 100644
--- a/projects/dex-ui/src/utils/query/queryKeys.ts
+++ b/projects/dex-ui/src/utils/query/queryKeys.ts
@@ -18,5 +18,9 @@ export const queryKeys = {
   // prices
   tokenPricesAll: ["prices", "token"],
   tokenPrices: (address: string[]) => ["prices", "token", ...address],
-  lpTokenPrices: (addresses: string[]) => ["prices", "lp-token", ...addresses]
+  lpTokenPrices: (addresses: string[]) => ["prices", "lp-token", ...addresses],
+
+  // token balance
+  tokenBalancesAll: ["token", "balance"],
+  tokenBalance: (symbol: string) => ["token", "balance", symbol]
 } as const;