From 4271eb9267a59dcc3b6ee2cd9992a1f5c8692d8b Mon Sep 17 00:00:00 2001
From: Tom Quirk <12551741+tomquirk@users.noreply.github.com>
Date: Mon, 19 Feb 2024 14:34:24 +1000
Subject: [PATCH] fix: fix migrated tokens showing up in legacy balance (race
 condition)

---
 .../useTotalLegacyTokenBalance.ts              | 18 ++++++++++++------
 .../useV2UnclaimedBalanceForV3Token.ts         | 14 +++++++++-----
 .../contractReader/useUnclaimedTokenBalance.ts |  8 +-------
 3 files changed, 22 insertions(+), 18 deletions(-)

diff --git a/src/hooks/JBV3Token/contractReader/useTotalLegacyTokenBalance.ts b/src/hooks/JBV3Token/contractReader/useTotalLegacyTokenBalance.ts
index f9751edc98..66d8ef0d1b 100644
--- a/src/hooks/JBV3Token/contractReader/useTotalLegacyTokenBalance.ts
+++ b/src/hooks/JBV3Token/contractReader/useTotalLegacyTokenBalance.ts
@@ -13,15 +13,21 @@ export function useTotalLegacyTokenBalance({
   const v1UnclaimedBalance = useV1UnclaimedBalanceForV3Token()
   const v1ClaimedBalance = v1TotalBalance?.sub(v1UnclaimedBalance)
 
-  const { data: v2TotalBalance } = useV2TotalBalance({ projectId })
-  const { data: v2UnclaimedBalance } = useV2UnclaimedBalanceForV3Token({
-    projectId,
-  })
-  const v2ClaimedBalance =
-    v2TotalBalance?.sub(v2UnclaimedBalance ?? 0) ?? BigNumber.from(0)
+  const { data: v2TotalBalance, loading: v2BalanceLoading } = useV2TotalBalance(
+    { projectId },
+  )
+  const { data: v2UnclaimedBalance, loading: unclaimedBalanceLoading } =
+    useV2UnclaimedBalanceForV3Token({
+      projectId,
+    })
+
+  if (v2BalanceLoading || unclaimedBalanceLoading) return BigNumber.from(0)
 
   const totalLegacyTokenBalance = v1TotalBalance?.add(v2TotalBalance ?? 0)
 
+  const v2ClaimedBalance =
+    v2TotalBalance?.sub(v2UnclaimedBalance ?? 0) ?? BigNumber.from(0)
+
   return {
     v1TotalBalance,
     v1UnclaimedBalance,
diff --git a/src/hooks/JBV3Token/contractReader/useV2UnclaimedBalanceForV3Token.ts b/src/hooks/JBV3Token/contractReader/useV2UnclaimedBalanceForV3Token.ts
index 50079a24de..50a486f1a6 100644
--- a/src/hooks/JBV3Token/contractReader/useV2UnclaimedBalanceForV3Token.ts
+++ b/src/hooks/JBV3Token/contractReader/useV2UnclaimedBalanceForV3Token.ts
@@ -1,5 +1,6 @@
+import { BigNumber } from 'ethers'
 import { useWallet } from 'hooks/Wallet'
-import { useUnclaimedTokenBalance } from 'hooks/v2v3/contractReader/useUnclaimedTokenBalance'
+import useV2ContractReader from 'hooks/v2v3/contractReader/useV2ContractReader'
 import { useJBTokenStoreForV3Token } from '../contracts/useJBTokenStoreForV3Token'
 
 export function useV2UnclaimedBalanceForV3Token({
@@ -10,9 +11,12 @@ export function useV2UnclaimedBalanceForV3Token({
   const { userAddress } = useWallet()
   const v2TokenStoreContract = useJBTokenStoreForV3Token()
 
-  return useUnclaimedTokenBalance({
-    userAddress,
-    projectId,
-    JBTokenStore: v2TokenStoreContract,
+  return useV2ContractReader<BigNumber>({
+    contract: v2TokenStoreContract,
+    functionName: 'unclaimedBalanceOf',
+    args:
+      userAddress && projectId && v2TokenStoreContract
+        ? [userAddress, projectId]
+        : null,
   })
 }
diff --git a/src/hooks/v2v3/contractReader/useUnclaimedTokenBalance.ts b/src/hooks/v2v3/contractReader/useUnclaimedTokenBalance.ts
index eeb93f4604..d0a2b426e9 100644
--- a/src/hooks/v2v3/contractReader/useUnclaimedTokenBalance.ts
+++ b/src/hooks/v2v3/contractReader/useUnclaimedTokenBalance.ts
@@ -1,7 +1,4 @@
 import { BigNumber } from 'ethers'
-import { bigNumbersDiff } from 'utils/bigNumbers'
-
-import { Contract } from 'ethers'
 import { V2V3ContractName } from 'models/v2v3/contracts'
 import useContractReader from './useV2ContractReader'
 
@@ -9,19 +6,16 @@ import useContractReader from './useV2ContractReader'
 export function useUnclaimedTokenBalance({
   projectId,
   userAddress,
-  JBTokenStore,
 }: {
   projectId: number | undefined
   userAddress: string | undefined
-  JBTokenStore?: Contract
 }) {
   return useContractReader<BigNumber>({
-    contract: JBTokenStore ?? V2V3ContractName.JBTokenStore,
+    contract: V2V3ContractName.JBTokenStore,
     functionName: 'unclaimedBalanceOf',
     args:
       userAddress && projectId
         ? [userAddress, BigNumber.from(projectId).toHexString()]
         : null,
-    valueDidChange: bigNumbersDiff,
   })
 }