rpc: optimize getTokenLargestAccounts

rpc/src/rpc.rs

@@ -101,7 +101,7 @@ use {
     std::{
         any::type_name,
         cmp::{max, min},
-        collections::{HashMap, HashSet},
+        collections::{BinaryHeap, HashMap, HashSet},
         convert::TryFrom,
         net::SocketAddr,
         str::FromStr,
@@ -1861,36 +1861,59 @@ impl JsonRpcRequestProcessor {
                 "Invalid param: not a Token mint".to_string(),
             ));
         }
-        let mut token_balances: Vec<_> = self
-            .get_filtered_spl_token_accounts_by_mint(&bank, &mint_owner, mint, vec![])?
-            .into_iter()
-            .map(|(address, account)| {
-                let amount = StateWithExtensions::<TokenAccount>::unpack(account.data())
-                    .map(|account| account.base.amount)
-                    .unwrap_or(0);
-                (address, amount)
-            })
-            .collect();
 
-        let sort_largest = |a: &(_, u64), b: &(_, u64)| b.1.cmp(&a.1);
+        #[derive(Debug, PartialEq, Eq)]
+        struct TokenAccountBalance {
+            address: Pubkey,
+            amount: u64,
+        }
+        impl PartialOrd for TokenAccountBalance {
+            fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
+                Some(self.amount.cmp(&other.amount).reverse())
+            }
+        }
+        impl Ord for TokenAccountBalance {
+            fn cmp(&self, other: &Self) -> std::cmp::Ordering {
+                self.amount.cmp(&other.amount).reverse()
+            }
+        }
 
-        let largest_token_balances = if token_balances.len() > NUM_LARGEST_ACCOUNTS {
-            token_balances
-                .select_nth_unstable_by(NUM_LARGEST_ACCOUNTS, sort_largest)
-                .0
-        } else {
-            token_balances.as_mut_slice()
-        };
-        largest_token_balances.sort_unstable_by(sort_largest);
+        let mut token_balances =
+            BinaryHeap::<TokenAccountBalance>::with_capacity(NUM_LARGEST_ACCOUNTS);
+        for (address, account) in
+            self.get_filtered_spl_token_accounts_by_mint(&bank, &mint_owner, mint, vec![])?
+        {
+            let amount = StateWithExtensions::<TokenAccount>::unpack(account.data())
+                .map(|account| account.base.amount)
+                .unwrap_or(0);
 
-        let largest_token_balances = largest_token_balances
-            .iter()
-            .map(|(address, amount)| RpcTokenAccountBalance {
-                address: address.to_string(),
-                amount: token_amount_to_ui_amount(*amount, decimals),
+            let requires_insert = if token_balances.len() < NUM_LARGEST_ACCOUNTS {
+                true
+            } else {
+                if let Some(peek) = token_balances.peek() {
+                    peek.amount < amount
+                } else {
+                    false
+                }
+            };
+
+            if requires_insert {
+                if token_balances.len() >= NUM_LARGEST_ACCOUNTS {
+                    token_balances.pop();
+                }
+                token_balances.push(TokenAccountBalance { address, amount });
+            }
+        }
+        let token_balances = token_balances
+            .into_sorted_vec()
+            .into_iter()
+            .map(|token| RpcTokenAccountBalance {
+                address: token.address.to_string(),
+                amount: token_amount_to_ui_amount(token.amount, decimals),
             })
             .collect();
-        Ok(new_response(&bank, largest_token_balances))
+
+        Ok(new_response(&bank, token_balances))
     }
 
     pub fn get_token_accounts_by_owner(