support bin divisions up to 65536 (#17563) (#17692)

* support bin divisions up to 65536

* add tests

(cherry picked from commit db8811eacd)

# Conflicts:
#	runtime/src/accounts_db.rs

Co-authored-by: Jeff Washington (jwash) <75863576+jeffwashington@users.noreply.github.com>

runtime/src/accounts_db.rs

@@ -29,6 +29,7 @@ use crate::{
     ancestors::Ancestors,
     append_vec::{AppendVec, StoredAccountMeta, StoredMeta, StoredMetaWriteVersion},
     contains::Contains,
+    pubkey_bins::PubkeyBinCalculator16,
     read_only_accounts_cache::ReadOnlyAccountsCache,
     sorted_storages::SortedStorages,
 };
@@ -4305,8 +4306,7 @@ impl AccountsDb {
         bin_range: &Range<usize>,
         check_hash: bool,
     ) -> Result<Vec<Vec<Vec<CalculateHashIntermediate>>>, BankHashVerificationError> {
-        let max_plus_1 = std::u8::MAX as usize + 1;
-        assert!(bins <= max_plus_1 && bins > 0);
+        let bin_calculator = PubkeyBinCalculator16::new(bins);
         assert!(bin_range.start < bins && bin_range.end <= bins && bin_range.start < bin_range.end);
         let mut time = Measure::start("scan all accounts");
         stats.num_snapshot_storage = storage.len();
@@ -4318,7 +4318,7 @@ impl AccountsDb {
              accum: &mut Vec<Vec<CalculateHashIntermediate>>,
              slot: Slot| {
                 let pubkey = loaded_account.pubkey();
-                let pubkey_to_bin_index = pubkey.as_ref()[0] as usize * bins / max_plus_1;
+                let pubkey_to_bin_index = bin_calculator.bin_from_pubkey(pubkey);
                 if !bin_range.contains(&pubkey_to_bin_index) {
                     return;
                 }
@@ -5673,24 +5673,6 @@ pub mod tests {
         SortedStorages::new(&[])
     }
 
-    #[test]
-    #[should_panic(expected = "assertion failed: bins <= max_plus_1 && bins > 0")]
-    fn test_accountsdb_scan_snapshot_stores_illegal_bins2() {
-        let mut stats = HashStats::default();
-        let bounds = Range { start: 0, end: 0 };
-
-        AccountsDb::scan_snapshot_stores(&empty_storages(), &mut stats, 257, &bounds, false)
-            .unwrap();
-    }
-    #[test]
-    #[should_panic(expected = "assertion failed: bins <= max_plus_1 && bins > 0")]
-    fn test_accountsdb_scan_snapshot_stores_illegal_bins() {
-        let mut stats = HashStats::default();
-        let bounds = Range { start: 0, end: 0 };
-
-        AccountsDb::scan_snapshot_stores(&empty_storages(), &mut stats, 0, &bounds, false).unwrap();
-    }
-
     #[test]
     #[should_panic(
         expected = "bin_range.start < bins && bin_range.end <= bins &&\\n    bin_range.start < bin_range.end"

runtime/src/lib.rs

@@ -28,6 +28,7 @@ pub mod log_collector;
 pub mod message_processor;
 mod native_loader;
 pub mod non_circulating_supply;
+mod pubkey_bins;
 mod read_only_accounts_cache;
 pub mod rent_collector;
 pub mod secondary_index;

runtime/src/pubkey_bins.rs

@@ -0,0 +1,143 @@
+use solana_sdk::pubkey::Pubkey;
+
+pub(crate) struct PubkeyBinCalculator16 {
+    // how many bits from the first 2 bytes to shift away to ignore when calculating bin
+    shift_bits: u32,
+}
+
+impl PubkeyBinCalculator16 {
+    const fn num_bits<T>() -> usize {
+        std::mem::size_of::<T>() * 8
+    }
+
+    fn log_2(x: u32) -> u32 {
+        assert!(x > 0);
+        Self::num_bits::<u32>() as u32 - x.leading_zeros() - 1
+    }
+
+    pub fn new(bins: usize) -> Self {
+        const MAX_BITS: u32 = 16;
+        assert!(bins > 0);
+        let max_plus_1 = 1 << MAX_BITS;
+        assert!(bins <= max_plus_1);
+        assert!(bins.is_power_of_two());
+        let bits = Self::log_2(bins as u32);
+        Self {
+            shift_bits: MAX_BITS - bits,
+        }
+    }
+
+    pub fn bin_from_pubkey(&self, pubkey: &Pubkey) -> usize {
+        let as_ref = pubkey.as_ref();
+        ((as_ref[0] as usize * 256 + as_ref[1] as usize) as usize) >> self.shift_bits
+    }
+}
+
+#[cfg(test)]
+pub mod tests {
+    use super::*;
+
+    #[test]
+    fn test_pubkey_bins_log2() {
+        assert_eq!(PubkeyBinCalculator16::num_bits::<u8>(), 8);
+        assert_eq!(PubkeyBinCalculator16::num_bits::<u32>(), 32);
+        for i in 0..32 {
+            assert_eq!(PubkeyBinCalculator16::log_2(2u32.pow(i)), i);
+        }
+    }
+
+    #[test]
+    fn test_pubkey_bins() {
+        for i in 0..=16 {
+            let bins = 2u32.pow(i);
+            let calc = PubkeyBinCalculator16::new(bins as usize);
+            assert_eq!(calc.shift_bits, 16 - i, "i: {}", i);
+        }
+    }
+
+    #[test]
+    fn test_pubkey_bins_pubkeys() {
+        let mut pk = Pubkey::new(&[0; 32]);
+        for i in 0..=8 {
+            let bins = 2usize.pow(i);
+            let calc = PubkeyBinCalculator16::new(bins);
+
+            let shift_bits = calc.shift_bits - 8; // we are only dealing with first byte
+
+            pk.as_mut()[0] = 0;
+            assert_eq!(0, calc.bin_from_pubkey(&pk));
+            pk.as_mut()[0] = 0xff;
+            assert_eq!(bins - 1, calc.bin_from_pubkey(&pk));
+
+            for bin in 0..bins {
+                pk.as_mut()[0] = (bin << shift_bits) as u8;
+                assert_eq!(
+                    bin,
+                    calc.bin_from_pubkey(&pk),
+                    "bin: {}/{}, bits: {}, val: {}",
+                    bin,
+                    bins,
+                    shift_bits,
+                    pk.as_ref()[0]
+                );
+                if bin > 0 {
+                    pk.as_mut()[0] = ((bin << shift_bits) - 1) as u8;
+                    assert_eq!(bin - 1, calc.bin_from_pubkey(&pk));
+                }
+            }
+        }
+
+        for i in 9..=16 {
+            let mut pk = Pubkey::new(&[0; 32]);
+            let bins = 2usize.pow(i);
+            let calc = PubkeyBinCalculator16::new(bins);
+
+            let shift_bits = calc.shift_bits;
+
+            pk.as_mut()[1] = 0;
+            assert_eq!(0, calc.bin_from_pubkey(&pk));
+            pk.as_mut()[0] = 0xff;
+            pk.as_mut()[1] = 0xff;
+            assert_eq!(bins - 1, calc.bin_from_pubkey(&pk));
+
+            let mut pk = Pubkey::new(&[0; 32]);
+            for bin in 0..bins {
+                let mut target = (bin << shift_bits) as u16;
+                pk.as_mut()[0] = (target / 256) as u8;
+                pk.as_mut()[1] = (target % 256) as u8;
+                assert_eq!(
+                    bin,
+                    calc.bin_from_pubkey(&pk),
+                    "bin: {}/{}, bits: {}, val: {}",
+                    bin,
+                    bins,
+                    shift_bits,
+                    pk.as_ref()[0]
+                );
+                if bin > 0 {
+                    target -= 1;
+                    pk.as_mut()[0] = (target / 256) as u8;
+                    pk.as_mut()[1] = (target % 256) as u8;
+                    assert_eq!(bin - 1, calc.bin_from_pubkey(&pk));
+                }
+            }
+        }
+    }
+
+    #[test]
+    #[should_panic(expected = "bins.is_power_of_two()")]
+    fn test_pubkey_bins_illegal_bins3() {
+        PubkeyBinCalculator16::new(3);
+    }
+
+    #[test]
+    #[should_panic(expected = "bins <= max_plus_1")]
+    fn test_pubkey_bins_illegal_bins2() {
+        PubkeyBinCalculator16::new(65537);
+    }
+    #[test]
+    #[should_panic(expected = "bins > 0")]
+    fn test_pubkey_bins_illegal_bins() {
+        PubkeyBinCalculator16::new(0);
+    }
+}