Fix token balances refetch on swap and join

balancer · Jul 11, 2024 · f4ec61e · f4ec61e
1 parent ef3f361
commit f4ec61e
Show file tree

Hide file tree

Showing 8 changed files with 72 additions and 40 deletions.
diff --git a/packages/nextjs/app/pools/_components/PoolActions.tsx b/packages/nextjs/app/pools/_components/PoolActions.tsx
@@ -3,7 +3,7 @@ import { AddLiquidityForm, RemoveLiquidityForm, SwapForm } from "./actions";
 import { useAccount } from "wagmi";
 import { ExclamationTriangleIcon } from "@heroicons/react/24/outline";
 import { useTokens } from "~~/hooks/balancer";
-import { type Pool } from "~~/hooks/balancer/types";
+import { type Pool, type TokenBalances } from "~~/hooks/balancer/types";
 import { type RefetchPool } from "~~/hooks/balancer/usePoolContract";
 import { useAccountBalance, useScaffoldContractWrite } from "~~/hooks/scaffold-eth";
 
@@ -12,12 +12,14 @@ type Action = "Swap" | "AddLiquidity" | "RemoveLiquidity";
 export interface PoolActionsProps {
   pool: Pool;
   refetchPool: RefetchPool;
+  tokenBalances: TokenBalances;
+  refetchTokenBalances: () => void;
 }
 
 /**
  * Allow user to swap, add liquidity, and remove liquidity from a pool
  */
