during range, stop scanning earlier than currently scanning (#18896)

* during range, stop scanning earlier than currently scanning

* refactor for better testing

* clean up a different edge case

runtime/src/accounts_index.rs

@@ -541,14 +541,37 @@ impl<'a, T> AccountsIndexIterator<'a, T> {
         }
     }
 
+    fn bin_from_bound(bound: &Bound<Pubkey>, unbounded_bin: usize) -> usize {
+        match bound {
+            Bound::Included(bound) | Bound::Excluded(bound) => get_bin_pubkey(bound),
+            Bound::Unbounded => unbounded_bin,
+        }
+    }
+
     fn start_bin(&self) -> usize {
         // start in bin where 'start_bound' would exist
-        match &self.start_bound {
-            Bound::Included(start_bound) | Bound::Excluded(start_bound) => {
-                get_bin_pubkey(start_bound)
-            }
-            Bound::Unbounded => 0,
-        }
+        Self::bin_from_bound(&self.start_bound, 0)
+    }
+
+    fn end_bin_inclusive(&self) -> usize {
+        // end in bin where 'end_bound' would exist
+        Self::bin_from_bound(&self.end_bound, usize::MAX)
+    }
+
+    fn bin_start_and_range(&self) -> (usize, usize) {
+        let start_bin = self.start_bin();
+        // calculate the max range of bins to look in
+        let end_bin_inclusive = self.end_bin_inclusive();
+        let bin_range = if start_bin > end_bin_inclusive {
+            0 // empty range
+        } else if end_bin_inclusive == usize::MAX {
+            usize::MAX
+        } else {
+            // the range is end_inclusive + 1 - start
+            // end_inclusive could be usize::MAX already if no bound was specified
+            end_bin_inclusive.saturating_add(1) - start_bin
+        };
+        (start_bin, bin_range)
     }
 
     pub fn new<R>(account_maps: &'a LockMapTypeSlice<T>, range: Option<R>) -> Self
@@ -576,10 +599,9 @@ impl<'a, T: 'static + Clone> Iterator for AccountsIndexIterator<'a, T> {
         if self.is_finished {
             return None;
         }
-
-        let start_bin = self.start_bin();
+        let (start_bin, bin_range) = self.bin_start_and_range();
         let mut chunk: Vec<(Pubkey, AccountMapEntry<T>)> = Vec::with_capacity(ITER_BATCH_SIZE);
-        'outer: for i in self.account_maps.iter().skip(start_bin) {
+        'outer: for i in self.account_maps.iter().skip(start_bin).take(bin_range) {
             for (pubkey, account_map_entry) in
                 i.read().unwrap().range((self.start_bound, self.end_bound))
             {
@@ -3816,31 +3838,74 @@ pub mod tests {
     }
 
     #[test]
-    fn test_start_bin() {
+    fn test_bin_start_and_range() {
+        let index = AccountsIndex::<bool>::default();
+        let iter = AccountsIndexIterator::new(&index.account_maps, None::<RangeInclusive<Pubkey>>);
+        assert_eq!((0, usize::MAX), iter.bin_start_and_range());
+
+        let key_0 = Pubkey::new(&[0; 32]);
+        let key_ff = Pubkey::new(&[0xff; 32]);
+
+        let iter = AccountsIndexIterator::new(
+            &index.account_maps,
+            Some(RangeInclusive::new(key_0, key_ff)),
+        );
+        assert_eq!((0, BINS), iter.bin_start_and_range());
+        let iter = AccountsIndexIterator::new(
+            &index.account_maps,
+            Some(RangeInclusive::new(key_ff, key_0)),
+        );
+        assert_eq!((BINS - 1, 0), iter.bin_start_and_range());
+        let iter =
+            AccountsIndexIterator::new(&index.account_maps, Some((Included(key_0), Unbounded)));
+        assert_eq!((0, usize::MAX), iter.bin_start_and_range());
+        let iter =
+            AccountsIndexIterator::new(&index.account_maps, Some((Included(key_ff), Unbounded)));
+        assert_eq!((BINS - 1, usize::MAX), iter.bin_start_and_range());
+
+        assert_eq!(
+            (0..2)
+                .into_iter()
+                .skip(1)
+                .take(usize::MAX)
+                .collect::<Vec<_>>(),
+            vec![1]
+        );
+    }
+
+    #[test]
+    fn test_start_end_bin() {
         let index = AccountsIndex::<bool>::default();
         let iter = AccountsIndexIterator::new(&index.account_maps, None::<RangeInclusive<Pubkey>>);
         assert_eq!(iter.start_bin(), 0); // no range, so 0
+        assert_eq!(iter.end_bin_inclusive(), usize::MAX); // no range, so max
 
         let key = Pubkey::new(&[0; 32]);
         let iter =
             AccountsIndexIterator::new(&index.account_maps, Some(RangeInclusive::new(key, key)));
         assert_eq!(iter.start_bin(), 0); // start at pubkey 0, so 0
+        assert_eq!(iter.end_bin_inclusive(), 0); // end at pubkey 0, so 0
         let iter =
             AccountsIndexIterator::new(&index.account_maps, Some((Included(key), Excluded(key))));
         assert_eq!(iter.start_bin(), 0); // start at pubkey 0, so 0
+        assert_eq!(iter.end_bin_inclusive(), 0); // end at pubkey 0, so 0
         let iter =
             AccountsIndexIterator::new(&index.account_maps, Some((Excluded(key), Excluded(key))));
         assert_eq!(iter.start_bin(), 0); // start at pubkey 0, so 0
+        assert_eq!(iter.end_bin_inclusive(), 0); // end at pubkey 0, so 0
 
         let key = Pubkey::new(&[0xff; 32]);
         let iter =
             AccountsIndexIterator::new(&index.account_maps, Some(RangeInclusive::new(key, key)));
         assert_eq!(iter.start_bin(), BINS - 1); // start at highest possible pubkey, so BINS - 1
+        assert_eq!(iter.end_bin_inclusive(), BINS - 1);
         let iter =
             AccountsIndexIterator::new(&index.account_maps, Some((Included(key), Excluded(key))));
         assert_eq!(iter.start_bin(), BINS - 1); // start at highest possible pubkey, so BINS - 1
+        assert_eq!(iter.end_bin_inclusive(), BINS - 1);
         let iter =
             AccountsIndexIterator::new(&index.account_maps, Some((Excluded(key), Excluded(key))));
         assert_eq!(iter.start_bin(), BINS - 1); // start at highest possible pubkey, so BINS - 1
+        assert_eq!(iter.end_bin_inclusive(), BINS - 1);
     }
 }