use slots returned from get_snapshot_storages to sort (#17638)

* use slots returned from get_snapshot_storages to sort

* add tests

runtime/src/sorted_storages.rs

@@ -7,7 +7,7 @@ use std::ops::Range;
 pub struct SortedStorages<'a> {
     range: Range<Slot>,
     storages: Vec<Option<&'a SnapshotStorage>>,
-    count: usize,
+    slot_count: usize,
 }
 
 impl<'a> SortedStorages<'a> {
@@ -28,35 +28,45 @@ impl<'a> SortedStorages<'a> {
         &self.range
     }
 
-    pub fn len(&self) -> usize {
-        self.count
-    }
-
-    pub fn is_empty(&self) -> bool {
-        self.len() == 0
+    pub fn slot_count(&self) -> usize {
+        self.slot_count
     }
 
+    // assumptions:
+    // 1. each SnapshotStorage.!is_empty()
+    // 2. SnapshotStorage.first().unwrap().get_slot() is unique from all other SnapshotStorage items.
     pub fn new(source: &'a [SnapshotStorage]) -> Self {
-        let mut min = Slot::MAX;
-        let mut max = Slot::MIN;
-        let mut count = 0;
-        let mut time = Measure::start("get slot");
         let slots = source
             .iter()
             .map(|storages| {
-                count += storages.len();
-                if !storages.is_empty() {
-                    storages.first().map(|storage| {
-                        let slot = storage.slot();
-                        min = std::cmp::min(slot, min);
-                        max = std::cmp::max(slot + 1, max);
-                        slot
-                    })
-                } else {
-                    None
-                }
+                let first = storages.first();
+                assert!(first.is_some(), "SnapshotStorage.is_empty()");
+                let storage = first.unwrap();
+                storage.slot() // this must be unique. Will be enforced in new_with_slots
             })
             .collect::<Vec<_>>();
+        Self::new_with_slots(source, &slots)
+    }
+
+    // source[i] is in slot slots[i]
+    // assumptions:
+    // 1. slots vector contains unique slot #s.
+    // 2. slots and source are the same len
+    pub fn new_with_slots(source: &'a [SnapshotStorage], slots: &[Slot]) -> Self {
+        assert_eq!(
+            source.len(),
+            slots.len(),
+            "source and slots are different lengths"
+        );
+        let mut min = Slot::MAX;
+        let mut max = Slot::MIN;
+        let slot_count = source.len();
+        let mut time = Measure::start("get slot");
+        slots.iter().for_each(|slot| {
+            let slot = *slot;
+            min = std::cmp::min(slot, min);
+            max = std::cmp::max(slot + 1, max);
+        });
         time.stop();
         let mut time2 = Measure::start("sort");
         let range;
@@ -75,11 +85,9 @@ impl<'a> SortedStorages<'a> {
                 .iter()
                 .zip(slots)
                 .for_each(|(original_storages, slot)| {
-                    if let Some(slot) = slot {
-                        let index = (slot - min) as usize;
-                        assert!(storages[index].is_none());
-                        storages[index] = Some(original_storages);
-                    }
+                    let index = (slot - min) as usize;
+                    assert!(storages[index].is_none(), "slots are not unique"); // we should not encounter the same slot twice
+                    storages[index] = Some(original_storages);
                 });
         }
         time2.stop();
@@ -87,7 +95,7 @@ impl<'a> SortedStorages<'a> {
         Self {
             range,
             storages,
-            count,
+            slot_count,
         }
     }
 }
@@ -102,7 +110,7 @@ pub mod tests {
                 start: min,
                 end: min + len as Slot,
             };
-            let count = source.len();
+            let slot_count = source.len();
             for (storage, slot) in source {
                 storages[*slot as usize] = Some(*storage);
             }
@@ -110,8 +118,79 @@ pub mod tests {
             Self {
                 range,
                 storages,
-                count,
+                slot_count,
             }
         }
     }
+
+    #[test]
+    #[should_panic(expected = "SnapshotStorage.is_empty()")]
+    fn test_sorted_storages_empty() {
+        SortedStorages::new(&[Vec::new()]);
+    }
+
+    #[test]
+    #[should_panic(expected = "slots are not unique")]
+    fn test_sorted_storages_duplicate_slots() {
+        SortedStorages::new_with_slots(&[Vec::new(), Vec::new()], &[0, 0]);
+    }
+
+    #[test]
+    #[should_panic(expected = "source and slots are different lengths")]
+    fn test_sorted_storages_mismatched_lengths() {
+        SortedStorages::new_with_slots(&[Vec::new()], &[0, 0]);
+    }
+
+    #[test]
+    fn test_sorted_storages_none() {
+        let result = SortedStorages::new_with_slots(&[], &[]);
+        assert_eq!(result.range, Range::default());
+        assert_eq!(result.slot_count, 0);
+        assert_eq!(result.storages.len(), 0);
+        assert!(result.get(0).is_none());
+    }
+
+    #[test]
+    fn test_sorted_storages_1() {
+        let vec = vec![];
+        let vec_check = vec.clone();
+        let slot = 4;
+        let vecs = [vec];
+        let result = SortedStorages::new_with_slots(&vecs, &[slot]);
+        assert_eq!(
+            result.range,
+            Range {
+                start: slot,
+                end: slot + 1
+            }
+        );
+        assert_eq!(result.slot_count, 1);
+        assert_eq!(result.storages.len(), 1);
+        assert_eq!(result.get(slot).unwrap().len(), vec_check.len());
+    }
+
+    #[test]
+    fn test_sorted_storages_2() {
+        let vec = vec![];
+        let vec_check = vec.clone();
+        let slots = [4, 7];
+        let vecs = [vec.clone(), vec];
+        let result = SortedStorages::new_with_slots(&vecs, &slots);
+        assert_eq!(
+            result.range,
+            Range {
+                start: slots[0],
+                end: slots[1] + 1,
+            }
+        );
+        assert_eq!(result.slot_count, 2);
+        assert_eq!(result.storages.len() as Slot, slots[1] - slots[0] + 1);
+        assert!(result.get(0).is_none());
+        assert!(result.get(3).is_none());
+        assert!(result.get(5).is_none());
+        assert!(result.get(6).is_none());
+        assert!(result.get(8).is_none());
+        assert_eq!(result.get(slots[0]).unwrap().len(), vec_check.len());
+        assert_eq!(result.get(slots[1]).unwrap().len(), vec_check.len());
+    }
 }