diff --git a/Cargo.toml b/Cargo.toml index b177fc8..7dffbfb 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -16,8 +16,5 @@ no_std = ["hashbrown"] [dependencies] hashbrown = { version = "0.14", optional = true } -[dev-dependencies] -libc-print = "0.1.23" - [build-dependencies] rustc_version = "0.4.1" diff --git a/src/lib.rs b/src/lib.rs index 22c97fd..3159206 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,4 +1,4 @@ -#![no_std] +#![cfg_attr(feature = "no_std", no_std)] #![feature(allocator_api)] extern crate alloc; diff --git a/src/vec.rs b/src/vec.rs index baf25dc..6de62d3 100644 --- a/src/vec.rs +++ b/src/vec.rs @@ -39,10 +39,37 @@ impl Vec { self.reserve(1)?; // SAFETY: we just reserved space for one more element. unsafe { - let len = self.inner.len(); - let end = self.inner.as_mut_ptr().add(len); - core::ptr::write(end, value); - self.inner.set_len(len + 1) + self.unsafe_push(value); + } + Ok(()) + } + + #[inline] + unsafe fn unsafe_push(&mut self, value: T) { + let len = self.inner.len(); + let end = self.inner.as_mut_ptr().add(len); + core::ptr::write(end, value); + self.inner.set_len(len + 1) + } + + pub fn extend(&mut self, iter: impl IntoIterator) -> Result<(), TryReserveError> { + let mut iter = iter.into_iter(); + let (lower_bound, _) = iter.size_hint(); + + // Extend N with pre-allocation from the iterator + self.reserve(lower_bound)?; + for _ in 0..lower_bound { + let Some(value) = iter.next() else { + return Ok(()); + }; + unsafe { + self.unsafe_push(value); + } + } + + // Dynamically append the rest + for value in iter { + self.push(value)?; } Ok(()) } @@ -378,4 +405,76 @@ mod tests { vec[3] = 8; assert_eq!(&*vec, &[5, 6, 7, 8]); } + + struct MyIter { + counter: usize, + min_size_hint: usize, + } + + impl MyIter { + fn new(min_size_hint: usize) -> Self { + Self { + counter: 0, + min_size_hint, + } + } + } + + impl Iterator for MyIter { + type Item = usize; + + fn next(&mut self) -> Option { + if self.counter >= 10 { + return None; + } + self.counter += 1; + Some(self.counter - 1) + } + + // This sort-of lies, but it's here for testing purposes. + // It states that the iterator has at least, just, 5 elements. + // This is done so we can get good code coverage and test both + // the optimised pre-reserve code path for `extend` and the + // slower dynamic re-allocation code path. + fn size_hint(&self) -> (usize, Option) { + (self.min_size_hint, None) + } + } + + #[test] + fn test_extend() { + // Test the optimised with mixed pre-reserved and dynamic allocation extend paths. + let wma = WatermarkAllocator::new(32 * size_of::()); + { + let mut vec = Vec::new_in(wma.clone()); + vec.extend(MyIter::new(5)).unwrap(); + assert_eq!(vec.inner.as_slice(), &[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]); + } + assert_eq!(wma.in_use(), 0); + + // Test with a fully pre-reserved path. + { + let mut vec = Vec::new_in(wma.clone()); + vec.extend(MyIter::new(10)).unwrap(); + assert_eq!(vec.inner.as_slice(), &[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]); + } + assert_eq!(wma.in_use(), 0); + + // Test with a fully pre-reserved path, but the `min` size_hint lies + // and exceeds the truth. + { + let mut vec = Vec::new_in(wma.clone()); + vec.extend(MyIter::new(20)).unwrap(); + assert_eq!(vec.inner.as_slice(), &[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]); + } + assert_eq!(wma.in_use(), 0); + + // The min size hint is zero, all dynamically allocated. + { + let mut vec = Vec::new_in(wma.clone()); + vec.extend(MyIter::new(0)).unwrap(); + assert_eq!(vec.inner.as_slice(), &[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]); + } + assert_eq!(wma.in_use(), 0); + } }