FEAT: close original ATA upon upgrade

packages/ui/src/entities/token/use-token-balance.tsx

@@ -4,10 +4,26 @@ import { useConnection } from "@solana/wallet-adapter-react"
 import { QueryObserverOptions, useQuery } from "@tanstack/react-query"
 import { useWallet } from "@solana/wallet-adapter-react"
 
+type InfoTokenAmount = {
+  amount: string
+  decimals: number
+  uiAmount: number
+  uiAmountString: string
+}
+
+const empty: InfoTokenAmount = {
+  amount: "0",
+  decimals: -1,
+  uiAmount: 0,
+  uiAmountString: "0",
+}
+
+export { empty as placeholderData }
+
 export function useTokenBalance(
   address: web3.PublicKey | string | undefined,
   opts?: Pick<
-    QueryObserverOptions<string>,
+    QueryObserverOptions<InfoTokenAmount>,
     "refetchInterval" | "refetchIntervalInBackground" | "placeholderData"
   >,
 ) {
@@ -18,10 +34,10 @@ export function useTokenBalance(
   const { data: balance, error } = useQuery({
     enabled: Boolean(address && publicKey),
     placeholderData: opts?.placeholderData,
-    queryFn: async (): Promise<string> => {
+    queryFn: async (): Promise<InfoTokenAmount> => {
       // no token selected yet
       if (!address) {
-        return "0"
+        return empty
       }
 
       if (typeof address === "string") {
@@ -30,16 +46,23 @@ export function useTokenBalance(
 
       if (!publicKey) {
         // no wallet connected yet on solana
-        return "0"
+        return empty
       }
 
       if (NATIVE_MINT.equals(new web3.PublicKey(address))) {
         const balance = await connection.getBalance(publicKey)
         const amount = balance / web3.LAMPORTS_PER_SOL
 
-        return amount.toFixed(
+        const a = amount.toFixed(
           Math.fround(Math.log(web3.LAMPORTS_PER_SOL) / Math.log(10)),
         )
+
+        return {
+          amount: String(balance),
+          decimals: 9,
+          uiAmount: amount,
+          uiAmountString: a,
+        }
       } else {
         // Get the initial solana token balance
         const results = await connection.getParsedTokenAccountsByOwner(
@@ -48,21 +71,19 @@ export function useTokenBalance(
         )
 
         for (const item of results.value) {
-          // TODO: use schema to parse values
+          // FEAT: use schema to parse values
           const tokenInfo = {
             mint: item.account.data.parsed.info.mint as string,
-            tokenAmount: item.account.data.parsed.info.tokenAmount as {
-              uiAmountString: string
-            },
+            tokenAmount: item.account.data.parsed.info
+              .tokenAmount as InfoTokenAmount,
           }
           const mintAddress = tokenInfo.mint
-          const amount = tokenInfo.tokenAmount.uiAmountString
           if (mintAddress === address.toString()) {
-            return amount
+            return tokenInfo.tokenAmount
           }
         }
       }
-      return "0"
+      return empty
     },
     queryKey: ["useTokenBalance", String(address), String(publicKey)],
     refetchInterval: opts?.refetchInterval ?? 6000,

packages/ui/src/entities/upgrade/index.tsx

@@ -20,6 +20,18 @@ export async function upgradeToken(
 ): Promise<[web3.Transaction, web3.Signer[]]> {
   const originalMint = await spl.getMint(connection, oldToken)
 
+  /**
+   *  Converting uiAmount to the number of lamports.
+   *
+   *  It's needed to mult/div the result to remove inaccuracy.
+   */
+  function fromUiAmount(amount: number) {
+    const _amount = 1e10 * amount
+    const result = (_amount * Math.pow(10, originalMint.decimals)) / 1e10
+
+    return result
+}
+
   /// Anciliary creation
   //  Store N amount of token to upgrade
   const anciliaryAccountKeypair = web3.Keypair.generate()
@@ -47,6 +59,17 @@ export async function upgradeToken(
     spl.ASSOCIATED_TOKEN_PROGRAM_ID,
   )
 
+  const originalAccount = await spl.getAccount(connection, holderATA)
+
+  if (Number(originalAccount.amount) < fromUiAmount(amount)) {
+    throw new Error("Insufficient amount of token")
+  }
+  /// Checking that upgrade is for full amount of token
+  let shouldCloseOriginalATA = false
+  if (Number(originalAccount.amount) === fromUiAmount(amount)) {
+    shouldCloseOriginalATA = true
+  }
+
   /// Calculating minimal rent
   //
   const mintAccountRentExtemption =
@@ -72,7 +95,7 @@ export async function upgradeToken(
       holderATA,
       anciliaryAccountKeypair.publicKey,
       holder,
-      amount * Math.pow(10, originalMint.decimals),
+      fromUiAmount(amount),
     ),
     // Create associated account if needed
     spl.createAssociatedTokenAccountIdempotentInstruction(
@@ -100,6 +123,12 @@ export async function upgradeToken(
     ),
   ]
 
+  if (shouldCloseOriginalATA) {
+    instructions.push(
+      spl.createCloseAccountInstruction(holderATA, holder, holder),
+    )
+  }
+
   const exchangeTx = new web3.Transaction().add(...instructions)
 
   return [exchangeTx, [anciliaryAccountKeypair]]

packages/ui/src/widgets/token-upgrade.tsx

@@ -9,7 +9,7 @@ import useTokenAmount from "../entities/token/use-token-amount"
 import { Container } from "./token-upgrade/container"
 import { UpgradeButton } from "../features/upgrade-button"
 import { useMint } from "../entities/token/use-mint"
-import { useTokenBalance } from "../entities/token/use-token-balance"
+import { useTokenBalance, placeholderData } from "../entities/token/use-token-balance"
 import { withErrorBoundary } from "react-error-boundary"
 import { useTokenUpgrade } from "../entities/use-token-upgrade"
 
@@ -39,7 +39,7 @@ export function TokenUpgradeBase({
   tokenUpgradeProgramId,
 }: TokenUpgradeProps) {
   const [{ amount, destination }, setAction] = useTokenAmount()
-  const { balance } = useTokenBalance(tokenAddress, { placeholderData: "0" })
+  const { balance } = useTokenBalance(tokenAddress, { placeholderData })
   const { mint } = useMint(tokenAddress)
   const { mutate } = useTokenUpgrade()
 
@@ -121,7 +121,7 @@ export function TokenUpgradeBase({
           <Form.Field className="pb-4 pt-3.5" name="amount">
             <Amount
               address={tokenAddress}
-              balance={balance}
+              balance={balance?.uiAmountString}
               disabled={!tokenAddress}
               onAmountChange={onAmountChange}
               onAmountMaxChange={onAmountChange}