Tune RepeatWith::try_fold and Take::for_each and Vec::extend_trusted

Author: scottmcm

library/alloc/src/vec/mod.rs

@@ -2874,13 +2874,12 @@ impl<T, A: Allocator> Vec<T, A> {
  );
  self.reserve(additional);
  unsafe {
- let mut ptr = self.as_mut_ptr().add(self.len());
+ let ptr = self.as_mut_ptr();
  let mut local_len = SetLenOnDrop::new(&mut self.len);
  iterator.for_each(move |element| {
- ptr::write(ptr, element);
- ptr = ptr.add(1);
- // Since the loop executes user code which can panic we have to bump the pointer
- // after each step.
+ ptr::write(ptr.add(local_len.current_len()), element);
+ // Since the loop executes user code which can panic we have to update
+ // the length every step to correctly drop what we've written.
  // NB can't overflow since we would have had to alloc the address space
  local_len.increment_len(1);
  });

library/alloc/src/vec/set_len_on_drop.rs

@@ -18,6 +18,11 @@ impl<'a> SetLenOnDrop<'a> {
  pub(super) fn increment_len(&mut self, increment: usize) {
  self.local_len += increment;
  }
+
+ #[inline]
+ pub(super) fn current_len(&self) -> usize {
+ self.local_len
+ }
 }
 
 impl Drop for SetLenOnDrop<'_> {

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

@@ -75,7 +75,6 @@ where
  #[inline]
  fn try_fold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R
  where
- Self: Sized,
  Fold: FnMut(Acc, Self::Item) -> R,
  R: Try<Output = Acc>,
  {
@@ -100,6 +99,26 @@ where
 
  impl_fold_via_try_fold! { fold -> try_fold }
 
+ #[inline]
+ fn for_each<F: FnMut(Self::Item)>(mut self, f: F) {
+ // The default implementation would use a unit accumulator, so we can
+ // avoid a stateful closure by folding over the remaining number
+ // of items we wish to return instead.
+ fn check<'a, Item>(
+ mut action: impl FnMut(Item) + 'a,
+ ) -> impl FnMut(usize, Item) -> Option<usize> + 'a {
+ move |more, x| {
+ action(x);
+ more.checked_sub(1)
+ }
+ }
+
+ let remaining = self.n;
+ if remaining > 0 {
+ self.iter.try_fold(remaining - 1, check(f));
+ }
+ }
+
  #[inline]
  #[rustc_inherit_overflow_checks]
  fn advance_by(&mut self, n: usize) -> Result<(), usize> {

library/core/src/iter/sources/repeat_with.rs

@@ -1,4 +1,5 @@
 use crate::iter::{FusedIterator, TrustedLen};
+use crate::ops::Try;
 
 /// Creates a new iterator that repeats elements of type `A` endlessly by
 /// applying the provided closure, the repeater, `F: FnMut() -> A`.
@@ -89,6 +90,22 @@ impl<A, F: FnMut() -> A> Iterator for RepeatWith<F> {
  fn size_hint(&self) -> (usize, Option<usize>) {
  (usize::MAX, None)
  }
+
+ #[inline]
+ fn try_fold<Acc, Fold, R>(&mut self, mut init: Acc, mut fold: Fold) -> R
+ where
+ Fold: FnMut(Acc, Self::Item) -> R,
+ R: Try<Output = Acc>,
+ {
+ // This override isn't strictly needed, but avoids the need to optimize
+ // away the `next`-always-returns-`Some` and emphasizes that the `?`
+ // is the only way to exit the loop.
+
+ loop {
+ let item = (self.repeater)();
+ init = fold(init, item)?;
+ }
+ }
 }
 
 #[stable(feature = "iterator_repeat_with", since = "1.28.0")]

library/core/tests/iter/adapters/take.rs

@@ -146,3 +146,23 @@ fn test_take_try_folds() {
  assert_eq!(iter.try_for_each(Err), Err(2));
  assert_eq!(iter.try_for_each(Err), Ok(()));
 }
+
+#[test]
+fn test_byref_take_consumed_items() {
+ let mut inner = 10..90;
+
+ let mut count = 0;
+ inner.by_ref().take(0).for_each(|_| count += 1);
+ assert_eq!(count, 0);
+ assert_eq!(inner, 10..90);
+
+ let mut count = 0;
+ inner.by_ref().take(10).for_each(|_| count += 1);
+ assert_eq!(count, 10);
+ assert_eq!(inner, 20..90);
+
+ let mut count = 0;
+ inner.by_ref().take(100).for_each(|_| count += 1);
+ assert_eq!(count, 70);
+ assert_eq!(inner, 90..90);
+}

src/test/codegen/repeat-trusted-len.rs

@@ -11,3 +11,10 @@ pub fn repeat_take_collect() -> Vec<u8> {
 // CHECK: call void @llvm.memset.{{.+}}({{i8\*|ptr}} {{.*}}align 1{{.*}} %{{[0-9]+}}, i8 42, i{{[0-9]+}} 100000, i1 false)
  iter::repeat(42).take(100000).collect()
 }
+
+// CHECK-LABEL: @repeat_with_take_collect
+#[no_mangle]
+pub fn repeat_with_take_collect() -> Vec<u8> {
+// CHECK: call void @llvm.memset.{{.+}}({{i8\*|ptr}} {{.*}}align 1{{.*}} %{{[0-9]+}}, i8 13, i{{[0-9]+}} 12345, i1 false)
+ iter::repeat_with(|| 13).take(12345).collect()
+}