diff --git a/library/alloc/src/vec/mod.rs b/library/alloc/src/vec/mod.rs index a08e99a277d70..76b3c6bc3ada2 100644 --- a/library/alloc/src/vec/mod.rs +++ b/library/alloc/src/vec/mod.rs @@ -4006,19 +4006,23 @@ impl Vec { // for item in iterator { // self.push(item); // } - while let Some(element) = iterator.next() { - let len = self.len(); - if len == self.capacity() { - let (lower, _) = iterator.size_hint(); - self.reserve(lower.saturating_add(1)); - } - unsafe { - ptr::write(self.as_mut_ptr().add(len), element); - // Since next() executes user code which can panic we have to bump the length - // after each step. - // NB can't overflow since we would have had to alloc the address space - self.set_len(len + 1); + let (lower, _) = iterator.size_hint(); + self.reserve(lower); + while let Err(element) = iterator.try_fold((), |(), element| { + if self.len < self.capacity() { + unsafe { + let end = self.as_mut_ptr().add(self.len); + ptr::write(end, element); + self.len += 1; + } + Ok(()) + } else { + Err(element) } + }) { + let (lower, _) = iterator.size_hint(); + self.reserve(lower.saturating_add(1)); + self.push(element); } } diff --git a/library/alloctests/benches/vec.rs b/library/alloctests/benches/vec.rs index 1dab71fa1f4f4..dcfc65450f6c0 100644 --- a/library/alloctests/benches/vec.rs +++ b/library/alloctests/benches/vec.rs @@ -193,6 +193,117 @@ fn bench_extend_1000_1000(b: &mut Bencher) { do_bench_extend(b, 1000, 1000) } +fn do_bench_extend_chain_filter(b: &mut Bencher, dst_len: usize, src_len: usize) { + let dst: Vec<_> = FromIterator::from_iter(0..dst_len); + let src1: Vec<_> = FromIterator::from_iter(dst_len..dst_len + src_len); + let src2: Vec<_> = FromIterator::from_iter(dst_len + src_len..dst_len + src_len * 2); + + b.bytes = src_len as u64 * 2; + + b.iter(|| { + let mut dst = dst.clone(); + dst.extend(src1.iter().chain(src2.iter()).filter(|&&x| x < usize::MAX)); + dst + }); +} + +#[bench] +fn bench_extend_chain_filter_0000_0000(b: &mut Bencher) { + do_bench_extend_chain_filter(b, 0, 0) +} + +#[bench] +fn bench_extend_chain_filter_0000_0010(b: &mut Bencher) { + do_bench_extend_chain_filter(b, 0, 10) +} + +#[bench] +fn bench_extend_chain_filter_0000_0100(b: &mut Bencher) { + do_bench_extend_chain_filter(b, 0, 100) +} + +#[bench] +fn bench_extend_chain_filter_0000_1000(b: &mut Bencher) { + do_bench_extend_chain_filter(b, 0, 1000) +} + +#[bench] +fn bench_extend_chain_filter_0010_0010(b: &mut Bencher) { + do_bench_extend_chain_filter(b, 10, 10) +} + +#[bench] +fn bench_extend_chain_filter_0100_0100(b: &mut Bencher) { + do_bench_extend_chain_filter(b, 100, 100) +} + +#[bench] +fn bench_extend_chain_filter_1000_1000(b: &mut Bencher) { + do_bench_extend_chain_filter(b, 1000, 1000) +} + +fn do_bench_collect_chain_filter(b: &mut Bencher, src_len: usize) { + let src1: Vec = FromIterator::from_iter(0..src_len); + let src2: Vec = FromIterator::from_iter(src_len..src_len * 2); + + b.bytes = src_len as u64 * 2; + + b.iter(|| src1.iter().chain(src2.iter()).filter(|&&x| x < usize::MAX).collect::>()); +} + +#[bench] +fn bench_collect_chain_filter_0010(b: &mut Bencher) { + do_bench_collect_chain_filter(b, 10) +} + +#[bench] +fn bench_collect_chain_filter_0100(b: &mut Bencher) { + do_bench_collect_chain_filter(b, 100) +} + +#[bench] +fn bench_collect_chain_filter_1000(b: &mut Bencher) { + do_bench_collect_chain_filter(b, 1000) +} + +fn do_bench_extend_flatten(b: &mut Bencher, inner_len: usize, outer_len: usize) { + let src: Vec> = + (0..outer_len).map(|i| (i as u32..i as u32 + inner_len as u32).collect()).collect(); + + b.bytes = (inner_len * outer_len) as u64; + + b.iter(|| { + let mut v: Vec<&u32> = Vec::new(); + v.extend(src.iter().flatten()); + v + }); +} + +#[bench] +fn bench_extend_flatten_0001_0000(b: &mut Bencher) { + do_bench_extend_flatten(b, 1, 0) +} + +#[bench] +fn bench_extend_flatten_0001_0010(b: &mut Bencher) { + do_bench_extend_flatten(b, 1, 10) +} + +#[bench] +fn bench_extend_flatten_0010_0100(b: &mut Bencher) { + do_bench_extend_flatten(b, 10, 100) +} + +#[bench] +fn bench_extend_flatten_0100_0100(b: &mut Bencher) { + do_bench_extend_flatten(b, 100, 100) +} + +#[bench] +fn bench_extend_flatten_1000_0010(b: &mut Bencher) { + do_bench_extend_flatten(b, 1000, 10) +} + fn do_bench_extend_from_slice(b: &mut Bencher, dst_len: usize, src_len: usize) { let dst: Vec<_> = FromIterator::from_iter(0..dst_len); let src: Vec<_> = FromIterator::from_iter(dst_len..dst_len + src_len);