Merge pull request #65 from nulllpc/npc/fetch-balance-for-multi-indices

Support fetching token balances for multiple account indices

Author: jonathunne

src/hooks/useBalance.ts

@@ -155,6 +155,48 @@ export const balanceQueryKeys = {
       QUERY_KEY_TAGS.TOKEN,
       assetId,
     ] as const,
+  byTokensAndIndices: (walletId: string, accountIndices: number[], assetIds: string[]) => {
+    const sortedIndices = [...accountIndices].sort((a, b) => a - b)
+    const sortedAssetIds = [...assetIds].sort()
+
+    return [
+      QUERY_KEY_TAGS.BALANCES,
+      QUERY_KEY_TAGS.WALLET,
+      walletId,
+      sortedIndices,
+      sortedAssetIds,
+      'multi'
+    ]
+  }
+}
+
+async function promisePool<T>(
+  promiseFns: (() => Promise<T>)[],
+  concurrency = 2
+): Promise<T[]> {
+  const results: T[] = new Array(promiseFns.length)
+  let index = 0
+
+  const runNext = async (): Promise<void> => {
+    if (index >= promiseFns.length) {
+      return
+    }
+
+    const currentIndex = index++
+    const promiseFn = promiseFns[currentIndex]!
+
+    results[currentIndex] = await promiseFn()
+
+    await runNext()
+  }
+
+  const workers = Array(Math.min(concurrency, promiseFns.length))
+    .fill(null)
+    .map(() => runNext())
+
+  await Promise.all(workers)
+
+  return results
 }
 
 /**
@@ -283,7 +325,7 @@ export function useBalancesForWallet(
     error: addressesError,
   } = useMultiAddressLoader({
     networks: uniqueNetworks,
-    accountIndex,
+    accountIndices: [accountIndex],
     enabled: options?.enabled,
   });
 
@@ -332,6 +374,97 @@ export function useBalancesForWallet(
   return { ...query, isLoading, error: error as Error | null };
 }
 
+export type UseBalancesForWalletsResult = Omit<
+  UseQueryResult<BalanceFetchResult[], Error>,
+  'isLoading' | 'error'
+> & {
+  isLoading: boolean;
+  error: Error | null;
+};
+
+export function useBalancesForWallets(
+  accountIndices: number[],
+  assetConfigs: IAsset[],
+  options?: BalanceQueryOptions & { concurrencyLimit?: number }
+): UseBalancesForWalletsResult {
+  const uniqueNetworks = useMemo(
+    () => [...new Set(assetConfigs.map((asset) => asset.getNetwork()))],
+    [assetConfigs],
+  );
+
+  const {
+    isLoading: areAddressesLoading,
+    error: addressesError,
+  } = useMultiAddressLoader({
+    networks: uniqueNetworks,
+    accountIndices: accountIndices,
+    enabled: options?.enabled,
+  });
+
+  const walletId = getWalletStore()((state) => state.activeWalletId);
+
+  const initialData: BalanceFetchResult[] | undefined = (() => {
+    if (!walletId || assetConfigs.length === 0) {
+      return undefined;
+    }
+
+    return accountIndices.flatMap((accountIndex) => {
+      return assetConfigs.map((asset) => {
+        const balance = BalanceService.getBalance(
+          accountIndex,
+          asset.getNetwork(),
+          asset.getId(),
+          walletId,
+        );
+  
+        return {
+          success: true,
+          network: asset.getNetwork(),
+          accountIndex: accountIndex,
+          assetId: asset.getId(),
+          balance,
+        };
+      })
+    });
+  })();
+
+  const query = useQuery({
+    queryKey: balanceQueryKeys.byTokensAndIndices(walletId || '', accountIndices, assetConfigs.map((asset) => asset.getId())),
+    queryFn: async () => {
+      const concurrency = options?.concurrencyLimit ?? 2;
+      const tasks = accountIndices.map((accountIndex) => async (): Promise<BalanceFetchResult[]> => {
+        try {
+          return await fetchBalances(accountIndex, assetConfigs);
+        } catch (error) {
+          return assetConfigs.map((asset) => ({
+            success: false,
+            network: asset.getNetwork(),
+            accountIndex,
+            assetId: asset.getId(),
+            balance: null,
+            error: error instanceof Error ? error.message : String(error),
+          }));
+        }
+      });
+      const results = await promisePool(tasks, concurrency);
+      return results.flat();
+    },
+    enabled: isQueryEnabled(
+      options?.enabled,
+      !!walletId && !areAddressesLoading && assetConfigs.length > 0,
+    ),
+    refetchInterval: options?.refetchInterval,
+    staleTime: options?.staleTime ?? DEFAULT_QUERY_STALE_TIME_MS,
+    gcTime: DEFAULT_QUERY_GC_TIME_MS,
+    initialData,
+  });
+
+  const isLoading = areAddressesLoading || query.isLoading;
+  const error = addressesError || query.error;
+
+  return { ...query, isLoading, error: error as Error | null };
+}
+
 /**
  * Invalidate balance queries based on refresh type
  */