-export const PoolActions: React.FC<PoolActionsProps> = ({ pool, refetchPool }) => {
+export const PoolActions: React.FC<{ pool: Pool; refetchPool: RefetchPool }> = ({ pool, refetchPool }) => {
   const [activeTab, setActiveTab] = useState<Action>("Swap");
 
   const { address } = useAccount();
@@ -26,12 +28,12 @@ export const PoolActions: React.FC<PoolActionsProps> = ({ pool, refetchPool }) =
   const tokens = pool.poolTokens.map(token => ({
     address: token.address as `0x${string}`,
     decimals: token.decimals,
-    rawAmount: 0n,
+    rawAmount: 0n, // Quirky solution cus useTokens expects type InputAmount[] cus originally built for AddLiquidityForm :D
   }));
 
   const { tokenBalances, refetchTokenBalances } = useTokens(tokens);
 
-  const userHasNoTokens = tokenBalances?.every(balance => balance === 0n);
+  const userHasNoTokens = Object.values(tokenBalances).every(balance => balance === 0n);
 
   const { writeAsync: mintToken1 } = useScaffoldContractWrite({
     contractName: "MockToken1",
@@ -46,9 +48,30 @@ export const PoolActions: React.FC<PoolActionsProps> = ({ pool, refetchPool }) =
   });
 
   const tabs = {
-    Swap: <SwapForm pool={pool} refetchPool={refetchPool} />,
-    AddLiquidity: <AddLiquidityForm pool={pool} refetchPool={refetchPool} />,
-    RemoveLiquidity: <RemoveLiquidityForm pool={pool} refetchPool={refetchPool} />,
+    Swap: (
+      <SwapForm
+        pool={pool}
+        refetchPool={refetchPool}
+        tokenBalances={tokenBalances}
+        refetchTokenBalances={refetchTokenBalances}
+      />
+    ),
+    AddLiquidity: (
+      <AddLiquidityForm
+        pool={pool}
+        refetchPool={refetchPool}
+        tokenBalances={tokenBalances}
+        refetchTokenBalances={refetchTokenBalances}
+      />
+    ),
+    RemoveLiquidity: (
+      <RemoveLiquidityForm
+        pool={pool}
+        refetchPool={refetchPool}
+        tokenBalances={tokenBalances}
+        refetchTokenBalances={refetchTokenBalances}
+      />
+    ),
   };
 
   return (

diff --git a/packages/nextjs/app/pools/_components/actions/AddLiquidityForm.tsx b/packages/nextjs/app/pools/_components/actions/AddLiquidityForm.tsx
@@ -17,7 +17,12 @@ import { formatToHuman } from "~~/utils/formatToHuman";
  * 3. Send transaction to add liquidity to the pool
  * 4. Display the transaction results to the user
  */
-export const AddLiquidityForm: React.FC<PoolActionsProps> = ({ pool, refetchPool }) => {
+export const AddLiquidityForm: React.FC<PoolActionsProps> = ({
+  pool,
+  refetchPool,
+  tokenBalances,
+  refetchTokenBalances,
+}) => {
   const initialTokenInputs = pool.poolTokens.map(token => ({
     address: token.address as `0x${string}`,
     decimals: token.decimals,
@@ -33,7 +38,7 @@ export const AddLiquidityForm: React.FC<PoolActionsProps> = ({ pool, refetchPool
   const [isAddingLiquidity, setIsAddingLiquidity] = useState(false);
   const [addLiquidityReceipt, setAddLiquidityReceipt] = useState<PoolActionReceipt>(null);
 
-  const { tokenAllowances, refetchTokenAllowances, tokenBalances } = useTokens(tokenInputs);
+  const { tokenAllowances, refetchTokenAllowances } = useTokens(tokenInputs);
   const { queryAddLiquidity, addLiquidity } = useAddLiquidity(pool, tokenInputs);
   const writeTx = useTransactor(); // scaffold hook for tx status toast notifications
   const publicClient = usePublicClient();
@@ -134,6 +139,7 @@ export const AddLiquidityForm: React.FC<PoolActionsProps> = ({ pool, refetchPool
       setIsAddingLiquidity(true);
       await addLiquidity();
       refetchTokenAllowances();
+      refetchTokenBalances();
       refetchPool();
     } catch (e) {
       console.error("error", e);
@@ -172,7 +178,7 @@ export const AddLiquidityForm: React.FC<PoolActionsProps> = ({ pool, refetchPool
             }
             onAmountChange={value => handleInputChange(index, value)}
             allowance={tokenAllowances && formatToHuman(tokenAllowances[index] || 0n, token.decimals)}
-            balance={tokenBalances && formatToHuman(tokenBalances[index] || 0n, token.decimals)}
+            balance={tokenBalances && formatToHuman(tokenBalances[token.address] || 0n, token.decimals)}
           />
         ))}
       </div>

diff --git a/packages/nextjs/app/pools/_components/actions/SwapForm.tsx b/packages/nextjs/app/pools/_components/actions/SwapForm.tsx
@@ -35,7 +35,7 @@ const initialSwapConfig = {
  * 3. Approve the vault for the tokenIn used in the swap transaction (if necessary)
  * 4. Send transaction to swap the tokens
  */
-export const SwapForm: React.FC<PoolActionsProps> = ({ pool, refetchPool }) => {
+export const SwapForm: React.FC<PoolActionsProps> = ({ pool, refetchPool, tokenBalances, refetchTokenBalances }) => {
   const [queryResponse, setQueryResponse] = useState<QuerySwapResponse | null>(null);
   const [queryError, setQueryError] = useState<QueryPoolActionError>(null);
   const [swapConfig, setSwapConfig] = useState<SwapConfig>(initialSwapConfig);
@@ -54,7 +54,7 @@ export const SwapForm: React.FC<PoolActionsProps> = ({ pool, refetchPool }) => {
   const { querySwap, swap } = useSwap(pool, swapConfig);
   const { approveSpenderOnToken: approvePermit2OnToken } = useApprove(tokenIn.address, PERMIT2[chainId]);
   const { approveSpenderOnPermit2: approveRouterOnPermit2 } = useApprove(tokenIn.address, BALANCER_ROUTER[chainId]);
-  const { tokenAllowance, refetchTokenAllowance, tokenBalance, refetchTokenBalance } = useToken(tokenIn.address);
+  const { tokenAllowance, refetchTokenAllowance } = useToken(tokenIn.address);
 
   const sufficientAllowance = useMemo(() => {
     return tokenAllowance && tokenAllowance >= swapConfig.tokenIn.rawAmount;
@@ -171,14 +171,15 @@ export const SwapForm: React.FC<PoolActionsProps> = ({ pool, refetchPool }) => {
 
   const handleSwap = async () => {
     try {
+      const tokenBalance = tokenBalances[tokenIn.address];
       if (tokenBalance === null || tokenBalance === undefined || tokenBalance < swapConfig.tokenIn.rawAmount) {
         throw new Error("Insufficient user balance");
       }
       setIsSwapping(true);
       await swap();
       refetchPool();
       refetchTokenAllowance();
-      refetchTokenBalance();
+      refetchTokenBalances();
     } catch (e) {
       if (e instanceof Error) {
         console.error("error", e);
@@ -230,7 +231,7 @@ export const SwapForm: React.FC<PoolActionsProps> = ({ pool, refetchPool }) => {
         setTokenDropdownOpen={setTokenInDropdownOpen}
         selectableTokens={pool.poolTokens.filter(token => token.symbol !== tokenIn.symbol)}
         allowance={formatToHuman(tokenAllowance, tokenIn.decimals)}
-        balance={formatToHuman(tokenBalance, tokenIn.decimals)}
+        balance={formatToHuman(tokenBalances[tokenIn.address] ?? 0n, tokenIn.decimals)}
         isHighlighted={queryResponse?.swapKind === SwapKind.GivenIn}
       />
       <TokenField
@@ -243,6 +244,7 @@ export const SwapForm: React.FC<PoolActionsProps> = ({ pool, refetchPool }) => {
         setTokenDropdownOpen={setTokenOutDropdownOpen}
         selectableTokens={pool.poolTokens.filter(token => token.symbol !== tokenOut.symbol)}
         isHighlighted={queryResponse?.swapKind === SwapKind.GivenOut}
+        balance={formatToHuman(tokenBalances[tokenOut.address] ?? 0n, tokenOut.decimals)}
       />
 
       {!expectedAmount || (expectedAmount && swapReceipt) ? (

diff --git a/packages/nextjs/app/pools/page.tsx b/packages/nextjs/app/pools/page.tsx
@@ -82,7 +82,9 @@ const PoolDashboard = ({ pool, refetchPool }: { pool: Pool; refetchPool: Refetch
             <PoolAttributes pool={pool} />
           </div>
           <div className="flex flex-col gap-7">
-            {pool.poolConfig?.isPoolInitialized && <PoolActions pool={pool} refetchPool={refetchPool} />}
+            {pool.poolConfig?.isPoolInitialized && (
+              <PoolActions key={pool.address} pool={pool} refetchPool={refetchPool} />
+            )}
             <HooksConfig pool={pool} />
             <PoolConfig pool={pool} />
           </div>

diff --git a/packages/nextjs/components/ScaffoldEthAppWithProviders.tsx b/packages/nextjs/components/ScaffoldEthAppWithProviders.tsx
@@ -10,21 +10,10 @@ import { Footer } from "~~/components/Footer";
 import { Header } from "~~/components/Header";
 import { BlockieAvatar } from "~~/components/scaffold-eth";
 import { ProgressBar } from "~~/components/scaffold-eth/ProgressBar";
-import { useNativeCurrencyPrice } from "~~/hooks/scaffold-eth";
-import { useGlobalState } from "~~/services/store/store";
 import { wagmiConfig } from "~~/services/web3/wagmiConfig";
 import { appChains } from "~~/services/web3/wagmiConnectors";
 
 const ScaffoldEthApp = ({ children }: { children: React.ReactNode }) => {
-  const price = useNativeCurrencyPrice();
-  const setNativeCurrencyPrice = useGlobalState(state => state.setNativeCurrencyPrice);
-
-  useEffect(() => {
-    if (price > 0) {
-      setNativeCurrencyPrice(price);
-    }
-  }, [setNativeCurrencyPrice, price]);
-
   return (
     <>
       <div className="flex flex-col min-h-screen">

diff --git a/packages/nextjs/contracts/deployedContracts.ts b/packages/nextjs/contracts/deployedContracts.ts
@@ -7,7 +7,7 @@ import { GenericContractsDeclaration } from "~~/utils/scaffold-eth/contract";
 const deployedContracts = {
   31337: {
     ConstantSumFactory: {
-      address: "0xe91cEDac266d3A331D307CA1834A29B1156126C4",
+      address: "0x579bB3a775c831526D5E94EF9ccd6601F62790Bc",
       abi: [
         {
           type: "constructor",
@@ -387,7 +387,7 @@ const deployedContracts = {
       },
     },
     ConstantProductFactory: {
-      address: "0x4655Ed483C8aE255e2098A171624bd012aBC6b40",
+      address: "0x3626DEff4AFB3Acd8f217288cd33FBE3a337Ce0B",
       abi: [
         {
           type: "constructor",
@@ -767,7 +767,7 @@ const deployedContracts = {
       },
     },
     MockToken1: {
-      address: "0xC17a00DfCb10af9cB41645186e89969C96D1a3B6",
+      address: "0x69221a99e5Bc30E0cf891992e958E3Ba3815bfc6",
       abi: [
         {
           type: "constructor",
@@ -1125,7 +1125,7 @@ const deployedContracts = {
       },
     },
     MockToken2: {
-      address: "0xe093601b7740105Ce4f3A3740f387c9B7EBb683b",
+      address: "0x66B4cF3Be49431371E4241C462B4A75ae3a6E986",
       abi: [
         {
           type: "constructor",
@@ -1483,7 +1483,7 @@ const deployedContracts = {
       },
     },
     MockVeBAL: {
-      address: "0x4d75F68A6b570d2FE0C84666e3Aec200e654D506",
+      address: "0x01eE8e2C72a2Cd4C5206823ce393520758F5Bc79",
       abi: [
         {
           type: "constructor",
@@ -1841,7 +1841,7 @@ const deployedContracts = {
       },
     },
     VeBALFeeDiscountHook: {
-      address: "0xF81277cEC0240B78aaCB542Eb0f0A1B9f6D7b5fC",
+      address: "0xB93eD7BbBFfcedfA8125f0653Bc07E7D229ffE21",
       abi: [
         {
           type: "constructor",

diff --git a/packages/nextjs/hooks/balancer/types.ts b/packages/nextjs/hooks/balancer/types.ts
@@ -139,10 +139,12 @@ export type UseToken = {
   refetchTokenBalance: () => void;
 };
 
+export type TokenBalances = { [key: Address]: bigint };
+
 export type UseTokens = {
   tokenAllowances: (bigint | undefined)[] | undefined;
   refetchTokenAllowances: () => void;
-  tokenBalances?: (bigint | undefined)[];
+  tokenBalances: TokenBalances;
   refetchTokenBalances: () => void;
 };
 

diff --git a/packages/nextjs/hooks/balancer/useTokens.ts b/packages/nextjs/hooks/balancer/useTokens.ts
@@ -4,6 +4,7 @@ import { zeroAddress } from "viem";
 import { useContractReads, useWalletClient } from "wagmi";
 import { useTargetFork } from "~~/hooks/balancer";
 import { Permit2Allowance, UseTokens } from "~~/hooks/balancer/types";
+import { type TokenBalances } from "~~/hooks/balancer/types";
 
 /**
  * Custom hook for dealing with multiple tokens
@@ -23,13 +24,20 @@ export const useTokens = (amountsIn: InputAmount[]): UseTokens => {
   });
 
   const tokenBalances = useMemo(() => {
-    return balances?.map(balance => {
-      if (typeof balance.result === "bigint") {
-        return balance.result;
-      }
-      return undefined;
-    });
-  }, [balances]); // Only recompute if tokenAllowances changes
+    const balancesObject: TokenBalances = {};
+    if (balances) {
+      balances.forEach((res, idx) => {
+        const address = amountsIn[idx].address;
+        const balance = (res.result as bigint) ?? 0n;
+        balancesObject[address] = balance;
+      });
+    } else {
+      amountsIn.forEach(token => {
+        balancesObject[token.address] = 0n;
+      });
+    }
+    return balancesObject;
+  }, [balances, amountsIn]);
 
   const { data: allowances, refetch: refetchTokenAllowances } = useContractReads({
     contracts: amountsIn.map(token => ({