support in-place collecting additional FlatMap shapes

This commit is contained in:
The 8472 2023-04-29 15:53:26 +02:00
parent 7047fb41b6
commit 439f63019f
2 changed files with 70 additions and 15 deletions

View File

@ -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]

View File

@ -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,
}