in-place collect for Vec. Box<[]> and BinaryHeap IntoIter and some adapters

11 files changed, 280 insertions, 43 deletions

diff --git a/library/alloc/src/collections/binary_heap.rs b/library/alloc/src/collections/binary_heap.rs
index 477a598ff5b..f7012a03425 100644
--- a/library/alloc/src/collections/binary_heap.rs
+++ b/library/alloc/src/collections/binary_heap.rs
@@ -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;
@@ -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> {
diff --git a/library/alloc/src/lib.rs b/library/alloc/src/lib.rs
index 755a21934f0..72aa7fea4cf 100644
--- a/library/alloc/src/lib.rs
+++ b/library/alloc/src/lib.rs
@@ -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)]
diff --git a/library/alloc/src/vec.rs b/library/alloc/src/vec.rs
index 719ee8a3b0f..ad639ca320a 100644
--- a/library/alloc/src/vec.rs
+++ b/library/alloc/src/vec.rs
@@ -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};
@@ -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())
     }
 }
 
@@ -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>,
 {
@@ -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)
     }
@@ -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();
@@ -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())
     }
@@ -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`].
diff --git a/library/alloc/tests/binary_heap.rs b/library/alloc/tests/binary_heap.rs
index 62084ccf53c..ce794a9a4af 100644
--- a/library/alloc/tests/binary_heap.rs
+++ b/library/alloc/tests/binary_heap.rs
@@ -231,6 +231,18 @@ fn test_to_vec() {
 }
 
 #[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();
     assert!(heap.pop().is_none());
diff --git a/library/alloc/tests/lib.rs b/library/alloc/tests/lib.rs
index f2ba1ab6481..e0e146dc427 100644
--- a/library/alloc/tests/lib.rs
+++ b/library/alloc/tests/lib.rs
@@ -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};
diff --git a/library/alloc/tests/slice.rs b/library/alloc/tests/slice.rs
index 3c7d57f8b32..1f561bebd90 100644
--- a/library/alloc/tests/slice.rs
+++ b/library/alloc/tests/slice.rs
@@ -1460,6 +1460,15 @@ fn test_to_vec() {
 }
 
 #[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]];
     let data2 = data.clone().into_boxed_slice().clone().to_vec();
diff --git a/library/alloc/tests/vec.rs b/library/alloc/tests/vec.rs
index ffff543b07f..271a705cf06 100644
--- a/library/alloc/tests/vec.rs
+++ b/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};
@@ -776,6 +777,27 @@ fn test_into_iter_leak() {
 }
 
 #[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)"];
     let owned = vec!["owned", "(vec)"];
diff --git a/library/core/src/iter/adapters/mod.rs b/library/core/src/iter/adapters/mod.rs
index f32c3963abe..c7488ef2720 100644
--- a/library/core/src/iter/adapters/mod.rs
+++ b/library/core/src/iter/adapters/mod.rs
@@ -4,7 +4,9 @@ use crate::intrinsics;
 use crate::ops::{Add, AddAssign, ControlFlow, Try};
 
 use super::from_fn;
-use super::{DoubleEndedIterator, ExactSizeIterator, FusedIterator, Iterator, TrustedLen};
+use super::{
+    DoubleEndedIterator, ExactSizeIterator, FusedIterator, InPlaceIterable, Iterator, TrustedLen,
+};
 
 mod chain;
 mod flatten;
@@ -19,6 +21,40 @@ use self::zip::try_get_unchecked;
 pub(crate) use self::zip::TrustedRandomAccess;
 pub use self::zip::Zip;
 
+/// This trait provides access to to the backing source of an interator-adapter pipeline
+/// under the conditions that
+/// * the iterator source `S` itself implements `SourceIter<Source = S>`
+/// * there is a delegating implementation of this trait for each adapter in the pipeline
+///
+/// This is useful for specializing [`FromIterator`] implementations or to retrieve
+/// the remaining values from a source of a partially consumed iterator.
+///
+/// # Examples
+///
+/// Retrieving a partially consumed source and wrapping it into a different pipeline:
+///
+/// ```
+/// # #![feature(inplace_iteration)]
+/// # use std::iter::SourceIter;
+///
+/// let mut iter = vec![9, 9, 9].into_iter().map(|i| i * i);
+/// let first = iter.next().unwrap();
+/// let mut remainder = std::mem::replace(iter.as_inner(), Vec::new().into_iter());
+/// let second = remainder.map(|i| i + 1).next().unwrap();
+/// assert_eq!(first, 81);
+/// assert_eq!(second, 10);
+/// ```
+///
+/// [`FromIterator`]: trait.FromIterator.html
+#[unstable(issue = "0", feature = "inplace_iteration")]
+pub trait SourceIter {
+    /// The source iterator of the adapter.
+    type Source: Iterator;
+
+    /// Recursively extract the source of an iterator pipeline.
+    fn as_inner(&mut self) -> &mut Self::Source;
+}
+
 /// A double-ended iterator with the direction inverted.
 ///
 /// This `struct` is created by the [`rev`] method on [`Iterator`]. See its