src/hooks/useMultiAddressLoader.ts

@@ -13,18 +13,19 @@
 // limitations under the License.
 
 import { useState, useEffect } from 'react';
-import { AccountService } from '../services/accountService';
 import { getWalletStore } from '../store/walletStore';
 import { logError } from '../utils/logger';
+import { AddressService } from '../services/addressService';
 
 export interface AddressResult {
   network: string;
-  address: string;
+  address: string | null;
+  accountIndex: number;
 }
 
 interface UseMultiAddressLoaderParams {
   networks: string[];
-  accountIndex: number;
+  accountIndices: number[];
   enabled?: boolean;
 }
 
@@ -42,7 +43,7 @@ interface UseMultiAddressLoaderResult {
  */
 export function useMultiAddressLoader({
   networks,
-  accountIndex,
+  accountIndices,
   enabled = true,
 }: UseMultiAddressLoaderParams): UseMultiAddressLoaderResult {
   const [addresses, setAddresses] = useState<AddressResult[] | null>(null);
@@ -51,13 +52,18 @@ export function useMultiAddressLoader({
   const activeWalletId = getWalletStore()((state) => state.activeWalletId);
 
   const networksKey = JSON.stringify([...networks].sort());
+  const activeIndices = JSON.stringify([...accountIndices].sort());
 
   useEffect(() => {
+    let isStale = false;
+
     const loadAddresses = async () => {
       if (!enabled || networks.length === 0 || !activeWalletId) {
-        setIsLoading(false);
-        setError(null);
-        setAddresses(null);
+        if (!isStale) {
+          setIsLoading(false);
+          setError(null);
+          setAddresses(null);
+        }
         return;
       }
 
@@ -66,40 +72,46 @@ export function useMultiAddressLoader({
       setAddresses(null);
 
       try {
-        const uniqueNetworks = [...new Set(networks)];
-        const addressPromises = uniqueNetworks.map((network) =>
-          AccountService.callAccountMethod<'getAddress'>(
-            network,
-            accountIndex,
-            'getAddress'
-          ),
-        );
-
-        const loadedAddresses = await Promise.all(addressPromises);
-
-        const addressMap = new Map<string, string>();
-        uniqueNetworks.forEach((network, index) => {
-          addressMap.set(network, loadedAddresses[index] as string);
-        });
-
-        const finalAddresses = networks.map((network) => ({
-          network,
-          address: addressMap.get(network)!,
-        }));
+        const addressesResult = await AddressService.getAddresses(accountIndices, networks)
+
+        if (isStale) return;
+
+        const finalAddresses: AddressResult[] = addressesResult.map((addressInfo) => {
+          if (addressInfo.success === true) {
+            return {
+              network: addressInfo.network,
+              accountIndex: addressInfo.accountIndex,
+              address: addressInfo.address
+            }
+          } else {
+            return {
+              network: addressInfo.network,
+              accountIndex: addressInfo.accountIndex,
+              address: null
+            }
+          }
+        })
 
         setAddresses(finalAddresses);
       } catch (e) {
+        if (isStale) return;
         const err = e instanceof Error ? e : new Error('Failed to load addresses');
         logError('useMultiAddressLoader failed:', err);
         setError(err);
       } finally {
-        setIsLoading(false);
+        if (!isStale) {
+          setIsLoading(false);
+        }
       }
     };
 
     loadAddresses();
+
+    return () => {
+      isStale = true;
+    };
     // eslint-disable-next-line react-hooks/exhaustive-deps
-  }, [networksKey, accountIndex, enabled, activeWalletId]);
+  }, [networksKey, activeIndices, enabled, activeWalletId]);
 
   return { addresses, isLoading, error };
 }

src/index.ts

@@ -48,6 +48,7 @@ export type { UseWalletManagerResult, WalletInfo } from './hooks/useWalletManage
 export {
   useBalance,
   useBalancesForWallet,
+  useBalancesForWallets,
   useRefreshBalance,
   balanceQueryKeys,
 } from './hooks/useBalance'