From 4ed0fcdde6bfc25de361e205a376f81be4bd88e7 Mon Sep 17 00:00:00 2001
From: "mergify[bot]" <37929162+mergify[bot]@users.noreply.github.com>
Date: Wed, 16 Jun 2021 20:48:18 +0000
Subject: [PATCH] avoid unnecessary empty arrays when binning (#17944) (#18011)

(cherry picked from commit eee5414c64fcd93a5642c9c0e6f03f5c8bdd076f)

Co-authored-by: Jeff Washington (jwash) <75863576+jeffwashington@users.noreply.github.com>
---
 runtime/src/accounts_db.rs | 36 +++++++++++++++++++++---------------
 1 file changed, 21 insertions(+), 15 deletions(-)

diff --git a/runtime/src/accounts_db.rs b/runtime/src/accounts_db.rs
index d4b25df2cf..b2117fefbd 100644
--- a/runtime/src/accounts_db.rs
+++ b/runtime/src/accounts_db.rs
@@ -4563,6 +4563,7 @@ impl AccountsDb {
         let mut time = Measure::start("scan all accounts");
         stats.num_snapshot_storage = storage.slot_count();
         let mismatch_found = AtomicU64::new(0);
+        let range = bin_range.end - bin_range.start;
 
         let result: Vec<Vec<Vec<CalculateHashIntermediate>>> = Self::scan_account_storage_no_bank(
             accounts_cache_and_ancestors,
@@ -4571,11 +4572,14 @@ impl AccountsDb {
              accum: &mut Vec<Vec<CalculateHashIntermediate>>,
              slot: Slot| {
                 let pubkey = loaded_account.pubkey();
-                let pubkey_to_bin_index = bin_calculator.bin_from_pubkey(pubkey);
+                let mut pubkey_to_bin_index = bin_calculator.bin_from_pubkey(pubkey);
                 if !bin_range.contains(&pubkey_to_bin_index) {
                     return;
                 }
 
+                // when we are scanning with bin ranges, we don't need to use exact bin numbers. Subtract to make first bin we care about at index 0.
+                pubkey_to_bin_index -= bin_range.start;
+
                 let raw_lamports = loaded_account.lamports();
                 let zero_raw_lamports = raw_lamports == 0;
                 let balance = if zero_raw_lamports {
@@ -4603,7 +4607,7 @@ impl AccountsDb {
 
                 let max = accum.len();
                 if max == 0 {
-                    accum.extend(vec![Vec::new(); bins]);
+                    accum.extend(vec![Vec::new(); range]);
                 }
                 accum[pubkey_to_bin_index].push(source_item);
             },
@@ -6157,18 +6161,19 @@ pub mod tests {
 
         // just the first bin of 2
         let bins = 2;
+        let half_bins = bins / 2;
         let result = AccountsDb::scan_snapshot_stores(
             &get_storage_refs(&storages),
             &mut stats,
             bins,
             &Range {
                 start: 0,
-                end: bins / 2,
+                end: half_bins,
             },
             false,
         )
         .unwrap();
-        let mut expected = vec![Vec::new(); bins];
+        let mut expected = vec![Vec::new(); half_bins];
         expected[0].push(raw_expected[0].clone());
         expected[0].push(raw_expected[1].clone());
         assert_eq!(result, vec![expected]);
@@ -6186,14 +6191,15 @@ pub mod tests {
         )
         .unwrap();
 
-        let mut expected = vec![Vec::new(); bins];
-        expected[bins - 1].push(raw_expected[2].clone());
-        expected[bins - 1].push(raw_expected[3].clone());
+        let mut expected = vec![Vec::new(); half_bins];
+        let starting_bin_index = 0;
+        expected[starting_bin_index].push(raw_expected[2].clone());
+        expected[starting_bin_index].push(raw_expected[3].clone());
         assert_eq!(result, vec![expected]);
 
         // 1 bin at a time of 4
         let bins = 4;
-        for bin in 0..bins {
+        for (bin, expected_item) in raw_expected.iter().enumerate().take(bins) {
             let result = AccountsDb::scan_snapshot_stores(
                 &get_storage_refs(&storages),
                 &mut stats,
@@ -6205,8 +6211,8 @@ pub mod tests {
                 false,
             )
             .unwrap();
-            let mut expected = vec![Vec::new(); bins];
-            expected[bin].push(raw_expected[bin].clone());
+            let mut expected = vec![Vec::new(); 1];
+            expected[0].push(expected_item.clone());
             assert_eq!(result, vec![expected]);
         }
 
@@ -6226,8 +6232,8 @@ pub mod tests {
             .unwrap();
             let mut expected = vec![];
             if let Some(index) = bin_locations.iter().position(|&r| r == bin) {
-                expected = vec![Vec::new(); bins];
-                expected[bin].push(raw_expected[index].clone());
+                expected = vec![Vec::new(); 1];
+                expected[0].push(raw_expected[index].clone());
             }
             assert_eq!(result, vec![expected]);
         }
@@ -6260,9 +6266,9 @@ pub mod tests {
         .unwrap();
         assert_eq!(result.len(), 2); // 2 chunks
         assert_eq!(result[0].len(), 0); // nothing found in first slots
-        let mut expected = vec![Vec::new(); bins];
-        expected[127].push(raw_expected[1].clone());
-        assert_eq!(result[1].len(), bins);
+        let mut expected = vec![Vec::new(); 1];
+        expected[0].push(raw_expected[1].clone());
+        assert_eq!(result[1].len(), 1);
         assert_eq!(result[1], expected);
     }