@@ -939,6 +975,23 @@ where
     }
 }
 
+#[unstable(issue = "0", feature = "inplace_iteration")]
+impl<S: Iterator, B, I: Iterator, F> SourceIter for Map<I, F>
+where
+    F: FnMut(I::Item) -> B,
+    I: SourceIter<Source = S>,
+{
+    type Source = S;
+
+    #[inline]
+    fn as_inner(&mut self) -> &mut S {
+        SourceIter::as_inner(&mut self.iter)
+    }
+}
+
+#[unstable(issue = "0", feature = "inplace_iteration")]
+unsafe impl<B, I: InPlaceIterable, F> InPlaceIterable for Map<I, F> where F: FnMut(I::Item) -> B {}
+
 /// An iterator that filters the elements of `iter` with `predicate`.
 ///
 /// This `struct` is created by the [`filter`] method on [`Iterator`]. See its
@@ -1414,6 +1467,22 @@ impl<I> FusedIterator for Enumerate<I> where I: FusedIterator {}
 #[unstable(feature = "trusted_len", issue = "37572")]
 unsafe impl<I> TrustedLen for Enumerate<I> where I: TrustedLen {}
 
+#[unstable(issue = "0", feature = "inplace_iteration")]
+impl<S: Iterator, I: Iterator> SourceIter for Enumerate<I>
+where
+    I: SourceIter<Source = S>,
+{
+    type Source = S;
+
+    #[inline]
+    fn as_inner(&mut self) -> &mut S {
+        SourceIter::as_inner(&mut self.iter)
+    }
+}
+
+#[unstable(issue = "0", feature = "inplace_iteration")]
+unsafe impl<I: InPlaceIterable> InPlaceIterable for Enumerate<I> {}
+
 /// An iterator with a `peek()` that returns an optional reference to the next
 /// element.
 ///
@@ -1692,6 +1761,25 @@ impl<I: Iterator> Peekable<I> {
     }
 }
 
+#[unstable(feature = "trusted_len", issue = "37572")]
+unsafe impl<I> TrustedLen for Peekable<I> where I: TrustedLen {}
+
+#[unstable(issue = "0", feature = "inplace_iteration")]
+impl<S: Iterator, I: Iterator> SourceIter for Peekable<I>
+where
+    I: SourceIter<Source = S>,
+{
+    type Source = S;
+
+    #[inline]
+    fn as_inner(&mut self) -> &mut S {
+        SourceIter::as_inner(&mut self.iter)
+    }
+}
+
+#[unstable(issue = "0", feature = "inplace_iteration")]
+unsafe impl<I: InPlaceIterable> InPlaceIterable for Peekable<I> {}
+
 /// An iterator that rejects elements while `predicate` returns `true`.
 ///
 /// This `struct` is created by the [`skip_while`] method on [`Iterator`]. See its
@@ -2167,6 +2255,22 @@ where
 #[stable(feature = "fused", since = "1.26.0")]
 impl<I> FusedIterator for Skip<I> where I: FusedIterator {}
 
