Skip to content

Panic in .extend() and from_iter on capacity exceeded #175

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Mar 23, 2021
Merged

Panic in .extend() and from_iter on capacity exceeded #175

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
110 changes: 70 additions & 40 deletions src/arrayvec.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -704,7 +704,7 @@ impl<T, const CAP: usize> std::convert::TryFrom<&[T]> for ArrayVec<T, CAP>
Err(CapacityError::new(()))
} else {
let mut array = Self::new();
array.extend(slice.iter().cloned());
array.extend_from_slice(slice);
Ok(array)
}
}
Expand Down Expand Up @@ -931,39 +931,75 @@ impl<T, Data, F> Drop for ScopeExitGuard<T, Data, F>

/// Extend the `ArrayVec` with an iterator.
///
/// Does not extract more items than there is space for. No error
/// occurs if there are more iterator elements.
/// ***Panics*** if extending the vector exceeds its capacity.
impl<T, const CAP: usize> Extend<T> for ArrayVec<T, CAP> {
/// Extend the `ArrayVec` with an iterator.
///
/// ***Panics*** if extending the vector exceeds its capacity.
fn extend<I: IntoIterator<Item=T>>(&mut self, iter: I) {
let take = self.capacity() - self.len();
unsafe {
let len = self.len();
let mut ptr = raw_ptr_add(self.as_mut_ptr(), len);
let end_ptr = raw_ptr_add(ptr, take);
// Keep the length in a separate variable, write it back on scope
// exit. To help the compiler with alias analysis and stuff.
// We update the length to handle panic in the iteration of the
// user's iterator, without dropping any elements on the floor.
let mut guard = ScopeExitGuard {
value: &mut self.len,
data: len,
f: move |&len, self_len| {
**self_len = len;
}
};
let mut iter = iter.into_iter();
loop {
if ptr == end_ptr { break; }
if let Some(elt) = iter.next() {
raw_ptr_write(ptr, elt);
ptr = raw_ptr_add(ptr, 1);
guard.data += 1;
} else {
break;
}
self.extend_from_iter::<_, true>(iter)
}
}
}

#[inline(never)]
#[cold]
fn extend_panic() {
panic!("ArrayVec: capacity exceeded in extend/from_iter");
}

impl<T, const CAP: usize> ArrayVec<T, CAP> {
/// Extend the arrayvec from the iterable.
///
/// ## Safety
///
/// Unsafe because if CHECK is false, the length of the input is not checked.
/// The caller must ensure the length of the input fits in the capacity.
pub(crate) unsafe fn extend_from_iter<I, const CHECK: bool>(&mut self, iterable: I)
where I: IntoIterator<Item = T>
{
let take = self.capacity() - self.len();
let len = self.len();
let mut ptr = raw_ptr_add(self.as_mut_ptr(), len);
let end_ptr = raw_ptr_add(ptr, take);
// Keep the length in a separate variable, write it back on scope
// exit. To help the compiler with alias analysis and stuff.
// We update the length to handle panic in the iteration of the
// user's iterator, without dropping any elements on the floor.
let mut guard = ScopeExitGuard {
value: &mut self.len,
data: len,
f: move |&len, self_len| {
**self_len = len;
}
};
let mut iter = iterable.into_iter();
loop {
if let Some(elt) = iter.next() {
if ptr == end_ptr && CHECK { extend_panic(); }
debug_assert_ne!(ptr, end_ptr);
ptr.write(elt);
ptr = raw_ptr_add(ptr, 1);
guard.data += 1;
} else {
return; // success
}
}
}

/// Extend the ArrayVec with clones of elements from the slice;
/// the length of the slice must be <= the remaining capacity in the arrayvec.
pub(crate) fn extend_from_slice(&mut self, slice: &[T])
where T: Clone
{
let take = self.capacity() - self.len();
debug_assert!(slice.len() <= take);
unsafe {
let slice = if take < slice.len() { &slice[..take] } else { slice };
self.extend_from_iter::<_, false>(slice.iter().cloned());
}
}
}

/// Rawptr add but uses arithmetic distance for ZST
Expand All @@ -976,19 +1012,13 @@ unsafe fn raw_ptr_add<T>(ptr: *mut T, offset: usize) -> *mut T {
}
}

unsafe fn raw_ptr_write<T>(ptr: *mut T, value: T) {
if mem::size_of::<T>() == 0 {
/* nothing */
} else {
ptr::write(ptr, value)
}
}

