solana_accounts_db/
pubkey_bins.rs1use solana_pubkey::Pubkey;
2
3#[derive(Debug)]
4pub struct PubkeyBinCalculator24 {
5 shift_bits: u32,
7}
8
9impl PubkeyBinCalculator24 {
10 const fn num_bits<T>() -> usize {
11 std::mem::size_of::<T>() * 8
12 }
13
14 pub(crate) fn log_2(x: u32) -> u32 {
15 assert!(x > 0);
16 Self::num_bits::<u32>() as u32 - x.leading_zeros() - 1
17 }
18
19 pub fn new(bins: usize) -> Self {
20 const MAX_BITS: u32 = 24;
21 assert!(bins > 0);
22 let max_plus_1 = 1 << MAX_BITS;
23 assert!(bins <= max_plus_1);
24 assert!(bins.is_power_of_two());
25 let bits = Self::log_2(bins as u32);
26 Self {
27 shift_bits: MAX_BITS - bits,
28 }
29 }
30
31 #[inline]
32 pub fn bin_from_pubkey(&self, pubkey: &Pubkey) -> usize {
33 let as_ref = pubkey.as_ref();
34 ((as_ref[0] as usize) << 16 | (as_ref[1] as usize) << 8 | (as_ref[2] as usize))
35 >> self.shift_bits
36 }
37
38 #[cfg(test)]
39 pub(crate) fn lowest_pubkey_from_bin(&self, mut bin: usize, bins: usize) -> Pubkey {
40 assert!(bin < bins);
41 bin <<= self.shift_bits;
42 let mut pubkey = Pubkey::from([0; 32]);
43 pubkey.as_mut()[0] = ((bin / 256 / 256) & 0xff) as u8;
44 pubkey.as_mut()[1] = ((bin / 256) & 0xff) as u8;
45 pubkey.as_mut()[2] = (bin & 0xff) as u8;
46 pubkey
47 }
48}
49
50#[cfg(test)]
51pub mod tests {
52 use super::*;
53
54 #[test]
55 fn test_pubkey_bins_log2() {
56 assert_eq!(PubkeyBinCalculator24::num_bits::<u8>(), 8);
57 assert_eq!(PubkeyBinCalculator24::num_bits::<u32>(), 32);
58 for i in 0..32 {
59 assert_eq!(PubkeyBinCalculator24::log_2(2u32.pow(i)), i);
60 }
61 }
62
63 #[test]
64 fn test_pubkey_bins() {
65 for i in 0..=24 {
66 let bins = 2u32.pow(i);
67 let calc = PubkeyBinCalculator24::new(bins as usize);
68 assert_eq!(calc.shift_bits, 24 - i, "i: {i}");
69 for bin in 0..bins {
70 assert_eq!(
71 bin as usize,
72 calc.bin_from_pubkey(&calc.lowest_pubkey_from_bin(bin as usize, bins as usize))
73 );
74 }
75 }
76 }
77
78 #[test]
79 fn test_pubkey_bins_pubkeys() {
80 let mut pk = Pubkey::from([0; 32]);
81 for i in 0..=8 {
82 let bins = 2usize.pow(i);
83 let calc = PubkeyBinCalculator24::new(bins);
84
85 let shift_bits = calc.shift_bits - 16; pk.as_mut()[0] = 0;
88 assert_eq!(0, calc.bin_from_pubkey(&pk));
89 pk.as_mut()[0] = 0xff;
90 assert_eq!(bins - 1, calc.bin_from_pubkey(&pk));
91
92 for bin in 0..bins {
93 pk.as_mut()[0] = (bin << shift_bits) as u8;
94 assert_eq!(
95 bin,
96 calc.bin_from_pubkey(&pk),
97 "bin: {}/{}, shift_bits: {}, val: {}",
98 bin,
99 bins,
100 shift_bits,
101 pk.as_ref()[0]
102 );
103 if bin > 0 {
104 pk.as_mut()[0] = ((bin << shift_bits) - 1) as u8;
105 assert_eq!(bin - 1, calc.bin_from_pubkey(&pk));
106 }
107 }
108 }
109
110 for i in 9..=16 {
111 let mut pk = Pubkey::from([0; 32]);
112 let bins = 2usize.pow(i);
113 let calc = PubkeyBinCalculator24::new(bins);
114
115 let shift_bits = calc.shift_bits - 8;
116
117 pk.as_mut()[1] = 0;
118 assert_eq!(0, calc.bin_from_pubkey(&pk));
119 pk.as_mut()[0] = 0xff;
120 pk.as_mut()[1] = 0xff;
121 assert_eq!(bins - 1, calc.bin_from_pubkey(&pk));
122
123 let mut pk = Pubkey::from([0; 32]);
124 for bin in 0..bins {
125 let mut target = (bin << shift_bits) as u16;
126 pk.as_mut()[0] = (target / 256) as u8;
127 pk.as_mut()[1] = (target % 256) as u8;
128 assert_eq!(
129 bin,
130 calc.bin_from_pubkey(&pk),
131 "bin: {}/{}, shift_bits: {}, val: {}",
132 bin,
133 bins,
134 shift_bits,
135 pk.as_ref()[0]
136 );
137 if bin > 0 {
138 target -= 1;
139 pk.as_mut()[0] = (target / 256) as u8;
140 pk.as_mut()[1] = (target % 256) as u8;
141 assert_eq!(bin - 1, calc.bin_from_pubkey(&pk));
142 }
143 }
144 }
145
146 for i in 17..=24 {
147 let mut pk = Pubkey::from([0; 32]);
148 let bins = 2usize.pow(i);
149 let calc = PubkeyBinCalculator24::new(bins);
150
151 let shift_bits = calc.shift_bits;
152
153 pk.as_mut()[1] = 0;
154 assert_eq!(0, calc.bin_from_pubkey(&pk));
155 pk.as_mut()[0] = 0xff;
156 pk.as_mut()[1] = 0xff;
157 pk.as_mut()[2] = 0xff;
158 assert_eq!(bins - 1, calc.bin_from_pubkey(&pk));
159
160 let mut pk = Pubkey::from([0; 32]);
161 for bin in 0..bins {
162 let mut target = (bin << shift_bits) as u32;
163 pk.as_mut()[0] = (target / 256 / 256) as u8;
164 pk.as_mut()[1] = ((target / 256) % 256) as u8;
165 pk.as_mut()[2] = (target % 256) as u8;
166 assert_eq!(
167 bin,
168 calc.bin_from_pubkey(&pk),
169 "bin: {}/{}, shift_bits: {}, val: {:?}",
170 bin,
171 bins,
172 shift_bits,
173 &pk.as_ref()[0..3],
174 );
175 if bin > 0 {
176 target -= 1;
177 pk.as_mut()[0] = (target / 256 / 256) as u8;
178 pk.as_mut()[1] = ((target / 256) % 256) as u8;
179 pk.as_mut()[2] = (target % 256) as u8;
180 assert_eq!(bin - 1, calc.bin_from_pubkey(&pk));
181 }
182 }
183 }
184 }
185
186 #[test]
187 #[should_panic(expected = "bins.is_power_of_two()")]
188 fn test_pubkey_bins_illegal_bins3() {
189 PubkeyBinCalculator24::new(3);
190 }
191
192 #[test]
193 #[should_panic(expected = "bins <= max_plus_1")]
194 fn test_pubkey_bins_illegal_bins2() {
195 PubkeyBinCalculator24::new(65536 * 256 + 1);
196 }
197 #[test]
198 #[should_panic(expected = "bins > 0")]
199 fn test_pubkey_bins_illegal_bins() {
200 PubkeyBinCalculator24::new(0);
201 }
202}