Merge pull request #19 from hbina/fix/crashing-when-draining-empty-mi…

…nivec-small Fix/crashing when draining empty minivec small
cmazakas · Feb 18, 2021 · 09fd54a · 09fd54a
2 parents 29caf18 + 4acabc9
commit 09fd54a
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 8 deletions.
diff --git a/src/impl/drain.rs b/src/impl/drain.rs
@@ -18,13 +18,24 @@ pub fn make_drain_iterator<'a, T>(
     start_idx: usize,
     end_idx: usize,
 ) -> Drain<'a, T> {
-    Drain {
-        vec_: core::ptr::NonNull::from(vec),
-        drain_pos_: unsafe { core::ptr::NonNull::new_unchecked(data.add(start_idx)) },
-        drain_end_: unsafe { core::ptr::NonNull::new_unchecked(data.add(end_idx)) },
-        remaining_pos_: unsafe { core::ptr::NonNull::new_unchecked(data.add(end_idx)) },
-        remaining_: remaining,
-        marker_: core::marker::PhantomData,
+    if data.is_null() {
+        Drain {
+            vec_: core::ptr::NonNull::from(vec),
+            drain_pos_: core::ptr::NonNull::dangling(),
+            drain_end_: core::ptr::NonNull::dangling(),
+            remaining_pos_: core::ptr::NonNull::dangling(),
+            remaining_: remaining,
+            marker_: core::marker::PhantomData,
+        }
+    } else {
+        Drain {
+            vec_: core::ptr::NonNull::from(vec),
+            drain_pos_: unsafe { core::ptr::NonNull::new_unchecked(data.add(start_idx)) },
+            drain_end_: unsafe { core::ptr::NonNull::new_unchecked(data.add(end_idx)) },
+            remaining_pos_: unsafe { core::ptr::NonNull::new_unchecked(data.add(end_idx)) },
+            remaining_: remaining,
+            marker_: core::marker::PhantomData,
+        }
     }
 }
 

diff --git a/src/lib.rs b/src/lib.rs
@@ -456,7 +456,9 @@ impl<T> MiniVec<T> {
 
         let data = self.as_mut_ptr();
 
-        unsafe { self.set_len(start_idx) };
+        if !data.is_null() {
+            unsafe { self.set_len(start_idx) };
+        }
 
         make_drain_iterator(self, data, len - end_idx, start_idx, end_idx)
     }

diff --git a/tests/mine.rs b/tests/mine.rs
@@ -97,6 +97,15 @@ fn minivec_dedup_needs_drop() {
     assert_eq!(v.len(), 3);
 }
 
+#[test]
+fn minivec_drain_into_vec() {
+    let mut vec: MiniVec<usize> = mini_vec![];
+
+    let drain = vec.drain(..).collect::<MiniVec<_>>();
+
+    assert_eq!(drain.len(), 0);
+}
+
 #[test]
 fn minvec_drain() {
     let mut vec = mini_vec![1, 2, 3];