diff --git a/rpc/src/rpc.rs b/rpc/src/rpc.rs index 82eda9489..7bde6b837 100644 --- a/rpc/src/rpc.rs +++ b/rpc/src/rpc.rs @@ -100,8 +100,8 @@ use { }, std::{ any::type_name, - cmp::{max, min}, - collections::{HashMap, HashSet}, + cmp::{max, min, Reverse}, + collections::{BinaryHeap, HashMap, HashSet}, convert::TryFrom, net::SocketAddr, str::FromStr, @@ -1861,36 +1861,39 @@ 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![])? + + let mut token_balances = + BinaryHeap::>::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::::unpack(account.data()) + .map(|account| account.base.amount) + .unwrap_or(0); + + let new_entry = (amount, address); + if token_balances.len() >= NUM_LARGEST_ACCOUNTS { + let Reverse(entry) = token_balances + .peek() + .expect("BinaryHeap::peek should succeed when len > 0"); + if *entry >= new_entry { + continue; + } + token_balances.pop(); + } + token_balances.push(Reverse(new_entry)); + } + + let token_balances = token_balances + .into_sorted_vec() .into_iter() - .map(|(address, account)| { - let amount = StateWithExtensions::::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); - - 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 largest_token_balances = largest_token_balances - .iter() - .map(|(address, amount)| RpcTokenAccountBalance { + .map(|Reverse((amount, address))| RpcTokenAccountBalance { address: address.to_string(), - amount: token_amount_to_ui_amount(*amount, decimals), + amount: token_amount_to_ui_amount(amount, decimals), }) .collect(); - Ok(new_response(&bank, largest_token_balances)) + + Ok(new_response(&bank, token_balances)) } pub fn get_token_accounts_by_owner(