From 51d546f4aa8a94b81d2a580518d95d1ab12a3655 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sebastian=20Dr=C3=B6ge?= Date: Tue, 2 Jan 2018 02:13:20 +0200 Subject: [PATCH] Add slice::ExactChunks and ::ExactChunksMut iterators These guarantee that always the requested slice size will be returned and any leftoever elements at the end will be ignored. It allows llvm to get rid of bounds checks in the code using the iterator. This is inspired by the same iterators provided by ndarray. See https://github.com/rust-lang/rust/issues/47115 --- src/liballoc/lib.rs | 1 + src/liballoc/slice.rs | 56 ++++++++++ src/libcore/slice/mod.rs | 230 +++++++++++++++++++++++++++++++++++++++ 3 files changed, 287 insertions(+) diff --git a/src/liballoc/lib.rs b/src/liballoc/lib.rs index 3cc3ea467966..d8ce28695ab6 100644 --- a/src/liballoc/lib.rs +++ b/src/liballoc/lib.rs @@ -124,6 +124,7 @@ #![feature(unsize)] #![feature(allocator_internals)] #![feature(on_unimplemented)] +#![feature(exact_chunks)] #![cfg_attr(not(test), feature(fused, fn_traits, placement_new_protocol, swap_with_slice, i128))] #![cfg_attr(test, feature(test, box_heap))] diff --git a/src/liballoc/slice.rs b/src/liballoc/slice.rs index fa73197885be..bae366736379 100644 --- a/src/liballoc/slice.rs +++ b/src/liballoc/slice.rs @@ -123,6 +123,8 @@ pub use core::slice::{from_raw_parts, from_raw_parts_mut}; pub use core::slice::{from_ref, from_ref_mut}; #[unstable(feature = "slice_get_slice", issue = "35729")] pub use core::slice::SliceIndex; +#[unstable(feature = "exact_chunks", issue = "47115")] +pub use core::slice::{ExactChunks, ExactChunksMut}; //////////////////////////////////////////////////////////////////////////////// // Basic slice extension methods @@ -631,6 +633,31 @@ impl [T] { core_slice::SliceExt::chunks(self, chunk_size) } + /// Returns an iterator over `chunk_size` elements of the slice at a + /// time. The chunks are slices and do not overlap. If `chunk_size` does + /// not divide the length of the slice, then the last up to `chunk_size-1` + /// elements will be omitted. + /// + /// # Panics + /// + /// Panics if `chunk_size` is 0. + /// + /// # Examples + /// + /// ``` + /// let slice = ['l', 'o', 'r', 'e', 'm']; + /// let mut iter = slice.chunks(2); + /// assert_eq!(iter.next().unwrap(), &['l', 'o']); + /// assert_eq!(iter.next().unwrap(), &['r', 'e']); + /// assert_eq!(iter.next().unwrap(), &['m']); + /// assert!(iter.next().is_none()); + /// ``` + #[unstable(feature = "exact_chunks", issue = "47115")] + #[inline] + pub fn exact_chunks(&self, chunk_size: usize) -> ExactChunks { + core_slice::SliceExt::exact_chunks(self, chunk_size) + } + /// Returns an iterator over `chunk_size` elements of the slice at a time. /// The chunks are mutable slices, and do not overlap. If `chunk_size` does /// not divide the length of the slice, then the last chunk will not @@ -660,6 +687,35 @@ impl [T] { core_slice::SliceExt::chunks_mut(self, chunk_size) } + /// Returns an iterator over `chunk_size` elements of the slice at a time. + /// The chunks are mutable slices, and do not overlap. If `chunk_size` does + /// not divide the length of the slice, then the last up to `chunk_size-1` + /// elements will be omitted. + /// + /// # Panics + /// + /// Panics if `chunk_size` is 0. + /// + /// # Examples + /// + /// ``` + /// let v = &mut [0, 0, 0, 0, 0]; + /// let mut count = 1; + /// + /// for chunk in v.exact_chunks_mut(2) { + /// for elem in chunk.iter_mut() { + /// *elem += count; + /// } + /// count += 1; + /// } + /// assert_eq!(v, &[1, 1, 2, 2, 3]); + /// ``` + #[unstable(feature = "exact_chunks", issue = "47115")] + #[inline] + pub fn exact_chunks_mut(&mut self, chunk_size: usize) -> ExactChunksMut { + core_slice::SliceExt::exact_chunks_mut(self, chunk_size) + } + /// Divides one slice into two at an index. /// /// The first will contain all indices from `[0, mid)` (excluding diff --git a/src/libcore/slice/mod.rs b/src/libcore/slice/mod.rs index 72036d6d3a24..5791b1d6e79c 100644 --- a/src/libcore/slice/mod.rs +++ b/src/libcore/slice/mod.rs @@ -104,6 +104,9 @@ pub trait SliceExt { #[stable(feature = "core", since = "1.6.0")] fn chunks(&self, size: usize) -> Chunks; + #[unstable(feature = "exact_chunks", issue = "47115")] + fn exact_chunks(&self, size: usize) -> ExactChunks; + #[stable(feature = "core", since = "1.6.0")] fn get(&self, index: I) -> Option<&I::Output> where I: SliceIndex; @@ -181,6 +184,9 @@ pub trait SliceExt { #[stable(feature = "core", since = "1.6.0")] fn chunks_mut(&mut self, chunk_size: usize) -> ChunksMut; + #[unstable(feature = "exact_chunks", issue = "47115")] + fn exact_chunks_mut(&mut self, size: usize) -> ExactChunksMut; + #[stable(feature = "core", since = "1.6.0")] fn swap(&mut self, a: usize, b: usize); @@ -353,6 +359,14 @@ impl SliceExt for [T] { Chunks { v: self, chunk_size: chunk_size } } + #[inline] + fn exact_chunks(&self, chunk_size: usize) -> ExactChunks { + assert!(chunk_size != 0); + let rem = self.len() % chunk_size; + let len = self.len() - rem; + ExactChunks { v: &self[..len], chunk_size: chunk_size} + } + #[inline] fn get(&self, index: I) -> Option<&I::Output> where I: SliceIndex<[T]> @@ -536,6 +550,14 @@ impl SliceExt for [T] { ChunksMut { v: self, chunk_size: chunk_size } } + #[inline] + fn exact_chunks_mut(&mut self, chunk_size: usize) -> ExactChunksMut { + assert!(chunk_size != 0); + let rem = self.len() % chunk_size; + let len = self.len() - rem; + ExactChunksMut { v: &mut self[..len], chunk_size: chunk_size} + } + #[inline] fn swap(&mut self, a: usize, b: usize) { unsafe { @@ -2365,6 +2387,214 @@ unsafe impl<'a, T> TrustedRandomAccess for ChunksMut<'a, T> { fn may_have_side_effect() -> bool { false } } +/// An iterator over a slice in (non-overlapping) chunks (`chunk_size` elements at a +/// time). +/// +/// When the slice len is not evenly divided by the chunk size, the last +/// up to `chunk_size-1` elements will be omitted. +/// +/// This struct is created by the [`exact_chunks`] method on [slices]. +/// +/// [`exact_chunks`]: ../../std/primitive.slice.html#method.exact_chunks +/// [slices]: ../../std/primitive.slice.html +#[derive(Debug)] +#[unstable(feature = "exact_chunks", issue = "47115")] +pub struct ExactChunks<'a, T:'a> { + v: &'a [T], + chunk_size: usize +} + +// FIXME(#26925) Remove in favor of `#[derive(Clone)]` +#[unstable(feature = "exact_chunks", issue = "47115")] +impl<'a, T> Clone for ExactChunks<'a, T> { + fn clone(&self) -> ExactChunks<'a, T> { + ExactChunks { + v: self.v, + chunk_size: self.chunk_size, + } + } +} + +#[unstable(feature = "exact_chunks", issue = "47115")] +impl<'a, T> Iterator for ExactChunks<'a, T> { + type Item = &'a [T]; + + #[inline] + fn next(&mut self) -> Option<&'a [T]> { + if self.v.len() < self.chunk_size { + None + } else { + let (fst, snd) = self.v.split_at(self.chunk_size); + self.v = snd; + Some(fst) + } + } + + #[inline] + fn size_hint(&self) -> (usize, Option) { + let n = self.v.len() / self.chunk_size; + (n, Some(n)) + } + + #[inline] + fn count(self) -> usize { + self.len() + } + + #[inline] + fn nth(&mut self, n: usize) -> Option { + let (start, overflow) = n.overflowing_mul(self.chunk_size); + if start >= self.v.len() || overflow { + self.v = &[]; + None + } else { + let end = match start.checked_add(self.chunk_size) { + Some(sum) => cmp::min(self.v.len(), sum), + None => self.v.len(), + }; + + if end - start != self.chunk_size { + self.v = &[]; + None + } else { + let nth = &self.v[start..end]; + self.v = &self.v[end..]; + Some(nth) + } + } + } + + #[inline] + fn last(self) -> Option { + if self.v.len() < self.chunk_size { + None + } else { + let start = self.v.len() - self.chunk_size; + Some(&self.v[start..]) + } + } +} + +#[unstable(feature = "exact_chunks", issue = "47115")] +impl<'a, T> DoubleEndedIterator for ExactChunks<'a, T> { + #[inline] + fn next_back(&mut self) -> Option<&'a [T]> { + if self.v.len() < self.chunk_size { + None + } else { + let (fst, snd) = self.v.split_at(self.v.len() - self.chunk_size); + self.v = fst; + Some(snd) + } + } +} + +#[unstable(feature = "exact_chunks", issue = "47115")] +impl<'a, T> ExactSizeIterator for ExactChunks<'a, T> {} + +#[unstable(feature = "fused", issue = "35602")] +impl<'a, T> FusedIterator for ExactChunks<'a, T> {} + +/// An iterator over a slice in (non-overlapping) mutable chunks (`chunk_size` +/// elements at a time). When the slice len is not evenly divided by the chunk +/// size, the last up to `chunk_size-1` elements will be omitted. +/// +/// This struct is created by the [`exact_chunks_mut`] method on [slices]. +/// +/// [`exact_chunks_mut`]: ../../std/primitive.slice.html#method.exact_chunks_mut +/// [slices]: ../../std/primitive.slice.html +#[derive(Debug)] +#[unstable(feature = "exact_chunks", issue = "47115")] +pub struct ExactChunksMut<'a, T:'a> { + v: &'a mut [T], + chunk_size: usize +} + +#[unstable(feature = "exact_chunks", issue = "47115")] +impl<'a, T> Iterator for ExactChunksMut<'a, T> { + type Item = &'a mut [T]; + + #[inline] + fn next(&mut self) -> Option<&'a mut [T]> { + if self.v.len() < self.chunk_size { + None + } else { + let tmp = mem::replace(&mut self.v, &mut []); + let (head, tail) = tmp.split_at_mut(self.chunk_size); + self.v = tail; + Some(head) + } + } + + #[inline] + fn size_hint(&self) -> (usize, Option) { + let n = self.v.len() / self.chunk_size; + (n, Some(n)) + } + + #[inline] + fn count(self) -> usize { + self.len() + } + + #[inline] + fn nth(&mut self, n: usize) -> Option<&'a mut [T]> { + let (start, overflow) = n.overflowing_mul(self.chunk_size); + if start >= self.v.len() || overflow { + self.v = &mut []; + None + } else { + let end = match start.checked_add(self.chunk_size) { + Some(sum) => cmp::min(self.v.len(), sum), + None => self.v.len(), + }; + + if end - start != self.chunk_size { + self.v = &mut []; + None + } else { + let tmp = mem::replace(&mut self.v, &mut []); + let (head, tail) = tmp.split_at_mut(end); + let (_, nth) = head.split_at_mut(start); + self.v = tail; + Some(nth) + } + } + } + + #[inline] + fn last(self) -> Option { + if self.v.len() < self.chunk_size { + None + } else { + let start = (self.v.len() - self.chunk_size) / self.chunk_size * self.chunk_size; + Some(&mut self.v[start..]) + } + } +} + +#[unstable(feature = "exact_chunks", issue = "47115")] +impl<'a, T> DoubleEndedIterator for ExactChunksMut<'a, T> { + #[inline] + fn next_back(&mut self) -> Option<&'a mut [T]> { + if self.v.len() < self.chunk_size { + None + } else { + let tmp = mem::replace(&mut self.v, &mut []); + let tmp_len = tmp.len(); + let (head, tail) = tmp.split_at_mut(tmp_len - self.chunk_size); + self.v = head; + Some(tail) + } + } +} + +#[unstable(feature = "exact_chunks", issue = "47115")] +impl<'a, T> ExactSizeIterator for ExactChunksMut<'a, T> {} + +#[unstable(feature = "fused", issue = "35602")] +impl<'a, T> FusedIterator for ExactChunksMut<'a, T> {} + // // Free functions //