Skip to content

Commit

Permalink
in-place collect for Vec. Box<[]> and BinaryHeap IntoIter and some ad…
Browse files Browse the repository at this point in the history
…apters
  • Loading branch information
the8472 committed Sep 3, 2020
1 parent 038394a commit bb2d533
Show file tree
Hide file tree
Showing 11 changed files with 280 additions and 43 deletions.
15 changes: 14 additions & 1 deletion library/alloc/src/collections/binary_heap.rs
Expand Up @@ -145,7 +145,7 @@
#![stable(feature = "rust1", since = "1.0.0")]

use core::fmt;
use core::iter::{FromIterator, FusedIterator, TrustedLen};
use core::iter::{FromIterator, FusedIterator, InPlaceIterable, SourceIter, TrustedLen};
use core::mem::{self, size_of, swap, ManuallyDrop};
use core::ops::{Deref, DerefMut};
use core::ptr;
Expand Down Expand Up @@ -1173,6 +1173,19 @@ impl<T> ExactSizeIterator for IntoIter<T> {
#[stable(feature = "fused", since = "1.26.0")]
impl<T> FusedIterator for IntoIter<T> {}

#[unstable(issue = "0", feature = "inplace_iteration")]
impl<T> SourceIter for IntoIter<T> {
type Source = crate::vec::IntoIter<T>;

#[inline]
fn as_inner(&mut self) -> &mut Self::Source {
&mut self.iter
}
}

#[unstable(issue = "0", feature = "inplace_iteration")]
unsafe impl<I> InPlaceIterable for IntoIter<I> {}

#[unstable(feature = "binary_heap_into_iter_sorted", issue = "59278")]
#[derive(Clone, Debug)]
pub struct IntoIterSorted<T> {
Expand Down
1 change: 1 addition & 0 deletions library/alloc/src/lib.rs
Expand Up @@ -99,6 +99,7 @@
#![feature(fmt_internals)]
#![feature(fn_traits)]
#![feature(fundamental)]
#![feature(inplace_iteration)]
#![feature(internal_uninit_const)]
#![feature(lang_items)]
#![feature(layout_for_ptr)]
Expand Down
138 changes: 97 additions & 41 deletions library/alloc/src/vec.rs
Expand Up @@ -58,7 +58,7 @@ use core::cmp::{self, Ordering};
use core::fmt;
use core::hash::{Hash, Hasher};
use core::intrinsics::{arith_offset, assume};
use core::iter::{FromIterator, FusedIterator, TrustedLen};
use core::iter::{FromIterator, FusedIterator, InPlaceIterable, SourceIter, TrustedLen};
use core::marker::PhantomData;
use core::mem::{self, ManuallyDrop, MaybeUninit};
use core::ops::Bound::{Excluded, Included, Unbounded};
Expand Down Expand Up @@ -2012,7 +2012,7 @@ impl<T, I: SliceIndex<[T]>> IndexMut<I> for Vec<T> {
impl<T> FromIterator<T> for Vec<T> {
#[inline]
fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Vec<T> {
<Self as SpecExtend<T, I::IntoIter>>::from_iter(iter.into_iter())
<Self as SpecFrom<T, I::IntoIter>>::from_iter(iter.into_iter())
}
}

Expand Down Expand Up @@ -2094,13 +2094,12 @@ impl<T> Extend<T> for Vec<T> {
}
}

// Specialization trait used for Vec::from_iter and Vec::extend
trait SpecExtend<T, I> {
// Specialization trait used for Vec::from_iter
trait SpecFrom<T, I> {
fn from_iter(iter: I) -> Self;
fn spec_extend(&mut self, iter: I);
}

impl<T, I> SpecExtend<T, I> for Vec<T>
impl<T, I> SpecFrom<T, I> for Vec<T>
where
I: Iterator<Item = T>,
{
Expand All @@ -2125,7 +2124,86 @@ where
<Vec<T> as SpecExtend<T, I>>::spec_extend(&mut vector, iterator);
vector
}
}

fn from_into_iter_source<T, I>(mut iterator: I) -> Vec<T>
where
I: Iterator<Item = T> + InPlaceIterable + SourceIter<Source = IntoIter<T>>,
{
let mut insert_pos = 0;

// FIXME: how to drop values written into source when iteration panics?
// tail already gets cleaned by IntoIter::drop
while let Some(item) = iterator.next() {
let source_iter = iterator.as_inner();
let src_buf = source_iter.buf.as_ptr();
let src_idx = source_iter.ptr;
unsafe {
let dst = src_buf.offset(insert_pos as isize);
debug_assert!(
dst as *const _ < src_idx,
"InPlaceIterable implementation produced more\
items than it consumed from the source"
);
ptr::write(dst, item)
}
insert_pos += 1;
}

let src = iterator.as_inner();
let vec = unsafe { Vec::from_raw_parts(src.buf.as_ptr(), insert_pos, src.cap) };
mem::forget(iterator);
vec
}

impl<T> SpecFrom<T, IntoIter<T>> for Vec<T> {
fn from_iter(iterator: IntoIter<T>) -> Self {
// A common case is passing a vector into a function which immediately
// re-collects into a vector. We can short circuit this if the IntoIter
// has not been advanced at all.
if iterator.buf.as_ptr() as *const _ == iterator.ptr {
unsafe {
let it = ManuallyDrop::new(iterator);
return Vec::from_raw_parts(it.buf.as_ptr(), it.len(), it.cap);
}
}

from_into_iter_source(iterator)
}
}

// Further specialization potential once lattice specialization exists
// and https://github.com/rust-lang/rust/issues/62645 has been solved:
// This can be broadened to only require size and alignment equality between
// input and output Item types.
impl<T, I> SpecFrom<T, I> for Vec<T>
where
I: Iterator<Item = T> + InPlaceIterable + SourceIter<Source = IntoIter<T>>,
{
default fn from_iter(iterator: I) -> Self {
from_into_iter_source(iterator)
}
}

impl<'a, T: 'a, I> SpecFrom<&'a T, I> for Vec<T>
where
I: Iterator<Item = &'a T>,
T: Clone,
{
default fn from_iter(iterator: I) -> Self {
SpecFrom::from_iter(iterator.cloned())
}
}

// Specialization trait used for Vec::extend
trait SpecExtend<T, I> {
fn spec_extend(&mut self, iter: I);
}

impl<T, I> SpecExtend<T, I> for Vec<T>
where
I: Iterator<Item = T>,
{
default fn spec_extend(&mut self, iter: I) {
self.extend_desugared(iter)
}
Expand All @@ -2135,12 +2213,6 @@ impl<T, I> SpecExtend<T, I> for Vec<T>
where
I: TrustedLen<Item = T>,
{
default fn from_iter(iterator: I) -> Self {
let mut vector = Vec::new();
vector.spec_extend(iterator);
vector
}

default fn spec_extend(&mut self, iterator: I) {
// This is the case for a TrustedLen iterator.
let (low, high) = iterator.size_hint();
Expand Down Expand Up @@ -2170,40 +2242,11 @@ where
}
}

impl<T> SpecExtend<T, IntoIter<T>> for Vec<T> {
fn from_iter(iterator: IntoIter<T>) -> Self {
// A common case is passing a vector into a function which immediately
// re-collects into a vector. We can short circuit this if the IntoIter
// has not been advanced at all.
if iterator.buf.as_ptr() as *const _ == iterator.ptr {
unsafe {
let it = ManuallyDrop::new(iterator);
Vec::from_raw_parts(it.buf.as_ptr(), it.len(), it.cap)
}
} else {
let mut vector = Vec::new();
vector.spec_extend(iterator);
vector
}
}

fn spec_extend(&mut self, mut iterator: IntoIter<T>) {
unsafe {
self.append_elements(iterator.as_slice() as _);
}
iterator.ptr = iterator.end;
}
}

impl<'a, T: 'a, I> SpecExtend<&'a T, I> for Vec<T>
where
I: Iterator<Item = &'a T>,
T: Clone,
{
default fn from_iter(iterator: I) -> Self {
SpecExtend::from_iter(iterator.cloned())
}

default fn spec_extend(&mut self, iterator: I) {
self.spec_extend(iterator.cloned())
}
Expand Down Expand Up @@ -2779,6 +2822,19 @@ unsafe impl<#[may_dangle] T> Drop for IntoIter<T> {
}
}

#[unstable(issue = "0", feature = "inplace_iteration")]
unsafe impl<T> InPlaceIterable for IntoIter<T> {}

#[unstable(issue = "0", feature = "inplace_iteration")]
impl<T> SourceIter for IntoIter<T> {
type Source = IntoIter<T>;

#[inline]
fn as_inner(&mut self) -> &mut Self::Source {
self
}
}

/// A draining iterator for `Vec<T>`.
///
/// This `struct` is created by [`Vec::drain`].
Expand Down
12 changes: 12 additions & 0 deletions library/alloc/tests/binary_heap.rs
Expand Up @@ -230,6 +230,18 @@ fn test_to_vec() {
check_to_vec(vec![5, 4, 3, 2, 1, 5, 4, 3, 2, 1, 5, 4, 3, 2, 1]);
}

#[test]
fn test_in_place_iterator_specialization() {
let src: Vec<usize> = vec![1, 2, 3];
let src_ptr = src.as_ptr();
let heap: BinaryHeap<_> = src.into_iter().map(std::convert::identity).collect();
let heap_ptr = heap.iter().next().unwrap() as *const usize;
assert_eq!(src_ptr, heap_ptr);
let sink: Vec<_> = heap.into_iter().map(std::convert::identity).collect();
let sink_ptr = sink.as_ptr();
assert_eq!(heap_ptr, sink_ptr);
}

#[test]
fn test_empty_pop() {
let mut heap = BinaryHeap::<i32>::new();
Expand Down
1 change: 1 addition & 0 deletions library/alloc/tests/lib.rs
Expand Up @@ -14,6 +14,7 @@
#![feature(slice_ptr_get)]
#![feature(split_inclusive)]
#![feature(binary_heap_retain)]
#![feature(inplace_iteration)]

use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};
Expand Down
9 changes: 9 additions & 0 deletions library/alloc/tests/slice.rs
Expand Up @@ -1459,6 +1459,15 @@ fn test_to_vec() {
assert_eq!(ys, [1, 2, 3]);
}

#[test]
fn test_in_place_iterator_specialization() {
let src: Box<[usize]> = box [1, 2, 3];
let src_ptr = src.as_ptr();
let sink: Box<_> = src.into_vec().into_iter().map(std::convert::identity).collect();
let sink_ptr = sink.as_ptr();
assert_eq!(src_ptr, sink_ptr);
}

#[test]
fn test_box_slice_clone() {
let data = vec![vec![0, 1], vec![0], vec![1]];
Expand Down
22 changes: 22 additions & 0 deletions library/alloc/tests/vec.rs
@@ -1,6 +1,7 @@
use std::borrow::Cow;
use std::collections::TryReserveError::*;
use std::fmt::Debug;
use std::iter::InPlaceIterable;
use std::mem::size_of;
use std::panic::{catch_unwind, AssertUnwindSafe};
use std::vec::{Drain, IntoIter};
Expand Down Expand Up @@ -775,6 +776,27 @@ fn test_into_iter_leak() {
assert_eq!(unsafe { DROPS }, 3);
}

#[test]
fn test_from_iter_specialization() {
let src: Vec<usize> = vec![0usize; 1];
let srcptr = src.as_ptr();
let sink = src.into_iter().collect::<Vec<_>>();
let sinkptr = sink.as_ptr();
assert_eq!(srcptr, sinkptr);
}

#[test]
fn test_from_iter_specialization_with_iterator_adapters() {
fn assert_in_place_trait<T: InPlaceIterable>(_: &T) {};
let src: Vec<usize> = vec![0usize; 65535];
let srcptr = src.as_ptr();
let iter = src.into_iter().enumerate().map(|i| i.0 + i.1).peekable().skip(1);
assert_in_place_trait(&iter);
let sink = iter.collect::<Vec<_>>();
let sinkptr = sink.as_ptr();
assert_eq!(srcptr, sinkptr);
}

#[test]
fn test_cow_from() {
let borrowed: &[_] = &["borrowed", "(slice)"];
Expand Down

0 comments on commit bb2d533

Please sign in to comment.