From 4ff68b31a5384a8d243327338ef935341604c9fe Mon Sep 17 00:00:00 2001
From: saml33 <slam.uke@gmail.com>
Date: Mon, 3 Apr 2023 09:52:00 +1000
Subject: [PATCH] retain size when toggling margin

---
 components/shared/LeverageSlider.tsx   |  2 +
 components/trade/AdvancedTradeForm.tsx | 77 ++++++++++++++++++++++----
 components/trade/MaxSizeButton.tsx     |  4 +-
 3 files changed, 70 insertions(+), 13 deletions(-)

diff --git a/components/shared/LeverageSlider.tsx b/components/shared/LeverageSlider.tsx
index 834abce6..4508aa05 100644
--- a/components/shared/LeverageSlider.tsx
+++ b/components/shared/LeverageSlider.tsx
@@ -30,6 +30,8 @@ const LeverageSlider = ({
   useEffect(() => {
     if (amount) {
       setValue(amount)
+    } else {
+      setValue(0)
     }
   }, [amount])
 
diff --git a/components/trade/AdvancedTradeForm.tsx b/components/trade/AdvancedTradeForm.tsx
index cbfbdd00..3f1d4c4b 100644
--- a/components/trade/AdvancedTradeForm.tsx
+++ b/components/trade/AdvancedTradeForm.tsx
@@ -39,7 +39,7 @@ import SpotButtonGroup from './SpotButtonGroup'
 import PerpButtonGroup from './PerpButtonGroup'
 import SolBalanceWarnings from '@components/shared/SolBalanceWarnings'
 import useSelectedMarket from 'hooks/useSelectedMarket'
-import { getDecimalCount } from 'utils/numbers'
+import { floorToDecimal, getDecimalCount } from 'utils/numbers'
 import LogoWithFallback from '@components/shared/LogoWithFallback'
 import useIpAddress from 'hooks/useIpAddress'
 import ButtonGroup from '@components/forms/ButtonGroup'
@@ -174,15 +174,70 @@ const AdvancedTradeForm = () => {
     })
   }, [])
 
-  const handleSetMargin = useCallback((e: ChangeEvent<HTMLInputElement>) => {
-    if (!e.target.checked) {
-      set((s) => {
-        s.tradeForm.quoteSize = ''
-        s.tradeForm.baseSize = ''
-      })
-    }
-    setUseMargin(e.target.checked)
-  }, [])
+  const handleSetMargin = useCallback(
+    (e: ChangeEvent<HTMLInputElement>) => {
+      if (!e.target.checked) {
+        const price =
+          tradeForm.tradeType === 'Market' ? oraclePrice : tradeForm.price
+        const group = mangoStore.getState().group
+        if (
+          group &&
+          mangoAccount &&
+          price &&
+          selectedMarket &&
+          selectedMarket instanceof Serum3Market
+        ) {
+          if (tradeForm.side === 'buy') {
+            const balance = mangoAccount.getTokenBalanceUi(
+              group.getFirstBankByTokenIndex(selectedMarket.quoteTokenIndex)
+            )
+            const max = Math.max(balance, 0)
+            if (parseFloat(tradeForm.quoteSize) > max) {
+              set((s) => {
+                if (max > 0) {
+                  s.tradeForm.quoteSize = floorToDecimal(
+                    max,
+                    tickDecimals
+                  ).toFixed()
+                  s.tradeForm.baseSize = floorToDecimal(
+                    max / Number(price),
+                    minOrderDecimals
+                  ).toFixed()
+                } else {
+                  s.tradeForm.quoteSize = ''
+                  s.tradeForm.baseSize = ''
+                }
+              })
+            }
+          } else {
+            const balance = mangoAccount.getTokenBalanceUi(
+              group.getFirstBankByTokenIndex(selectedMarket.baseTokenIndex)
+            )
+            const max = Math.max(balance, 0)
+            if (parseFloat(tradeForm.baseSize) > max) {
+              set((s) => {
+                if (max > 0) {
+                  s.tradeForm.baseSize = floorToDecimal(
+                    max,
+                    minOrderDecimals
+                  ).toFixed()
+                  s.tradeForm.quoteSize = floorToDecimal(
+                    max * Number(price),
+                    tickDecimals
+                  ).toFixed()
+                } else {
+                  s.tradeForm.baseSize = ''
+                  s.tradeForm.quoteSize = ''
+                }
+              })
+            }
+          }
+        }
+      }
+      setUseMargin(e.target.checked)
+    },
+    [mangoAccount, oraclePrice, selectedMarket, set, tradeForm]
+  )
 
   const [tickDecimals, tickSize] = useMemo(() => {
     const group = mangoStore.getState().group
@@ -606,7 +661,7 @@ const AdvancedTradeForm = () => {
                   ? 'bg-th-up-dark text-white md:hover:bg-th-up'
                   : 'bg-th-down-dark text-white md:hover:bg-th-down'
               }`}
-              disabled={connected && !tradeForm.baseSize}
+              disabled={connected && (!tradeForm.baseSize || !tradeForm.price)}
               size="large"
               type="submit"
             >
diff --git a/components/trade/MaxSizeButton.tsx b/components/trade/MaxSizeButton.tsx
index 8f070ee0..c3e9675e 100644
--- a/components/trade/MaxSizeButton.tsx
+++ b/components/trade/MaxSizeButton.tsx
@@ -79,12 +79,12 @@ const MaxSizeButton = ({
         if (tradeType === 'Market' || !price) {
           state.tradeForm.quoteSize = floorToDecimal(
             max * oraclePrice,
-            minOrderDecimals
+            tickDecimals
           ).toFixed()
         } else {
           state.tradeForm.quoteSize = floorToDecimal(
             max * parseFloat(price),
-            minOrderDecimals
+            tickDecimals
           ).toFixed()
         }
       }