Fix #7966: Improve NFT balance fetching performance (#7968)

brave · Aug 25, 2023 · c1df10e · c1df10e
1 parent 3dce92a
commit c1df10e
Show file tree

Hide file tree

Showing 4 changed files with 64 additions and 48 deletions.
diff --git a/Sources/BraveWallet/Crypto/Accounts/AccountSelectionRootView.swift b/Sources/BraveWallet/Crypto/Accounts/AccountSelectionRootView.swift
@@ -92,6 +92,7 @@ private struct AccountListRowView: View {
           AccountView(address: account.address, name: account.name)
           checkmark
         }
+        .contentShape(Rectangle())
       }
       .buttonStyle(FadeButtonStyle())
     }

diff --git a/Sources/BraveWallet/Crypto/FiltersDisplaySettingsView.swift b/Sources/BraveWallet/Crypto/FiltersDisplaySettingsView.swift
@@ -529,6 +529,7 @@ private struct FilterDetailRowView<SelectionView: View>: View {
         .font(.body.weight(.semibold))
         .foregroundColor(Color(.separator))
     }
+    .contentShape(Rectangle())
   }
 }
 

diff --git a/Sources/BraveWallet/Crypto/Stores/NFTStore.swift b/Sources/BraveWallet/Crypto/Stores/NFTStore.swift
@@ -127,22 +127,7 @@ public class NFTStore: ObservableObject {
       let selectedAccounts = filters.accounts.filter(\.isSelected).map(\.model)
       let selectedNetworks = filters.networks.filter(\.isSelected).map(\.model)
       let allVisibleUserAssets = assetManager.getAllVisibleAssetsInNetworkAssets(networks: selectedNetworks)
-      var updatedUserVisibleNFTs: [NFTAssetViewModel] = []
-      for networkAssets in allVisibleUserAssets {
-        for token in networkAssets.tokens {
-          if token.isErc721 || token.isNft {
-            updatedUserVisibleNFTs.append(
-              NFTAssetViewModel(
-                token: token,
-                network: networkAssets.network,
-                balanceForAccounts: nftBalancesCache[token.contractAddress] ?? [:],
-                nftMetadata: metadataCache[token.id]
-              )
-            )
-          }
-        }
-      }
-      self.userVisibleNFTs = updatedUserVisibleNFTs
+      self.userVisibleNFTs = buildAssetViewModels(allVisibleUserAssets: allVisibleUserAssets)
         .optionallyFilterUnownedNFTs(
           isHidingUnownedNFTs: filters.isHidingUnownedNFTs,
           selectedAccounts: selectedAccounts
@@ -154,43 +139,56 @@ public class NFTStore: ObservableObject {
       if filters.isHidingUnownedNFTs {
         // fetch balance for all NFTs
         let allAccounts = filters.accounts.map(\.model)
-        for nft in allNFTs {
-          guard let networkForNFT = allNetworks.first(where: { $0.chainId == nft.chainId }) else {
-            continue
-          }
-          var nftBalances = nftBalancesCache[nft.contractAddress] ?? [:]
-          for account in allAccounts where account.coin == nft.coin {
-            guard let balance = await rpcService.balance(for: nft, in: account, network: networkForNFT) else {
-              continue
+        nftBalancesCache = await withTaskGroup(
+          of: [String: [String: Int]].self,
+          body: { @MainActor [nftBalancesCache, rpcService] group in
+            for nft in allNFTs { // for each NFT
+              guard let networkForNFT = allNetworks.first(where: { $0.chainId == nft.chainId }) else {
+                continue
+              }
+              group.addTask { @MainActor in
+                let updatedBalances = await withTaskGroup(
+                  of: [String: Int].self,
+                  body: { @MainActor group in
+                    for account in allAccounts where account.coin == nft.coin {
+                      group.addTask { @MainActor in
+                        let balanceForToken = await rpcService.balance(
+                          for: nft,
+                          in: account,
+                          network: networkForNFT
+                        )
+                        return [account.address: Int(balanceForToken ?? 0)]
+                      }
+                    }
+                    return await group.reduce(into: [String: Int](), { partialResult, new in
+                      partialResult.merge(with: new)
+                    })
+                  })
+                var tokenBalances = nftBalancesCache[nft.id] ?? [:]
+                tokenBalances.merge(with: updatedBalances)
+                return [nft.id: tokenBalances]
+              }
             }
-            nftBalances.merge(with: [account.address: Int(balance)])
-          }
-          nftBalancesCache[nft.contractAddress] = nftBalances
-        }
+            return await group.reduce(into: [String: [String: Int]](), { partialResult, new in
+              partialResult.merge(with: new)
+            })
+          })
       }
+      guard !Task.isCancelled else { return }
+      self.userVisibleNFTs = buildAssetViewModels(allVisibleUserAssets: allVisibleUserAssets)
+        .optionallyFilterUnownedNFTs(
+          isHidingUnownedNFTs: filters.isHidingUnownedNFTs,
+          selectedAccounts: selectedAccounts
+        )
+
       // fetch nft metadata for all NFTs
       let allMetadata = await rpcService.fetchNFTMetadata(tokens: allNFTs, ipfsApi: ipfsApi)
       for (key, value) in allMetadata { // update cached values
         metadataCache[key] = value
       }
 
       guard !Task.isCancelled else { return }
-      updatedUserVisibleNFTs.removeAll()
-      for networkAssets in allVisibleUserAssets {
-        for token in networkAssets.tokens {
-          if token.isErc721 || token.isNft {
-            updatedUserVisibleNFTs.append(
-              NFTAssetViewModel(
-                token: token,
-                network: networkAssets.network,
-                balanceForAccounts: nftBalancesCache[token.contractAddress] ?? [:],
-                nftMetadata: metadataCache[token.id]
-              )
-            )
-          }
-        }
-      }
-      self.userVisibleNFTs = updatedUserVisibleNFTs
+      self.userVisibleNFTs = buildAssetViewModels(allVisibleUserAssets: allVisibleUserAssets)
         .optionallyFilterUnownedNFTs(
           isHidingUnownedNFTs: filters.isHidingUnownedNFTs,
           selectedAccounts: selectedAccounts
@@ -207,6 +205,22 @@ public class NFTStore: ObservableObject {
     }
   }
 
+  private func buildAssetViewModels(
+    allVisibleUserAssets: [NetworkAssets]
+  ) -> [NFTAssetViewModel] {
+    allVisibleUserAssets.flatMap { networkAssets in
+      networkAssets.tokens.compactMap { token in
+        guard token.isErc721 || token.isNft else { return nil }
+        return NFTAssetViewModel(
+          token: token,
+          network: networkAssets.network,
+          balanceForAccounts: nftBalancesCache[token.id] ?? [:],
+          nftMetadata: metadataCache[token.id]
+        )
+      }
+    }
+  }
+
   @MainActor func isNFTDiscoveryEnabled() async -> Bool {
     await walletService.nftDiscoveryEnabled()
   }

diff --git a/Tests/BraveWalletTests/NFTStoreTests.swift b/Tests/BraveWalletTests/NFTStoreTests.swift
@@ -162,10 +162,10 @@ class NFTStoreTests: XCTestCase {
     XCTAssertTrue(store.userVisibleNFTs.isEmpty)  // Initial state
     store.$userVisibleNFTs
       .dropFirst()
-      .collect(2)
+      .collect(3)
       .sink { userVisibleNFTs in
         defer { userVisibleNFTsException.fulfill() }
-        XCTAssertEqual(userVisibleNFTs.count, 2) // empty nfts, populated nfts
+        XCTAssertEqual(userVisibleNFTs.count, 3) // empty nfts, populated w/ balance nfts, populated w/ metadata
         guard let lastUpdatedVisibleNFTs = userVisibleNFTs.last else {
           XCTFail("Unexpected test result")
           return
@@ -225,10 +225,10 @@ class NFTStoreTests: XCTestCase {
     let hidingUnownedExpectation = expectation(description: "update-hidingUnowned")
     store.$userVisibleNFTs
       .dropFirst()
-      .collect(2)
+      .collect(3)
       .sink { userVisibleNFTs in
         defer { hidingUnownedExpectation.fulfill() }
-        XCTAssertEqual(userVisibleNFTs.count, 2) // empty nfts, populated nfts
+        XCTAssertEqual(userVisibleNFTs.count, 3) // empty nfts, populated w/ balance nfts, populated w/ metadata
         guard let lastUpdatedVisibleNFTs = userVisibleNFTs.last else {
           XCTFail("Unexpected test result")
           return
-Original file line number
+Diff line change
@@ Expand Up @@
             .font(.body.weight(.semibold))
             .foregroundColor(Color(.separator))
         }
+        .contentShape(Rectangle())
       }
     }
@@ Expand Down @@