From db52c0bb24b12cc4dcd5bcf6e73f23ad6eb7878b Mon Sep 17 00:00:00 2001
From: Ben Goldberg <bengoldberg@rainbow.me>
Date: Sat, 25 May 2024 10:13:15 -0700
Subject: [PATCH] need to remove useAssetsToSell

---
 .../screens/Swap/components/SearchInput.tsx   |  1 -
 .../Swap/components/SwapInputAsset.tsx        | 18 +++------------
 .../Swap/components/SwapOutputAsset.tsx       | 22 ++++++-------------
 src/state/assets/userAssets.ts                |  2 ++
 4 files changed, 12 insertions(+), 31 deletions(-)

diff --git a/src/__swaps__/screens/Swap/components/SearchInput.tsx b/src/__swaps__/screens/Swap/components/SearchInput.tsx
index ec56734a43a..3fb718c8808 100644
--- a/src/__swaps__/screens/Swap/components/SearchInput.tsx
+++ b/src/__swaps__/screens/Swap/components/SearchInput.tsx
@@ -93,7 +93,6 @@ export const SearchInput = ({
                   animatedProps={searchInputValue}
                   onChange={onSearchQueryChange}
                   onBlur={() => {
-                    console.log('here');
                     onSearchQueryChange({
                       nativeEvent: {
                         text: '',
diff --git a/src/__swaps__/screens/Swap/components/SwapInputAsset.tsx b/src/__swaps__/screens/Swap/components/SwapInputAsset.tsx
index 6ca074c82a6..beea2cae743 100644
--- a/src/__swaps__/screens/Swap/components/SwapInputAsset.tsx
+++ b/src/__swaps__/screens/Swap/components/SwapInputAsset.tsx
@@ -15,9 +15,8 @@ import { TokenList } from '@/__swaps__/screens/Swap/components/TokenList/TokenLi
 import { BASE_INPUT_WIDTH, INPUT_INNER_WIDTH, INPUT_PADDING, THICK_BORDER_WIDTH } from '@/__swaps__/screens/Swap/constants';
 import { IS_ANDROID } from '@/env';
 import { useSwapContext } from '@/__swaps__/screens/Swap/providers/swap-provider';
-import { useAssetsToSell } from '@/__swaps__/screens/Swap/hooks/useAssetsToSell';
-import { isSameAssetWorklet } from '@/__swaps__/utils/assets';
 import { AmimatedSwapCoinIcon } from './AnimatedSwapCoinIcon';
+import { swapsStore } from '@/state/swaps/swapsStore';
 
 function SwapInputActionButton() {
   const { isDarkMode } = useColorMode();
@@ -80,21 +79,10 @@ function SwapInputIcon() {
 }
 
 function InputAssetBalanceBadge() {
-  const { internalSelectedInputAsset } = useSwapContext();
-
-  const userAssets = useAssetsToSell();
+  const selectedAssetBalance = swapsStore(state => state.inputAsset?.balance.display);
 
   const label = useDerivedValue(() => {
-    const asset = internalSelectedInputAsset.value;
-    if (!asset) return 'No balance';
-
-    const userAsset = userAssets.find(userAsset =>
-      isSameAssetWorklet(userAsset, {
-        address: asset.address,
-        chainId: asset.chainId,
-      })
-    );
-    return userAsset?.balance.display ?? 'No balance';
+    return selectedAssetBalance || 'No Balance';
   });
 
   return <BalanceBadge label={label} />;
diff --git a/src/__swaps__/screens/Swap/components/SwapOutputAsset.tsx b/src/__swaps__/screens/Swap/components/SwapOutputAsset.tsx
index 8b3ebf05028..28e1bfa7bce 100644
--- a/src/__swaps__/screens/Swap/components/SwapOutputAsset.tsx
+++ b/src/__swaps__/screens/Swap/components/SwapOutputAsset.tsx
@@ -15,9 +15,9 @@ import { TokenList } from '@/__swaps__/screens/Swap/components/TokenList/TokenLi
 import { BASE_INPUT_WIDTH, INPUT_INNER_WIDTH, INPUT_PADDING, THICK_BORDER_WIDTH } from '@/__swaps__/screens/Swap/constants';
 import { IS_ANDROID } from '@/env';
 import { useSwapContext } from '@/__swaps__/screens/Swap/providers/swap-provider';
-import { isSameAssetWorklet } from '@/__swaps__/utils/assets';
-import { useAssetsToSell } from '@/__swaps__/screens/Swap/hooks/useAssetsToSell';
 import { AmimatedSwapCoinIcon } from './AnimatedSwapCoinIcon';
+import { swapsStore } from '@/state/swaps/swapsStore';
+import { userAssetsStore } from '@/state/assets/userAssets';
 
 function SwapOutputActionButton() {
   const { isDarkMode } = useColorMode();
@@ -80,21 +80,13 @@ function SwapInputIcon() {
 }
 
 function OutputAssetBalanceBadge() {
-  const { internalSelectedOutputAsset } = useSwapContext();
-
-  const userAssets = useAssetsToSell();
+  const selectedAssetId = swapsStore(state => state.outputAsset?.uniqueId);
+  const selectedAssetBalance = userAssetsStore(state =>
+    selectedAssetId && state.hasUserAsset(selectedAssetId) ? state.getUserAsset(selectedAssetId).balance.display : undefined
+  );
 
   const label = useDerivedValue(() => {
-    const asset = internalSelectedOutputAsset.value;
-    if (!asset) return 'No balance';
-
-    const userAsset = userAssets.find(userAsset =>
-      isSameAssetWorklet(userAsset, {
-        address: asset.address,
-        chainId: asset.chainId,
-      })
-    );
-    return userAsset?.balance.display ?? 'No balance';
+    return selectedAssetBalance || 'No Balance';
   });
 
   return <BalanceBadge label={label} />;
diff --git a/src/state/assets/userAssets.ts b/src/state/assets/userAssets.ts
index 6c62bfa4899..57b456fb87e 100644
--- a/src/state/assets/userAssets.ts
+++ b/src/state/assets/userAssets.ts
@@ -5,6 +5,7 @@ import { RainbowError, logger } from '@/logger';
 export interface UserAssetsState {
   userAssets: Map<UniqueId, ParsedSearchAsset>;
 
+  hasUserAsset: (uniqueId: UniqueId) => boolean;
   getUserAsset: (uniqueId: UniqueId) => ParsedSearchAsset;
 }
 
@@ -65,6 +66,7 @@ export const userAssetsStore = createRainbowStore<UserAssetsState>(
     userAssets: new Map(),
 
     getUserAsset: (uniqueId: UniqueId) => get().userAssets.get(uniqueId) as ParsedSearchAsset,
+    hasUserAsset: (uniqueId: UniqueId) => get().userAssets.has(uniqueId),
   }),
   {
     storageKey: 'userAssets',