From 8ad125d0c0688aaf2b62bb95b535ff988ed7f9ac Mon Sep 17 00:00:00 2001
From: Kirill Fomichev <fanatid@ya.ru>
Date: Tue, 27 Feb 2024 02:08:29 -0500
Subject: [PATCH] rpc: optimize `getTokenLargestAccounts` (#35315)

* rpc: optimize `getTokenLargestAccounts`

* use tuple instead of struct

* untuple

Co-authored-by: Tyera <teulberg@gmail.com>

---------

Co-authored-by: Tyera <teulberg@gmail.com>
---
 rpc/src/rpc.rs | 55 ++++++++++++++++++++++++++------------------------
 1 file changed, 29 insertions(+), 26 deletions(-)

diff --git a/rpc/src/rpc.rs b/rpc/src/rpc.rs
index 82eda9489ef247..7bde6b837f2a13 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![])?
-            .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);
+        let mut token_balances =
+            BinaryHeap::<Reverse<(u64, Pubkey)>>::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 = 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 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 largest_token_balances = largest_token_balances
-            .iter()
-            .map(|(address, amount)| RpcTokenAccountBalance {
+        let token_balances = token_balances
+            .into_sorted_vec()
+            .into_iter()
+            .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(