support in-place collecting additional FlatMap shapes

Author: the8472

library/alloc/tests/vec.rs

@@ -1212,6 +1212,17 @@ fn test_in_place_specialization_step_up_down() {
     let sink_bytes = sink.capacity() * 4;
     assert_ne!(src_bytes, sink_bytes);
     assert_eq!(sink.len(), 2);
+
+    let src = vec![[0u8; 4]; 256];
+    let srcptr = src.as_ptr();
+    let iter = src
+        .into_iter()
+        .flat_map(|a| {
+            a.into_iter().map(|b| b.wrapping_add(1))
+        });
+    assert_in_place_trait(&iter);
+    let sink = iter.collect::<Vec<_>>();
+    assert_eq!(srcptr as *const u8, sink.as_ptr());
 }
 
 #[test]

library/core/src/iter/adapters/flatten.rs

@@ -1,12 +1,13 @@
 use crate::iter::adapters::SourceIter;
 use crate::iter::{
-    DoubleEndedIterator, Fuse, FusedIterator, InPlaceIterable, Iterator, Map, TrustedFused,
-    TrustedLen,
+    Cloned, Copied, DoubleEndedIterator, Filter, FilterMap, Fuse, FusedIterator, InPlaceIterable,
+    Iterator, Map, TrustedFused, TrustedLen,
 };
+use crate::iter::{Once, OnceWith};
 use crate::num::NonZeroUsize;
 use crate::ops::{ControlFlow, Try};
-use crate::{fmt, option};
-use core::iter::Once;
+use crate::result;
+use crate::{array, fmt, option};
 
 /// An iterator that maps each element to an iterator, and yields the elements
 /// of the produced iterators.
@@ -154,10 +155,10 @@ where
 unsafe impl<I, U, F> InPlaceIterable for FlatMap<I, U, F>
 where
     I: InPlaceIterable,
-    U: KnownExpansionFactor + IntoIterator,
+    U: BoundedSize + IntoIterator,
 {
     const EXPAND_BY: Option<NonZeroUsize> = const {
-        match (I::EXPAND_BY, U::FACTOR) {
+        match (I::EXPAND_BY, U::UPPER_BOUND) {
             (Some(m), Some(n)) => m.checked_mul(n),
             _ => None,
         }
@@ -180,16 +181,59 @@ where
     }
 }
 
+/// Marker trait for iterators/iterables which have a statically known upper
+/// bound of the number of items they can produce.
+///
+/// # Safety
+///
+/// Implementations must not yield more elements than indicated by UPPER_BOUND if it is `Some`.
+/// Used in specializations.  Implementations must not be conditional on lifetimes or
+/// user-implementable traits.
 #[rustc_specialization_trait]
-trait KnownExpansionFactor {
-    const FACTOR: Option<NonZeroUsize> = NonZeroUsize::new(1);
+#[unstable(issue = "none", feature = "inplace_iteration")]
+unsafe trait BoundedSize {
+    const UPPER_BOUND: Option<NonZeroUsize> = NonZeroUsize::new(1);
 }
 
-impl<T> KnownExpansionFactor for Option<T> {}
-impl<T> KnownExpansionFactor for option::IntoIter<T> {}
-impl<T> KnownExpansionFactor for Once<T> {}
-impl<T, const N: usize> KnownExpansionFactor for [T; N] {
-    const FACTOR: Option<NonZeroUsize> = NonZeroUsize::new(N);
+#[unstable(issue = "none", feature = "inplace_iteration")]
+unsafe impl<T> BoundedSize for Option<T> {}
+#[unstable(issue = "none", feature = "inplace_iteration")]
+unsafe impl<T> BoundedSize for option::IntoIter<T> {}
+#[unstable(issue = "none", feature = "inplace_iteration")]
+unsafe impl<T, U> BoundedSize for Result<T, U> {}
+#[unstable(issue = "none", feature = "inplace_iteration")]
+unsafe impl<T> BoundedSize for result::IntoIter<T> {}
+#[unstable(issue = "none", feature = "inplace_iteration")]
+unsafe impl<T> BoundedSize for Once<T> {}
+#[unstable(issue = "none", feature = "inplace_iteration")]
+unsafe impl<T> BoundedSize for OnceWith<T> {}
+#[unstable(issue = "none", feature = "inplace_iteration")]
+unsafe impl<T, const N: usize> BoundedSize for [T; N] {
+    const UPPER_BOUND: Option<NonZeroUsize> = NonZeroUsize::new(N);
+}
+#[unstable(issue = "none", feature = "inplace_iteration")]
+unsafe impl<T, const N: usize> BoundedSize for array::IntoIter<T, N> {
+    const UPPER_BOUND: Option<NonZeroUsize> = NonZeroUsize::new(N);
+}
+#[unstable(issue = "none", feature = "inplace_iteration")]
+unsafe impl<I: BoundedSize, P> BoundedSize for Filter<I, P> {
+    const UPPER_BOUND: Option<NonZeroUsize> = I::UPPER_BOUND;
+}
+#[unstable(issue = "none", feature = "inplace_iteration")]
+unsafe impl<I: BoundedSize, P> BoundedSize for FilterMap<I, P> {
+    const UPPER_BOUND: Option<NonZeroUsize> = I::UPPER_BOUND;
+}
+#[unstable(issue = "none", feature = "inplace_iteration")]
+unsafe impl<I: BoundedSize, F> BoundedSize for Map<I, F> {
+    const UPPER_BOUND: Option<NonZeroUsize> = I::UPPER_BOUND;
+}
+#[unstable(issue = "none", feature = "inplace_iteration")]
+unsafe impl<I: BoundedSize> BoundedSize for Copied<I> {
+    const UPPER_BOUND: Option<NonZeroUsize> = I::UPPER_BOUND;
+}
+#[unstable(issue = "none", feature = "inplace_iteration")]
+unsafe impl<I: BoundedSize> BoundedSize for Cloned<I> {
+    const UPPER_BOUND: Option<NonZeroUsize> = I::UPPER_BOUND;
 }
 
 /// An iterator that flattens one level of nesting in an iterator of things
@@ -340,10 +384,10 @@ where
 unsafe impl<I> InPlaceIterable for Flatten<I>
 where
     I: InPlaceIterable + Iterator,
-    <I as Iterator>::Item: IntoIterator + KnownExpansionFactor,
+    <I as Iterator>::Item: IntoIterator + BoundedSize,
 {
     const EXPAND_BY: Option<NonZeroUsize> = const {
-        match (I::EXPAND_BY, I::Item::FACTOR) {
+        match (I::EXPAND_BY, I::Item::UPPER_BOUND) {
             (Some(m), Some(n)) => m.checked_mul(n),
             _ => None,
         }