add function for token max

components/swap/LeverageSlider.tsx

@@ -1,7 +1,7 @@
 import Decimal from 'decimal.js'
-import { ChangeEvent, useEffect, useRef, useState } from 'react'
+import { ChangeEvent, useEffect, useMemo, useRef, useState } from 'react'
 import mangoStore from '../../store/mangoStore'
-import { useTokenMax } from './useTokenMax'
+import { getTokenInMax } from './useTokenMax'
 
 const LeverageSlider = ({
   amount,
@@ -74,7 +74,17 @@ export const SwapLeverageSlider = ({
   onChange: (x: string) => void
 }) => {
   const mangoAccount = mangoStore((s) => s.mangoAccount.current)
-  const { amountWithBorrow } = useTokenMax()
+  const inputTokenInfo = mangoStore((s) => s.swap.inputTokenInfo)
+
+  const { amountWithBorrow } = useMemo(() => {
+    const group = mangoStore.getState().group
+    if (inputTokenInfo && group) {
+      return getTokenInMax(inputTokenInfo.address, group, true)
+    }
+    return {
+      amountWithBorrow: new Decimal(0),
+    }
+  }, [inputTokenInfo])
 
   return (
     <>

components/swap/SwapForm.tsx

@@ -32,7 +32,7 @@ import {
   INPUT_TOKEN_DEFAULT,
   OUTPUT_TOKEN_DEFAULT,
 } from '../../utils/constants'
-import { useTokenMax } from './useTokenMax'
+import { getTokenInMax } from './useTokenMax'
 import WalletIcon from '../icons/WalletIcon'
 
 const MAX_DIGITS = 11
@@ -60,11 +60,22 @@ const SwapForm = () => {
   const jupiterTokens = mangoStore((s) => s.jupiterTokens)
   const connected = mangoStore((s) => s.connected)
   const [debouncedAmountIn] = useDebounce(amountInFormValue, 300)
+
   const {
     amount: tokenMax,
     amountWithBorrow,
     decimals,
-  } = useTokenMax(useMargin)
+  } = useMemo(() => {
+    const group = mangoStore.getState().group
+    if (inputTokenInfo && group) {
+      return getTokenInMax(inputTokenInfo.address, group, useMargin)
+    }
+    return {
+      amount: new Decimal(0),
+      amountWithBorrow: new Decimal(0),
+      decimals: 6,
+    }
+  }, [inputTokenInfo, useMargin])
 
   const amountIn: Decimal | null = useMemo(() => {
     return Number(debouncedAmountIn)

components/swap/SwapFormTokenList.tsx

@@ -12,6 +12,7 @@ import {
 import { useTranslation } from 'next-i18next'
 import { floorToDecimal } from '../../utils/numbers'
 import Decimal from 'decimal.js'
+import { getTokenInMax } from './useTokenMax'
 
 const generateSearchTerm = (item: Token, searchValue: string) => {
   const normalizedSearchValue = searchValue.toLowerCase()
@@ -53,8 +54,8 @@ const TokenItem = ({
   const { address, symbol, logoURI, name } = token
 
   const isDisabled =
-    (type === 'input' && useMargin && token.maxAmountWithBorrow!.eq(0)) ||
-    (type === 'input' && !useMargin && token.maxAmount!.eq(0))
+    (type === 'input' && useMargin && token.amountWithBorrow!.eq(0)) ||
+    (type === 'input' && !useMargin && token.amount!.eq(0))
 
   return (
     <div>
@@ -79,8 +80,8 @@ const TokenItem = ({
         {type === 'input' ? (
           <p className="text-sm text-th-fgd-2">
             {useMargin
-              ? token.maxAmountWithBorrow!.toString()
-              : token.maxAmount!.toString()}
+              ? token.amountWithBorrow!.toString()
+              : token.amount!.toString()}
           </p>
         ) : null}
       </button>
@@ -139,45 +140,14 @@ const SwapFormTokenList = ({
     ) {
       const filteredSortedTokens = tokens
         .map((token) => {
-          const bank = group.banksMapByMint.get(token.address)![0]
-          const tokenBalance = floorToDecimal(
-            mangoAccount.getTokenBalanceUi(bank),
-            bank.mintDecimals
-          )
-          const maxAmountWithoutMargin = tokenBalance.gt(0)
-            ? tokenBalance
-            : new Decimal(0)
-          const maxUiAmountWithBorrow = floorToDecimal(
-            mangoAccount?.getMaxSourceUiForTokenSwap(
-              group,
-              bank.mint,
-              outputBank.mint,
-              1
-            )!,
-            bank.mintDecimals
-          )
-          const inputBankVaultBalance = group.getTokenVaultBalanceByMintUi(
-            bank.mint
-          )
-          const maxAmount = useMargin
-            ? Decimal.min(
-                maxAmountWithoutMargin,
-                inputBankVaultBalance,
-                maxUiAmountWithBorrow!
-              )
-            : Decimal.min(maxAmountWithoutMargin, inputBankVaultBalance)
-
-          const maxAmountWithBorrow = Decimal.min(
-            maxUiAmountWithBorrow!,
-            inputBankVaultBalance
-          )
-          return { ...token, maxAmount, maxAmountWithBorrow }
+          const max = getTokenInMax(token.address, group, useMargin)
+          return { ...token, ...max }
         })
         .filter((token) => (token.symbol === outputBank?.name ? false : true))
         .sort((a, b) =>
           useMargin
-            ? Number(b.maxAmountWithBorrow) - Number(a.maxAmountWithBorrow)
-            : Number(b.maxAmount) - Number(a.maxAmount)
+            ? Number(b.amountWithBorrow) - Number(a.amountWithBorrow)
+            : Number(b.amount) - Number(a.amount)
         )
 
       return filteredSortedTokens
@@ -185,8 +155,8 @@ const SwapFormTokenList = ({
       const filteredTokens = tokens
         .map((token) => ({
           ...token,
-          maxAmount: new Decimal(0),
-          maxAmountWithBorrow: new Decimal(0),
+          amount: new Decimal(0),
+          amountWithBorrow: new Decimal(0),
         }))
         .filter((token) => (token.symbol === inputBank?.name ? false : true))
       return filteredTokens

components/swap/useTokenMax.tsx

@@ -22,6 +22,63 @@ export const getMaxWithdrawForBank = (
     : Decimal.min(accountBalance, vaultBalance, maxBorrow!)
 }
 
+export const getTokenInMax = (
+  inputTokenAddress: string,
+  group: Group,
+  useMargin: boolean
+) => {
+  const mangoAccount = mangoStore.getState().mangoAccount.current
+  const outputBank = mangoStore.getState().swap.outputBank
+  const inputBank = group.banksMapByMint.get(inputTokenAddress)![0]
+
+  if (!group || !inputBank || !mangoAccount || !outputBank) {
+    return {
+      amount: new Decimal(0.0),
+      decimals: 6,
+      amountWithBorrow: new Decimal(0.0),
+    }
+  }
+
+  const inputTokenBalance = floorToDecimal(
+    mangoAccount.getTokenBalanceUi(inputBank),
+    inputBank.mintDecimals
+  )
+  const maxAmountWithoutMargin = inputTokenBalance.gt(0)
+    ? inputTokenBalance
+    : new Decimal(0)
+  const maxUiAmountWithBorrow = floorToDecimal(
+    mangoAccount?.getMaxSourceUiForTokenSwap(
+      group,
+      inputBank.mint,
+      outputBank.mint,
+      1
+    )!,
+    inputBank.mintDecimals
+  )
+  const inputBankVaultBalance = group.getTokenVaultBalanceByMintUi(
+    inputBank.mint
+  )
+
+  const maxAmount = useMargin
+    ? Decimal.min(
+        maxAmountWithoutMargin,
+        inputBankVaultBalance,
+        maxUiAmountWithBorrow!
+      )
+    : Decimal.min(maxAmountWithoutMargin, inputBankVaultBalance)
+
+  const maxAmountWithBorrow = Decimal.min(
+    maxUiAmountWithBorrow!,
+    inputBankVaultBalance
+  )
+
+  return {
+    amount: maxAmount,
+    amountWithBorrow: maxAmountWithBorrow,
+    decimals: inputBank.mintDecimals,
+  }
+}
+
 export const useTokenMax = (useMargin = true) => {
   const { mangoAccount } = useMangoAccount()
   const inputBank = mangoStore((s) => s.swap.inputBank)

types/jupiter.ts

@@ -17,6 +17,6 @@ export interface Token {
   extensions?: {
     coingeckoId: string
   }
-  maxAmount?: Decimal
-  maxAmountWithBorrow?: Decimal
+  amount?: Decimal
+  amountWithBorrow?: Decimal
 }