userAssetsStore refactor

src/__swaps__/screens/Swap/Swap.tsx

@@ -28,6 +28,8 @@ import { useSwapsStore } from '@/state/swaps/swapsStore';
 import { SwapWarning } from './components/SwapWarning';
 import { clearCustomGasSettings } from './hooks/useCustomGas';
 import { SwapProvider, useSwapContext } from './providers/swap-provider';
+import { Address } from 'viem';
+import { useAccountSettings } from '@/hooks';
 
 /** README
  * This prototype is largely driven by Reanimated and Gesture Handler, which
@@ -103,9 +105,11 @@ const useMountSignal = () => {
 };
 
 const useCleanupOnUnmount = () => {
+  const { accountAddress } = useAccountSettings();
+
   useEffect(() => {
     return () => {
-      const highestValueEth = userAssetsStore.getState().getHighestValueEth();
+      const highestValueEth = userAssetsStore.getState(accountAddress as Address).getHighestValueEth();
       const parsedAsset = highestValueEth
         ? parseSearchAsset({
             assetWithPrice: undefined,
@@ -123,39 +127,39 @@ const useCleanupOnUnmount = () => {
         selectedOutputChainId: parsedAsset?.chainId ?? ChainId.mainnet,
       });
 
-      userAssetsStore.setState({ filter: 'all', inputSearchQuery: '' });
+      userAssetsStore.setState(accountAddress as Address, { filter: 'all', inputSearchQuery: '' });
 
       clearCustomGasSettings();
     };
-  }, []);
+  }, [accountAddress]);
 };
 
 const WalletAddressObserver = () => {
-  const currentWalletAddress = userAssetsStore(state => state.associatedWalletAddress);
+  const { accountAddress } = useAccountSettings();
   const { setAsset } = useSwapContext();
 
   const setNewInputAsset = useCallback(() => {
-    const newHighestValueEth = userAssetsStore.getState().getHighestValueEth();
+    const newHighestValueEth = userAssetsStore.getState(accountAddress as Address).getHighestValueEth();
 
-    if (userAssetsStore.getState().filter !== 'all') {
-      userAssetsStore.setState({ filter: 'all' });
+    if (userAssetsStore.getState(accountAddress as Address).filter !== 'all') {
+      userAssetsStore.setState(accountAddress as Address, { filter: 'all' });
     }
 
     setAsset({
       type: SwapAssetType.inputAsset,
       asset: newHighestValueEth,
     });
 
-    if (userAssetsStore.getState().userAssets.size === 0) {
+    if (userAssetsStore.getState(accountAddress as Address).userAssets.size === 0) {
       setAsset({
         type: SwapAssetType.outputAsset,
         asset: null,
       });
     }
-  }, [setAsset]);
+  }, [accountAddress, setAsset]);
 
   useAnimatedReaction(
-    () => currentWalletAddress,
+    () => accountAddress,
     (current, previous) => {
       const didWalletAddressChange = previous && current !== previous;
 

src/__swaps__/screens/Swap/components/CoinRow.tsx

@@ -11,13 +11,14 @@ import * as i18n from '@/languages';
 import { RainbowNetworks } from '@/networks';
 import { BASE_DEGEN_ADDRESS, DEGEN_CHAIN_DEGEN_ADDRESS, ETH_ADDRESS } from '@/references';
 import { toggleFavorite } from '@/resources/favorites';
-import { userAssetsStore } from '@/state/assets/userAssets';
+import { useUserAssetsStore } from '@/state/assets/userAssets';
 import { ethereumUtils, haptics, showActionSheetWithOptions } from '@/utils';
 import { startCase } from 'lodash';
 import React, { useCallback, useMemo } from 'react';
 import { GestureResponderEvent } from 'react-native';
 import { OnPressMenuItemEventObject } from 'react-native-ios-context-menu';
 import { SwapCoinIcon } from './SwapCoinIcon';
+import { Address } from 'viem';
 
 export const COIN_ROW_WITH_PADDING_HEIGHT = 56;
 
@@ -52,6 +53,7 @@ interface InputCoinRowProps {
   onPress: (asset: ParsedSearchAsset | null) => void;
   output?: false | undefined;
   uniqueId: string;
+  walletAddress: string;
 }
 
 type PartialAsset = Pick<SearchAsset, 'address' | 'chainId' | 'colors' | 'icon_url' | 'mainnetAddress' | 'name' | 'symbol' | 'uniqueId'>;
@@ -62,12 +64,13 @@ interface OutputCoinRowProps extends PartialAsset {
   output: true;
   nativePriceChange?: string;
   isTrending?: boolean;
+  walletAddress: string;
 }
 
 type CoinRowProps = InputCoinRowProps | OutputCoinRowProps;
 
-export function CoinRow({ isFavorite, onPress, output, uniqueId, ...assetProps }: CoinRowProps) {
-  const inputAsset = userAssetsStore(state => (output ? undefined : state.getUserAsset(uniqueId)));
+export function CoinRow({ isFavorite, onPress, output, uniqueId, walletAddress, ...assetProps }: CoinRowProps) {
+  const inputAsset = useUserAssetsStore(walletAddress as Address, state => (output ? undefined : state.getUserAsset(uniqueId)));
   const outputAsset = output ? (assetProps as PartialAsset) : undefined;
 
   const asset = output ? outputAsset : inputAsset;

src/__swaps__/screens/Swap/components/SearchInput.tsx

@@ -4,7 +4,7 @@ import { opacity } from '@/__swaps__/utils/swaps';
 import { Input } from '@/components/inputs';
 import { Bleed, Box, Column, Columns, Text, useColorMode, useForegroundColor } from '@/design-system';
 import * as i18n from '@/languages';
-import { userAssetsStore } from '@/state/assets/userAssets';
+import { userAssetsStore, useUserAssetsStore } from '@/state/assets/userAssets';
 import { useSwapsStore } from '@/state/swaps/swapsStore';
 import React from 'react';
 import Animated, {
@@ -17,6 +17,8 @@ import Animated, {
 } from 'react-native-reanimated';
 import { useDebouncedCallback } from 'use-debounce';
 import { SearchInputButton } from './SearchInputButton';
+import { useAccountSettings } from '@/hooks';
+import { Address } from 'viem';
 
 const AnimatedInput = Animated.createAnimatedComponent(Input);
 
@@ -34,12 +36,13 @@ export const SearchInput = ({
 }) => {
   const { isDarkMode } = useColorMode();
   const { inputProgress, inputSearchRef, outputProgress, outputSearchRef } = useSwapContext();
+  const { accountAddress } = useAccountSettings();
 
   const fillTertiary = useForegroundColor('fillTertiary');
   const label = useForegroundColor('label');
   const labelQuaternary = useForegroundColor('labelQuaternary');
 
-  const onInputSearchQueryChange = userAssetsStore(state => state.setSearchQuery);
+  const onInputSearchQueryChange = useUserAssetsStore(accountAddress as Address, state => state.setSearchQuery);
 
   const onOutputSearchQueryChange = useDebouncedCallback((text: string) => useSwapsStore.setState({ outputSearchQuery: text }), 100, {
     leading: false,
@@ -108,8 +111,8 @@ export const SearchInput = ({
                           useSwapsStore.setState({ outputSearchQuery: '' });
                         }
                       } else {
-                        if (userAssetsStore.getState().inputSearchQuery !== '') {
-                          userAssetsStore.getState().setSearchQuery('');
+                        if (userAssetsStore.getState(accountAddress as Address).inputSearchQuery !== '') {
+                          userAssetsStore.getState(accountAddress as Address).setSearchQuery('');
                         }
                       }
                     }

src/__swaps__/screens/Swap/components/TokenList/ChainSelection.tsx

@@ -13,12 +13,13 @@ import { analyticsV2 } from '@/analytics';
 import { ChainImage } from '@/components/coin-icon/ChainImage';
 import { ContextMenuButton } from '@/components/context-menu';
 import { AnimatedText, Bleed, Box, Inline, Text, TextIcon, globalColors, useColorMode } from '@/design-system';
-import { useAccountAccentColor } from '@/hooks';
+import { useAccountAccentColor, useAccountSettings } from '@/hooks';
 import { useSharedValueState } from '@/hooks/reanimated/useSharedValueState';
-import { userAssetsStore } from '@/state/assets/userAssets';
+import { userAssetsStore, useUserAssetsStore } from '@/state/assets/userAssets';
 import { swapsStore } from '@/state/swaps/swapsStore';
 import { showActionSheetWithOptions } from '@/utils';
 import { OnPressMenuItemEventObject } from 'react-native-ios-context-menu';
+import { Address } from 'viem';
 
 type ChainSelectionProps = {
   allText?: string;
@@ -28,11 +29,15 @@ type ChainSelectionProps = {
 export const ChainSelection = memo(function ChainSelection({ allText, output }: ChainSelectionProps) {
   const { isDarkMode } = useColorMode();
   const { accentColor: accountColor } = useAccountAccentColor();
+  const { accountAddress } = useAccountSettings();
   const { selectedOutputChainId, setSelectedOutputChainId } = useSwapContext();
 
   // chains sorted by balance on output, chains without balance hidden on input
-  const balanceSortedChainList = userAssetsStore(state => (output ? state.getBalanceSortedChainList() : state.getChainsWithBalance()));
-  const inputListFilter = useSharedValue(userAssetsStore.getState().filter);
+  const { balanceSortedChainList, filter } = useUserAssetsStore(accountAddress as Address, state => ({
+    balanceSortedChainList: output ? state.getBalanceSortedChainList() : state.getChainsWithBalance(),
+    filter: state.filter,
+  }));
+  const inputListFilter = useSharedValue(filter);
 
   const accentColor = useMemo(() => {
     if (c.contrast(accountColor, isDarkMode ? '#191A1C' : globalColors.white100) < (isDarkMode ? 2.125 : 1.5)) {
@@ -63,12 +68,12 @@ export const ChainSelection = memo(function ChainSelection({ allText, output }:
         setSelectedOutputChainId(Number(actionKey) as ChainId);
       } else {
         inputListFilter.value = actionKey === 'all' ? 'all' : (Number(actionKey) as ChainId);
-        userAssetsStore.setState({
+        userAssetsStore.setState(accountAddress as Address, {
           filter: actionKey === 'all' ? 'all' : (Number(actionKey) as ChainId),
         });
       }
     },
-    [inputListFilter, output, setSelectedOutputChainId]
+    [accountAddress, inputListFilter, output, setSelectedOutputChainId]
   );
 
   const menuConfig = useMemo(() => {
@@ -188,8 +193,9 @@ export const ChainSelection = memo(function ChainSelection({ allText, output }:
 
 const ChainButtonIcon = ({ output }: { output: boolean | undefined }) => {
   const { selectedOutputChainId: animatedSelectedOutputChainId } = useSwapContext();
+  const { accountAddress } = useAccountSettings();
 
-  const userAssetsFilter = userAssetsStore(state => (output ? undefined : state.filter));
+  const userAssetsFilter = useUserAssetsStore(accountAddress as Address, state => (output ? undefined : state.filter));
   const selectedOutputChainId = useSharedValueState(animatedSelectedOutputChainId, { pauseSync: !output });
 
   return (

src/__swaps__/screens/Swap/components/TokenList/TokenToBuyList.tsx

@@ -22,6 +22,8 @@ import { ScrollViewProps } from 'react-native';
 import Animated, { runOnUI, useAnimatedProps, useAnimatedStyle, withTiming } from 'react-native-reanimated';
 import { EXPANDED_INPUT_HEIGHT, FOCUSED_INPUT_HEIGHT } from '../../constants';
 import { ChainSelection } from './ChainSelection';
+import { useAccountSettings } from '@/hooks';
+import { Address } from 'viem';
 
 export const BUY_LIST_HEADER_HEIGHT = 20 + 10 + 8; // paddingTop + height + paddingBottom
 
@@ -86,6 +88,7 @@ const ScrollViewWithRef = forwardRef<Animated.ScrollView>(function ScrollViewWit
 export const TokenToBuyList = () => {
   const { internalSelectedInputAsset, internalSelectedOutputAsset, isFetching, isQuoteStale, outputProgress, setAsset } = useSwapContext();
   const { results: sections, isLoading } = useSearchCurrencyLists();
+  const { accountAddress } = useAccountSettings();
 
   const handleSelectToken = useCallback(
     (token: SearchAsset) => {
@@ -99,7 +102,7 @@ export const TokenToBuyList = () => {
         }
       })();
 
-      const userAsset = userAssetsStore.getState().getUserAsset(token.uniqueId);
+      const userAsset = userAssetsStore.getState(accountAddress as Address).getUserAsset(token.uniqueId);
       const parsedAsset = parseSearchAsset({
         assetWithPrice: undefined,
         searchAsset: token,
@@ -121,7 +124,7 @@ export const TokenToBuyList = () => {
         });
       }
     },
-    [internalSelectedInputAsset, internalSelectedOutputAsset, isFetching, isQuoteStale, setAsset]
+    [accountAddress, internalSelectedInputAsset, internalSelectedOutputAsset, isFetching, isQuoteStale, setAsset]
   );
 
   const animatedListPadding = useAnimatedStyle(() => {
@@ -172,6 +175,7 @@ export const TokenToBuyList = () => {
               output
               symbol={item.symbol}
               uniqueId={item.uniqueId}
+              walletAddress={accountAddress}
             />
           );
         }}

src/__swaps__/screens/Swap/components/TokenList/TokenToSellList.tsx

@@ -7,14 +7,16 @@ import { getStandardizedUniqueIdWorklet } from '@/__swaps__/utils/swaps';
 import { analyticsV2 } from '@/analytics';
 import { useDelayedMount } from '@/hooks/useDelayedMount';
 import * as i18n from '@/languages';
-import { userAssetsStore } from '@/state/assets/userAssets';
+import { userAssetsStore, useUserAssetsStore } from '@/state/assets/userAssets';
 import { swapsStore } from '@/state/swaps/swapsStore';
 import { DEVICE_WIDTH } from '@/utils/deviceUtils';
 import { FlashList } from '@shopify/flash-list';
 import React, { useCallback, useMemo } from 'react';
 import Animated, { runOnUI, useAnimatedProps, useAnimatedStyle } from 'react-native-reanimated';
 import { EXPANDED_INPUT_HEIGHT, FOCUSED_INPUT_HEIGHT } from '../../constants';
 import { ChainSelection } from './ChainSelection';
+import { useAccountSettings } from '@/hooks';
+import { Address } from 'viem';
 
 export const SELL_LIST_HEADER_HEIGHT = 20 + 10 + 14; // paddingTop + height + paddingBottom
 
@@ -31,8 +33,9 @@ export const TokenToSellList = () => {
 
 const TokenToSellListComponent = () => {
   const { inputProgress, internalSelectedInputAsset, internalSelectedOutputAsset, isFetching, isQuoteStale, setAsset } = useSwapContext();
+  const { accountAddress } = useAccountSettings();
 
-  const userAssetIds = userAssetsStore(state => state.getFilteredUserAssetIds());
+  const userAssetIds = useUserAssetsStore(accountAddress as Address, state => state.getFilteredUserAssetIds());
 
   const handleSelectToken = useCallback(
     (token: ParsedSearchAsset | null) => {
@@ -53,7 +56,7 @@ const TokenToSellListComponent = () => {
         asset: token,
       });
 
-      const { inputSearchQuery } = userAssetsStore.getState();
+      const { inputSearchQuery } = userAssetsStore.getState(accountAddress as Address);
 
       // track what search query the user had prior to selecting an asset
       if (inputSearchQuery.trim().length) {
@@ -63,7 +66,7 @@ const TokenToSellListComponent = () => {
         });
       }
     },
-    [internalSelectedInputAsset, internalSelectedOutputAsset, isFetching, isQuoteStale, setAsset]
+    [accountAddress, internalSelectedInputAsset, internalSelectedOutputAsset, isFetching, isQuoteStale, setAsset]
   );
 
   const animatedListPadding = useAnimatedStyle(() => {
@@ -92,7 +95,14 @@ const TokenToSellListComponent = () => {
       estimatedListSize={{ height: EXPANDED_INPUT_HEIGHT - 77, width: DEVICE_WIDTH - 24 }}
       keyExtractor={uniqueId => uniqueId}
       renderItem={({ item: uniqueId }) => {
-        return <CoinRow onPress={(asset: ParsedSearchAsset | null) => handleSelectToken(asset)} output={false} uniqueId={uniqueId} />;
+        return (
+          <CoinRow
+            onPress={(asset: ParsedSearchAsset | null) => handleSelectToken(asset)}
+            output={false}
+            uniqueId={uniqueId}
+            walletAddress={accountAddress}
+          />
+        );
       }}
       renderScrollComponent={props => {
         return (

src/__swaps__/screens/Swap/hooks/useAssetsToSell.ts

@@ -9,7 +9,7 @@ import {
 import { useUserAssets } from '@/__swaps__/screens/Swap/resources/assets';
 import { ParsedAssetsDictByChain, ParsedSearchAsset, UserAssetFilter } from '@/__swaps__/types/assets';
 import { useAccountSettings, useDebounce } from '@/hooks';
-import { userAssetsStore } from '@/state/assets/userAssets';
+import { useUserAssetsStore } from '@/state/assets/userAssets';
 
 const sortBy = (by: UserAssetFilter) => {
   switch (by) {
@@ -23,8 +23,10 @@ const sortBy = (by: UserAssetFilter) => {
 export const useAssetsToSell = () => {
   const { accountAddress: currentAddress, nativeCurrency: currentCurrency } = useAccountSettings();
 
-  const filter = userAssetsStore(state => state.filter);
-  const searchQuery = userAssetsStore(state => state.inputSearchQuery);
+  const { filter, searchQuery } = useUserAssetsStore(currentAddress as Address, state => ({
+    filter: state.filter,
+    searchQuery: state.inputSearchQuery,
+  }));
 
   const debouncedAssetToSellFilter = useDebounce(searchQuery, 200);
 

src/__swaps__/screens/Swap/providers/swap-provider.tsx

@@ -131,7 +131,7 @@ interface SwapProviderProps {
 }
 
 export const SwapProvider = ({ children }: SwapProviderProps) => {
-  const { nativeCurrency } = useAccountSettings();
+  const { accountAddress, nativeCurrency } = useAccountSettings();
 
   const isFetching = useSharedValue(false);
   const isQuoteStale = useSharedValue(0); // TODO: Convert this to a boolean
@@ -146,8 +146,8 @@ export const SwapProvider = ({ children }: SwapProviderProps) => {
   const lastTypedInput = useSharedValue<inputKeys>('inputAmount');
   const focusedInput = useSharedValue<inputKeys>('inputAmount');
 
-  const initialSelectedInputAsset = parseAssetAndExtend({ asset: swapsStore.getState().inputAsset });
-  const initialSelectedOutputAsset = parseAssetAndExtend({ asset: swapsStore.getState().outputAsset });
+  const initialSelectedInputAsset = parseAssetAndExtend({ asset: swapsStore.getState().inputAsset, walletAddress: accountAddress });
+  const initialSelectedOutputAsset = parseAssetAndExtend({ asset: swapsStore.getState().outputAsset, walletAddress: accountAddress });
 
   const internalSelectedInputAsset = useSharedValue<ExtendedAnimatedAssetWithColors | null>(initialSelectedInputAsset);
   const internalSelectedOutputAsset = useSharedValue<ExtendedAnimatedAssetWithColors | null>(initialSelectedOutputAsset);
@@ -552,7 +552,7 @@ export const SwapProvider = ({ children }: SwapProviderProps) => {
   const setAsset = useCallback(
     ({ type, asset }: { type: SwapAssetType; asset: ParsedSearchAsset | null }) => {
       const insertUserAssetBalance = type !== SwapAssetType.inputAsset;
-      const extendedAsset = parseAssetAndExtend({ asset, insertUserAssetBalance });
+      const extendedAsset = parseAssetAndExtend({ asset, insertUserAssetBalance, walletAddress: accountAddress });
 
       const otherSelectedAsset = type === SwapAssetType.inputAsset ? internalSelectedOutputAsset.value : internalSelectedInputAsset.value;
       const isSameAsOtherAsset = !!(otherSelectedAsset && otherSelectedAsset.uniqueId === extendedAsset?.uniqueId);
@@ -596,7 +596,11 @@ export const SwapProvider = ({ children }: SwapProviderProps) => {
 
       if (didSelectedAssetChange) {
         const assetToSet = insertUserAssetBalance
-          ? { ...asset, balance: (asset && userAssetsStore.getState().getUserAsset(asset.uniqueId)?.balance) || asset?.balance }
+          ? {
+              ...asset,
+              balance:
+                (asset && userAssetsStore.getState(accountAddress as Address).getUserAsset(asset.uniqueId)?.balance) || asset?.balance,
+            }
           : asset;
 
         if (isSameAsOtherAsset) {
@@ -647,9 +651,10 @@ export const SwapProvider = ({ children }: SwapProviderProps) => {
     },
     [
       SwapInputController.quoteFetchingInterval,
+      accountAddress,
       handleProgressNavigation,
-      internalSelectedInputAsset,
-      internalSelectedOutputAsset,
+      internalSelectedInputAsset.value,
+      internalSelectedOutputAsset.value,
       selectedOutputChainId,
       updateAssetValue,
     ]