From 64c915e3ade66d3335c445eed89509dc41502c13 Mon Sep 17 00:00:00 2001 From: Andre Bogus Date: Sat, 2 Feb 2019 01:29:39 +0100 Subject: [PATCH] override `VecDeque::try_rfold`, also update iterator This keeps the slice based iteration and updates the iterator state after each slice. It also uses a loop to reduce the amount of code. This uses unsafe code, so some thorough review would be appreciated. --- src/liballoc/benches/vec_deque.rs | 7 +++ src/liballoc/collections/vec_deque.rs | 51 ++++++++++++++++++--- src/liballoc/tests/vec_deque.rs | 64 +++++++++++++++++++++++++++ 3 files changed, 117 insertions(+), 5 deletions(-) diff --git a/src/liballoc/benches/vec_deque.rs b/src/liballoc/benches/vec_deque.rs index f7aadbdbd8266..7d2d3cfa61225 100644 --- a/src/liballoc/benches/vec_deque.rs +++ b/src/liballoc/benches/vec_deque.rs @@ -45,3 +45,10 @@ fn bench_mut_iter_1000(b: &mut Bencher) { black_box(sum); }) } + +#[bench] +fn bench_try_fold(b: &mut Bencher) { + let ring: VecDeque<_> = (0..1000).collect(); + + b.iter(|| black_box(ring.iter().try_fold(0, |a, b| Some(a + b)))) +} diff --git a/src/liballoc/collections/vec_deque.rs b/src/liballoc/collections/vec_deque.rs index 99fa54acb0836..3d91fd3daab86 100644 --- a/src/liballoc/collections/vec_deque.rs +++ b/src/liballoc/collections/vec_deque.rs @@ -2170,12 +2170,29 @@ impl<'a, T> Iterator for Iter<'a, T> { back.iter().fold(accum, &mut f) } - fn try_fold(&mut self, init: B, mut f: F) -> R where - Self: Sized, F: FnMut(B, Self::Item) -> R, R: Try + fn try_fold(&mut self, init: B, mut f: F) -> R + where + Self: Sized, + F: FnMut(B, Self::Item) -> R, + R: Try, { - let (front, back) = RingSlices::ring_slices(self.ring, self.head, self.tail); - let accum = front.iter().try_fold(init, &mut f)?; - back.iter().try_fold(accum, &mut f) + let (mut iter, final_res); + if self.tail <= self.head { + // single slice self.ring[self.tail..self.head] + iter = self.ring[self.tail..self.head].iter(); + final_res = iter.try_fold(init, &mut f); + } else { + // two slices: self.ring[self.tail..], self.ring[..self.head] + let (front, back) = self.ring.split_at(self.tail); + let mut back_iter = back.iter(); + let res = back_iter.try_fold(init, &mut f); + let len = self.ring.len(); + self.tail = (self.ring.len() - back_iter.len()) & (len - 1); + iter = front[..self.head].iter(); + final_res = iter.try_fold(res?, &mut f); + } + self.tail = self.head - iter.len(); + final_res } } @@ -2197,6 +2214,30 @@ impl<'a, T> DoubleEndedIterator for Iter<'a, T> { accum = back.iter().rfold(accum, &mut f); front.iter().rfold(accum, &mut f) } + + fn try_rfold(&mut self, init: B, mut f: F) -> R + where + Self: Sized, + F: FnMut(B, Self::Item) -> R, + R: Try, + { + let (mut iter, final_res); + if self.tail <= self.head { + // single slice self.ring[self.tail..self.head] + iter = self.ring[self.tail..self.head].iter(); + final_res = iter.try_rfold(init, &mut f); + } else { + // two slices: self.ring[self.tail..], self.ring[..self.head] + let (front, back) = self.ring.split_at(self.tail); + let mut front_iter = front[..self.head].iter(); + let res = front_iter.try_rfold(init, &mut f); + self.head = front_iter.len(); + iter = back.iter(); + final_res = iter.try_rfold(res?, &mut f); + } + self.head = self.tail + iter.len(); + final_res + } } #[stable(feature = "rust1", since = "1.0.0")] diff --git a/src/liballoc/tests/vec_deque.rs b/src/liballoc/tests/vec_deque.rs index aa49bdb009086..f801e42006bbb 100644 --- a/src/liballoc/tests/vec_deque.rs +++ b/src/liballoc/tests/vec_deque.rs @@ -1459,6 +1459,15 @@ fn test_try_fold_unit() { assert_eq!(Some(()), v.into_iter().try_fold((), |(), ()| Some(()))); } + +#[test] +fn test_try_fold_unit_none() { + let v: std::collections::VecDeque<()> = [(); 10].iter().cloned().collect(); + let mut iter = v.into_iter(); + assert!(iter.try_fold((), |_, _| None).is_none()); + assert_eq!(iter.len(), 9); +} + #[test] fn test_try_fold_rotated() { let mut v: VecDeque<_> = (0..12).collect(); @@ -1471,3 +1480,58 @@ fn test_try_fold_rotated() { assert_eq!(Ok::<_, ()>(66), v.iter().try_fold(0, |a, b| Ok(a + b))); } } + +#[test] +fn test_try_fold_moves_iter() { + let v: VecDeque<_> = [10, 20, 30, 40, 100, 60, 70, 80, 90].iter().collect(); + let mut iter = v.into_iter(); + assert_eq!(iter.try_fold(0_i8, |acc, &x| acc.checked_add(x)), None); + assert_eq!(iter.next(), Some(&60)); +} + +#[test] +fn test_try_fold_exhaust_wrap() { + let mut v = VecDeque::with_capacity(7); + v.push_back(1); + v.push_back(1); + v.push_back(1); + v.pop_front(); + v.pop_front(); + let mut iter = v.iter(); + let _ = iter.try_fold(0, |_, _| Some(1)); + assert!(iter.is_empty()); +} + +#[test] +fn test_try_fold_wraparound() { + let mut v = VecDeque::with_capacity(8); + v.push_back(7); + v.push_back(8); + v.push_back(9); + v.push_front(2); + v.push_front(1); + let mut iter = v.iter(); + let _ = iter.find(|&&x| x == 2); + assert_eq!(Some(&7), iter.next()); +} + +#[test] +fn test_try_rfold_rotated() { + let mut v: VecDeque<_> = (0..12).collect(); + for n in 0..10 { + if n & 1 == 0 { + v.rotate_left(n); + } else { + v.rotate_right(n); + } + assert_eq!(Ok::<_, ()>(66), v.iter().try_rfold(0, |a, b| Ok(a + b))); + } +} + +#[test] +fn test_try_rfold_moves_iter() { + let v : VecDeque<_> = [10, 20, 30, 40, 100, 60, 70, 80, 90].iter().collect(); + let mut iter = v.into_iter(); + assert_eq!(iter.try_rfold(0_i8, |acc, &x| acc.checked_add(x)), None); + assert_eq!(iter.next_back(), Some(&70)); +}