/// Create an `ArrayVec` from an iterator.
///
/// Does not extract more items than there is space for. No error
/// occurs if there are more iterator elements.
/// ***Panics*** if the number of elements in the iterator exceeds the arrayvec's capacity.
impl<T, const CAP: usize> iter::FromIterator<T> for ArrayVec<T, CAP> {
/// Create an `ArrayVec` from an iterator.
///
/// ***Panics*** if the number of elements in the iterator exceeds the arrayvec's capacity.
fn from_iter<I: IntoIterator<Item=T>>(iter: I) -> Self {
let mut array = ArrayVec::new();
array.extend(iter);
Expand All @@ -1014,8 +1044,8 @@ impl<T, const CAP: usize> Clone for ArrayVec<T, CAP>
self.pop();
}
} else {
let rhs_elems = rhs[self.len()..].iter().cloned();
self.extend(rhs_elems);
let rhs_elems = &rhs[self.len()..];
self.extend_from_slice(rhs_elems);
}
}
}
Expand Down
35 changes: 27 additions & 8 deletions tests/tests.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -251,11 +251,11 @@ fn test_drop_panics() {
fn test_extend() {
let mut range = 0..10;

let mut array: ArrayVec<_, 5> = range.by_ref().collect();
let mut array: ArrayVec<_, 5> = range.by_ref().take(5).collect();
assert_eq!(&array[..], &[0, 1, 2, 3, 4]);
assert_eq!(range.next(), Some(5));

array.extend(range.by_ref());
array.extend(range.by_ref().take(0));
assert_eq!(range.next(), Some(6));

let mut array: ArrayVec<_, 10> = (0..3).collect();
Expand All @@ -264,6 +264,25 @@ fn test_extend() {
assert_eq!(&array[..], &[0, 1, 2, 3, 4]);
}

#[should_panic]
#[test]
fn test_extend_capacity_panic_1() {
let mut range = 0..10;

let _: ArrayVec<_, 5> = range.by_ref().collect();
}

#[should_panic]
#[test]
fn test_extend_capacity_panic_2() {
let mut range = 0..10;

let mut array: ArrayVec<_, 5> = range.by_ref().take(5).collect();
assert_eq!(&array[..], &[0, 1, 2, 3, 4]);
assert_eq!(range.next(), Some(5));
array.extend(range.by_ref().take(1));
}

#[test]
fn test_is_send_sync() {
let data = ArrayVec::<Vec<i32>, 5>::new();
Expand Down Expand Up @@ -304,7 +323,7 @@ fn test_drain() {
v.drain(0..7);
assert_eq!(&v[..], &[]);

v.extend(0..);
v.extend(0..8);
v.drain(1..4);
assert_eq!(&v[..], &[0, 4, 5, 6, 7]);
let u: ArrayVec<_, 3> = v.drain(1..4).rev().collect();
Expand All @@ -320,7 +339,7 @@ fn test_drain_range_inclusive() {
v.drain(0..=7);
assert_eq!(&v[..], &[]);

v.extend(0..);
v.extend(0..8);
v.drain(1..=4);
assert_eq!(&v[..], &[0, 5, 6, 7]);
let u: ArrayVec<_, 3> = v.drain(1..=2).rev().collect();
Expand Down Expand Up @@ -436,9 +455,9 @@ fn test_into_inner_2() {
}

#[test]
fn test_into_inner_3_() {
fn test_into_inner_3() {
let mut v = ArrayVec::<i32, 4>::new();
v.extend(1..);
v.extend(1..=4);
assert_eq!(v.into_inner().unwrap(), [1, 2, 3, 4]);
}

Expand Down Expand Up @@ -672,11 +691,11 @@ fn test_extend_zst() {
#[derive(Copy, Clone, PartialEq, Debug)]
struct Z; // Zero sized type

let mut array: ArrayVec<_, 5> = range.by_ref().map(|_| Z).collect();
let mut array: ArrayVec<_, 5> = range.by_ref().take(5).map(|_| Z).collect();
assert_eq!(&array[..], &[Z; 5]);
assert_eq!(range.next(), Some(5));

array.extend(range.by_ref().map(|_| Z));
array.extend(range.by_ref().take(0).map(|_| Z));
assert_eq!(range.next(), Some(6));

let mut array: ArrayVec<_, 10> = (0..3).map(|_| Z).collect();
Expand Down