1use {
2 crate::accounts_db::SnapshotStorage,
3 log::*,
4 safecoin_measure::measure::Measure,
5 solana_sdk::clock::Slot,
6 std::ops::{Bound, Range, RangeBounds},
7};
8
9pub struct SortedStorages<'a> {
11 range: Range<Slot>,
13 storages: Vec<Option<&'a SnapshotStorage>>,
15 slot_count: usize,
16 storage_count: usize,
17}
18
19impl<'a> SortedStorages<'a> {
20 pub fn empty() -> Self {
22 SortedStorages {
23 range: Range::default(),
24 storages: Vec::default(),
25 slot_count: 0,
26 storage_count: 0,
27 }
28 }
29
30 pub fn iter_range<R>(&'a self, range: R) -> SortedStoragesIter<'a>
32 where
33 R: RangeBounds<Slot>,
34 {
35 SortedStoragesIter::new(self, range)
36 }
37
38 fn get(&self, slot: Slot) -> Option<&SnapshotStorage> {
39 if !self.range.contains(&slot) {
40 None
41 } else {
42 let index = (slot - self.range.start) as usize;
43 self.storages[index]
44 }
45 }
46
47 pub fn range_width(&self) -> Slot {
48 self.range.end - self.range.start
49 }
50
51 pub fn range(&self) -> &Range<Slot> {
52 &self.range
53 }
54
55 pub fn max_slot_inclusive(&self) -> Slot {
56 self.range.end.saturating_sub(1)
57 }
58
59 pub fn slot_count(&self) -> usize {
60 self.slot_count
61 }
62
63 pub fn storage_count(&self) -> usize {
64 self.storage_count
65 }
66
67 pub fn new(source: &'a [SnapshotStorage]) -> Self {
71 let slots = source.iter().map(|storages| {
72 let first = storages.first();
73 assert!(first.is_some(), "SnapshotStorage.is_empty()");
74 let storage = first.unwrap();
75 storage.slot() });
77 Self::new_with_slots(source.iter().zip(slots.into_iter()), None, None)
78 }
79
80 pub fn new_with_slots(
84 source: impl Iterator<Item = (&'a SnapshotStorage, Slot)> + Clone,
85 min_slot: Option<Slot>,
87 max_slot_inclusive: Option<Slot>,
92 ) -> Self {
93 let mut min = Slot::MAX;
94 let mut max = Slot::MIN;
95 let mut adjust_min_max = |slot| {
96 min = std::cmp::min(slot, min);
97 max = std::cmp::max(slot + 1, max);
98 };
99 if let Some(slot) = min_slot {
101 adjust_min_max(slot);
102 }
103 if let Some(slot) = max_slot_inclusive {
104 adjust_min_max(slot);
105 }
106
107 let mut slot_count = 0;
108 let mut time = Measure::start("get slot");
109 let source_ = source.clone();
110 let mut storage_count = 0;
111 source_.for_each(|(storages, slot)| {
112 storage_count += storages.len();
113 slot_count += 1;
114 adjust_min_max(slot);
115 });
116 time.stop();
117 let mut time2 = Measure::start("sort");
118 let range;
119 let mut storages;
120 if min > max {
121 range = Range::default();
122 storages = vec![];
123 } else {
124 range = Range {
125 start: min,
126 end: max,
127 };
128 let len = max - min;
129 storages = vec![None; len as usize];
130 source.for_each(|(original_storages, slot)| {
131 let index = (slot - min) as usize;
132 assert!(storages[index].is_none(), "slots are not unique"); storages[index] = Some(original_storages);
134 });
135 }
136 time2.stop();
137 debug!("SortedStorages, times: {}, {}", time.as_us(), time2.as_us());
138 Self {
139 range,
140 storages,
141 slot_count,
142 storage_count,
143 }
144 }
145}
146
147pub struct SortedStoragesIter<'a> {
151 range: Range<Slot>,
153 storages: &'a SortedStorages<'a>,
155 next_slot: Slot,
157}
158
159impl<'a> Iterator for SortedStoragesIter<'a> {
160 type Item = (Slot, Option<&'a SnapshotStorage>);
161
162 fn next(&mut self) -> Option<Self::Item> {
163 let slot = self.next_slot;
164 if slot < self.range.end {
165 self.next_slot += 1;
167 Some((slot, self.storages.get(slot)))
168 } else {
169 None
171 }
172 }
173}
174
175impl<'a> SortedStoragesIter<'a> {
176 pub fn new<R: RangeBounds<Slot>>(
177 storages: &'a SortedStorages<'a>,
178 range: R,
179 ) -> SortedStoragesIter<'a> {
180 let storage_range = storages.range();
181 let next_slot = match range.start_bound() {
182 Bound::Unbounded => {
183 storage_range.start }
185 Bound::Included(x) => *x,
186 Bound::Excluded(x) => *x + 1, };
188 let end_exclusive_slot = match range.end_bound() {
189 Bound::Unbounded => {
190 storage_range.end }
192 Bound::Included(x) => *x + 1, Bound::Excluded(x) => *x,
194 };
195 let range = next_slot..end_exclusive_slot;
199 SortedStoragesIter {
200 range,
201 storages,
202 next_slot,
203 }
204 }
205}
206
207#[cfg(test)]
208pub mod tests {
209 use super::*;
210 impl<'a> SortedStorages<'a> {
211 pub fn new_debug(source: &[(&'a SnapshotStorage, Slot)], min: Slot, len: usize) -> Self {
212 let mut storages = vec![None; len];
213 let range = Range {
214 start: min,
215 end: min + len as Slot,
216 };
217 let slot_count = source.len();
218 for (storage, slot) in source {
219 storages[*slot as usize] = Some(*storage);
220 }
221
222 Self {
223 range,
224 storages,
225 slot_count,
226 storage_count: 0,
227 }
228 }
229
230 pub fn new_for_tests(storages: &[&'a SnapshotStorage], slots: &[Slot]) -> Self {
231 assert_eq!(storages.len(), slots.len());
232 SortedStorages::new_with_slots(
233 storages.iter().cloned().zip(slots.iter().cloned()),
234 None,
235 None,
236 )
237 }
238 }
239
240 #[test]
241 fn test_sorted_storages_range_iter() {
242 let storages = SortedStorages::empty();
243 let check = |(slot, storages): (Slot, Option<&SnapshotStorage>)| {
244 assert!(storages.is_none());
245 slot
246 };
247 assert_eq!(
248 (0..5).into_iter().collect::<Vec<_>>(),
249 storages.iter_range(..5).map(check).collect::<Vec<_>>()
250 );
251 assert_eq!(
252 (1..5).into_iter().collect::<Vec<_>>(),
253 storages.iter_range(1..5).map(check).collect::<Vec<_>>()
254 );
255 assert_eq!(
256 (0..0).into_iter().collect::<Vec<_>>(),
257 storages.iter_range(..).map(check).collect::<Vec<_>>()
258 );
259 assert_eq!(
260 (0..0).into_iter().collect::<Vec<_>>(),
261 storages.iter_range(1..).map(check).collect::<Vec<_>>()
262 );
263
264 let s1 = Vec::new();
266 let storages = SortedStorages::new_for_tests(&[&s1], &[3]);
267 let check = |(slot, storages): (Slot, Option<&SnapshotStorage>)| {
268 assert!(
269 (slot != 3) ^ storages.is_some(),
270 "slot: {}, storages: {:?}",
271 slot,
272 storages
273 );
274 slot
275 };
276 for start in 0..5 {
277 for end in 0..5 {
278 assert_eq!(
279 (start..end).into_iter().collect::<Vec<_>>(),
280 storages
281 .iter_range(start..end)
282 .map(check)
283 .collect::<Vec<_>>()
284 );
285 }
286 }
287 assert_eq!(
288 (3..5).into_iter().collect::<Vec<_>>(),
289 storages.iter_range(..5).map(check).collect::<Vec<_>>()
290 );
291 assert_eq!(
292 (1..=3).into_iter().collect::<Vec<_>>(),
293 storages.iter_range(1..).map(check).collect::<Vec<_>>()
294 );
295 assert_eq!(
296 (3..=3).into_iter().collect::<Vec<_>>(),
297 storages.iter_range(..).map(check).collect::<Vec<_>>()
298 );
299
300 let s2 = Vec::with_capacity(2);
302 let s4 = Vec::with_capacity(4);
303 let storages = SortedStorages::new_for_tests(&[&s2, &s4], &[2, 4]);
304 let check = |(slot, storages): (Slot, Option<&SnapshotStorage>)| {
305 assert!(
306 (slot != 2 && slot != 4)
307 ^ storages
308 .map(|storages| storages.capacity() == (slot as usize))
309 .unwrap_or(false),
310 "slot: {}, storages: {:?}",
311 slot,
312 storages
313 );
314 slot
315 };
316 for start in 0..5 {
317 for end in 0..5 {
318 assert_eq!(
319 (start..end).into_iter().collect::<Vec<_>>(),
320 storages
321 .iter_range(start..end)
322 .map(check)
323 .collect::<Vec<_>>()
324 );
325 }
326 }
327 assert_eq!(
328 (2..5).into_iter().collect::<Vec<_>>(),
329 storages.iter_range(..5).map(check).collect::<Vec<_>>()
330 );
331 assert_eq!(
332 (1..=4).into_iter().collect::<Vec<_>>(),
333 storages.iter_range(1..).map(check).collect::<Vec<_>>()
334 );
335 assert_eq!(
336 (2..=4).into_iter().collect::<Vec<_>>(),
337 storages.iter_range(..).map(check).collect::<Vec<_>>()
338 );
339 }
340
341 #[test]
342 #[should_panic(expected = "SnapshotStorage.is_empty()")]
343 fn test_sorted_storages_empty() {
344 SortedStorages::new(&[Vec::new()]);
345 }
346
347 #[test]
348 #[should_panic(expected = "slots are not unique")]
349 fn test_sorted_storages_duplicate_slots() {
350 SortedStorages::new_for_tests(&[&Vec::new(), &Vec::new()], &[0, 0]);
351 }
352
353 #[test]
354 fn test_sorted_storages_none() {
355 let result = SortedStorages::empty();
356 assert_eq!(result.range, Range::default());
357 assert_eq!(result.slot_count, 0);
358 assert_eq!(result.storages.len(), 0);
359 assert!(result.get(0).is_none());
360 }
361
362 #[test]
363 fn test_sorted_storages_1() {
364 let vec = vec![];
365 let vec_check = vec.clone();
366 let slot = 4;
367 let vecs = [&vec];
368 let result = SortedStorages::new_for_tests(&vecs, &[slot]);
369 assert_eq!(
370 result.range,
371 Range {
372 start: slot,
373 end: slot + 1
374 }
375 );
376 assert_eq!(result.slot_count, 1);
377 assert_eq!(result.storages.len(), 1);
378 assert_eq!(result.get(slot).unwrap().len(), vec_check.len());
379 }
380
381 #[test]
382 fn test_sorted_storages_2() {
383 let vec = vec![];
384 let vec_check = vec.clone();
385 let slots = [4, 7];
386 let vecs = [&vec, &vec];
387 let result = SortedStorages::new_for_tests(&vecs, &slots);
388 assert_eq!(
389 result.range,
390 Range {
391 start: slots[0],
392 end: slots[1] + 1,
393 }
394 );
395 assert_eq!(result.slot_count, 2);
396 assert_eq!(result.storages.len() as Slot, slots[1] - slots[0] + 1);
397 assert!(result.get(0).is_none());
398 assert!(result.get(3).is_none());
399 assert!(result.get(5).is_none());
400 assert!(result.get(6).is_none());
401 assert!(result.get(8).is_none());
402 assert_eq!(result.get(slots[0]).unwrap().len(), vec_check.len());
403 assert_eq!(result.get(slots[1]).unwrap().len(), vec_check.len());
404 }
405}