about summary refs log tree commit diff
path: root/library/core/src/array
diff options
context:
space:
mode:
authorScott McMurray <scottmcm@users.noreply.github.com>2021-12-03 21:36:51 -0800
committerScott McMurray <scottmcm@users.noreply.github.com>2021-12-03 21:36:51 -0800
commiteb846dbaca87f156419724e77fd52eb91ca01161 (patch)
tree5730e03b0e05e3e6b0267aabbd1139f1cbbb9d0a /library/core/src/array
parent2a9e0831d6603d87220cedd1b1293e2eb82ef55c (diff)
downloadrust-eb846dbaca87f156419724e77fd52eb91ca01161.tar.gz
rust-eb846dbaca87f156419724e77fd52eb91ca01161.zip
Override `Iterator::advance(_back)_by` for `array::IntoIter`
Because I happened to notice that `nth` is currently getting codegen'd as a loop even for `Copy` types: <https://rust.godbolt.org/z/fPqv7Gvs7>
Diffstat (limited to 'library/core/src/array')
-rw-r--r--library/core/src/array/iter.rs44
1 files changed, 43 insertions, 1 deletions
diff --git a/library/core/src/array/iter.rs b/library/core/src/array/iter.rs
index 5d63cf03fcb..a146966c674 100644
--- a/library/core/src/array/iter.rs
+++ b/library/core/src/array/iter.rs
@@ -1,7 +1,7 @@
 //! Defines the `IntoIter` owned iterator for arrays.
 
 use crate::{
-    fmt,
+    cmp, fmt,
     iter::{self, ExactSizeIterator, FusedIterator, TrustedLen},
     mem::{self, MaybeUninit},
     ops::Range,
@@ -150,6 +150,27 @@ impl<T, const N: usize> Iterator for IntoIter<T, N> {
     fn last(mut self) -> Option<Self::Item> {
         self.next_back()
     }
+
+    fn advance_by(&mut self, n: usize) -> Result<(), usize> {
+        let len = self.len();
+
+        // The number of elements to drop.  Always in-bounds by construction.
+        let delta = cmp::min(n, len);
+
+        let range_to_drop = self.alive.start..(self.alive.start + delta);
+
+        // Moving the start marks them as conceptually "dropped", so if anything
+        // goes bad then our drop impl won't double-free them.
+        self.alive.start += delta;
+
+        // SAFETY: These elements are currently initialized, so it's fine to drop them.
+        unsafe {
+            let slice = self.data.get_unchecked_mut(range_to_drop);
+            ptr::drop_in_place(MaybeUninit::slice_assume_init_mut(slice));
+        }
+
+        if n > len { Err(len) } else { Ok(()) }
+    }
 }
 
 #[stable(feature = "array_value_iter_impls", since = "1.40.0")]
@@ -170,6 +191,27 @@ impl<T, const N: usize> DoubleEndedIterator for IntoIter<T, N> {
             unsafe { self.data.get_unchecked(idx).assume_init_read() }
         })
     }
+
+    fn advance_back_by(&mut self, n: usize) -> Result<(), usize> {
+        let len = self.len();
+
+        // The number of elements to drop.  Always in-bounds by construction.
+        let delta = cmp::min(n, len);
+
+        let range_to_drop = (self.alive.end - delta)..self.alive.end;
+
+        // Moving the end marks them as conceptually "dropped", so if anything
+        // goes bad then our drop impl won't double-free them.
+        self.alive.end -= delta;
+
+        // SAFETY: These elements are currently initialized, so it's fine to drop them.
+        unsafe {
+            let slice = self.data.get_unchecked_mut(range_to_drop);
+            ptr::drop_in_place(MaybeUninit::slice_assume_init_mut(slice));
+        }
+
+        if n > len { Err(len) } else { Ok(()) }
+    }
 }
 
 #[stable(feature = "array_value_iter_impls", since = "1.40.0")]