+#[unstable(issue = "0", feature = "inplace_iteration")]
+impl<S: Iterator, I: Iterator> SourceIter for Skip<I>
+where
+    I: SourceIter<Source = S>,
+{
+    type Source = S;
+
+    #[inline]
+    fn as_inner(&mut self) -> &mut S {
+        SourceIter::as_inner(&mut self.iter)
+    }
+}
+
+#[unstable(issue = "0", feature = "inplace_iteration")]
+unsafe impl<I: InPlaceIterable> InPlaceIterable for Skip<I> {}
+
 /// An iterator that only iterates over the first `n` iterations of `iter`.
 ///
 /// This `struct` is created by the [`take`] method on [`Iterator`]. See its
diff --git a/library/core/src/iter/mod.rs b/library/core/src/iter/mod.rs
index bab8dda2915..f35994560c5 100644
--- a/library/core/src/iter/mod.rs
+++ b/library/core/src/iter/mod.rs
@@ -342,6 +342,9 @@ pub use self::traits::{DoubleEndedIterator, Extend, FromIterator, IntoIterator};
 #[stable(feature = "rust1", since = "1.0.0")]
 pub use self::traits::{ExactSizeIterator, Product, Sum};
 
+#[unstable(issue = "0", feature = "inplace_iteration")]
+pub use self::traits::InPlaceIterable;
+
 #[stable(feature = "iter_cloned", since = "1.1.0")]
 pub use self::adapters::Cloned;
 #[stable(feature = "iter_copied", since = "1.36.0")]
@@ -350,6 +353,8 @@ pub use self::adapters::Copied;
 pub use self::adapters::Flatten;
 #[unstable(feature = "iter_map_while", reason = "recently added", issue = "68537")]
 pub use self::adapters::MapWhile;
+#[unstable(issue = "0", feature = "inplace_iteration")]
+pub use self::adapters::SourceIter;
 #[stable(feature = "iterator_step_by", since = "1.28.0")]
 pub use self::adapters::StepBy;
 #[stable(feature = "rust1", since = "1.0.0")]
diff --git a/library/core/src/iter/traits/marker.rs b/library/core/src/iter/traits/marker.rs
index 3c893c03992..549f7972689 100644
--- a/library/core/src/iter/traits/marker.rs
+++ b/library/core/src/iter/traits/marker.rs
@@ -42,3 +42,15 @@ pub unsafe trait TrustedLen: Iterator {}
 
 #[unstable(feature = "trusted_len", issue = "37572")]
 unsafe impl<I: TrustedLen + ?Sized> TrustedLen for &mut I {}
+
+/// An iterator that when yielding an item will have taken at least one element
+/// from its underlying [`SourceIter`].
+///
+/// Calling next() guarantees that at least one value of the iterator's underlying source
+/// has been moved out and the result of the iterator chain could be inserted in its place,
+/// assuming structural constraints of the source allow such an insertion.
+/// In other words this trait indicates that an iterator pipeline can be collected in place.
+///
+/// [`SourceIter`]: ../../std/iter/trait.SourceIter.html
+#[unstable(issue = "0", feature = "inplace_iteration")]
+pub unsafe trait InPlaceIterable: Iterator {}
diff --git a/library/core/src/iter/traits/mod.rs b/library/core/src/iter/traits/mod.rs
index efd1580a548..9ed2de7313d 100644
--- a/library/core/src/iter/traits/mod.rs
+++ b/library/core/src/iter/traits/mod.rs
@@ -11,5 +11,7 @@ pub use self::double_ended::DoubleEndedIterator;
 pub use self::exact_size::ExactSizeIterator;
 #[stable(feature = "rust1", since = "1.0.0")]
 pub use self::iterator::Iterator;
+#[unstable(issue = "0", feature = "inplace_iteration")]
+pub use self::marker::InPlaceIterable;
 #[stable(feature = "rust1", since = "1.0.0")]
 pub use self::marker::{FusedIterator, TrustedLen};