Skip to content

Commit

Permalink
Rollup merge of rust-lang#135104 - the8472:disable-in-place-iter-for-…
Browse files Browse the repository at this point in the history
…flatten, r=Mark-Simulacrum

do not in-place-iterate over flatmap/flatten

The implementation is unsound when a partially consumed iterator has some elements buffered in the front/back parts and cloning the Iterator removes the capacity from the backing vec::IntoIter.

This is a fix for rust-lang#135103 that removes the specialization trait impls without removing some supporting parts. I've kept it small so it can be easily backported. I'll either remove the remaining parts or think of a way to recover the optimization in a separate PR.
  • Loading branch information
workingjubilee authored Jan 5, 2025
2 parents dcb8be8 + 1ed0ea4 commit 5be6c9b
Show file tree
Hide file tree
Showing 2 changed files with 22 additions and 51 deletions.
39 changes: 20 additions & 19 deletions library/alloc/tests/vec.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1204,22 +1204,16 @@ fn test_from_iter_specialization_with_iterator_adapters() {
#[test]
fn test_in_place_specialization_step_up_down() {
fn assert_in_place_trait<T: InPlaceIterable>(_: &T) {}
let src = vec![[0u8; 4]; 256];
let srcptr = src.as_ptr();
let src_cap = src.capacity();
let iter = src.into_iter().flatten();
assert_in_place_trait(&iter);
let sink = iter.collect::<Vec<_>>();
let sinkptr = sink.as_ptr();
assert_eq!(srcptr as *const u8, sinkptr);
assert_eq!(src_cap * 4, sink.capacity());

let iter = sink.into_iter().array_chunks::<4>();
let src = vec![0u8; 1024];
let srcptr = src.as_ptr();
let src_bytes = src.capacity();
let iter = src.into_iter().array_chunks::<4>();
assert_in_place_trait(&iter);
let sink = iter.collect::<Vec<_>>();
let sinkptr = sink.as_ptr();
assert_eq!(srcptr, sinkptr);
assert_eq!(src_cap, sink.capacity());
assert_eq!(srcptr.addr(), sinkptr.addr());
assert_eq!(src_bytes, sink.capacity() * 4);

let mut src: Vec<u8> = Vec::with_capacity(17);
let src_bytes = src.capacity();
Expand All @@ -1236,13 +1230,6 @@ fn test_in_place_specialization_step_up_down() {
let sink: Vec<[u8; 2]> = iter.collect();
assert_eq!(sink.len(), 8);
assert!(sink.capacity() <= 25);

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]
Expand Down Expand Up @@ -1350,6 +1337,20 @@ fn test_collect_after_iterator_clone() {
assert_eq!(v, [1, 1, 1, 1, 1]);
assert!(v.len() <= v.capacity());
}

// regression test for #135103, similar to the one above Flatten/FlatMap had an unsound InPlaceIterable
// implementation.
#[test]
fn test_flatten_clone() {
const S: String = String::new();

let v = vec![[S, "Hello World!".into()], [S, S]];
let mut i = v.into_iter().flatten();
let _ = i.next();
let result: Vec<String> = i.clone().collect();
assert_eq!(result, ["Hello World!", "", ""]);
}

#[test]
fn test_cow_from() {
let borrowed: &[_] = &["borrowed", "(slice)"];
Expand Down
34 changes: 2 additions & 32 deletions library/core/src/iter/adapters/flatten.rs
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
use crate::iter::adapters::SourceIter;
use crate::iter::{
Cloned, Copied, Empty, Filter, FilterMap, Fuse, FusedIterator, InPlaceIterable, Map, Once,
OnceWith, TrustedFused, TrustedLen,
Cloned, Copied, Empty, Filter, FilterMap, Fuse, FusedIterator, Map, Once, OnceWith,
TrustedFused, TrustedLen,
};
use crate::num::NonZero;
use crate::ops::{ControlFlow, Try};
Expand Down Expand Up @@ -157,21 +157,6 @@ where
{
}

#[unstable(issue = "none", feature = "inplace_iteration")]
unsafe impl<I, U, F> InPlaceIterable for FlatMap<I, U, F>
where
I: InPlaceIterable,
U: BoundedSize + IntoIterator,
{
const EXPAND_BY: Option<NonZero<usize>> = const {
match (I::EXPAND_BY, U::UPPER_BOUND) {
(Some(m), Some(n)) => m.checked_mul(n),
_ => None,
}
};
const MERGE_BY: Option<NonZero<usize>> = I::MERGE_BY;
}

#[unstable(issue = "none", feature = "inplace_iteration")]
unsafe impl<I, U, F> SourceIter for FlatMap<I, U, F>
where
Expand Down Expand Up @@ -386,21 +371,6 @@ where
{
}

#[unstable(issue = "none", feature = "inplace_iteration")]
unsafe impl<I> InPlaceIterable for Flatten<I>
where
I: InPlaceIterable + Iterator,
<I as Iterator>::Item: IntoIterator + BoundedSize,
{
const EXPAND_BY: Option<NonZero<usize>> = const {
match (I::EXPAND_BY, I::Item::UPPER_BOUND) {
(Some(m), Some(n)) => m.checked_mul(n),
_ => None,
}
};
const MERGE_BY: Option<NonZero<usize>> = I::MERGE_BY;
}

#[unstable(issue = "none", feature = "inplace_iteration")]
unsafe impl<I> SourceIter for Flatten<I>
where
Expand Down

0 comments on commit 5be6c9b

Please sign in to comment.