| // Copyright 2022 The Fuchsia Authors. All rights reserved. |
| // Use of this source code is governed by a BSD-style license that can be |
| // found in the LICENSE file. |
| |
| use rayon::{ |
| iter::plumbing::{bridge, Consumer, Producer, ProducerCallback, UnindexedConsumer}, |
| prelude::*, |
| }; |
| |
| #[derive(Clone, Debug, Default)] |
| pub struct GroupedIter<'s> { |
| sums: &'s [u32], |
| group_start: u32, |
| group_end: u32, |
| start: u32, |
| end: u32, |
| } |
| |
| impl Iterator for GroupedIter<'_> { |
| type Item = (u32, u32); |
| |
| #[inline] |
| fn next(&mut self) -> Option<Self::Item> { |
| loop { |
| if self.start >= self.end { |
| return None; |
| } |
| |
| let exclusive_sum = |
| self.group_start.checked_sub(1).map(|i| self.sums[i as usize]).unwrap_or_default(); |
| let inclusive_sum = self.sums[self.group_start as usize]; |
| |
| if exclusive_sum == inclusive_sum { |
| self.group_start += 1; |
| continue; |
| } |
| |
| let result = (self.group_start, self.start as u32 - exclusive_sum); |
| |
| self.start += 1; |
| |
| if self.start == inclusive_sum { |
| self.group_start += 1; |
| } |
| |
| return Some(result); |
| } |
| } |
| } |
| |
| impl DoubleEndedIterator for GroupedIter<'_> { |
| #[inline] |
| fn next_back(&mut self) -> Option<Self::Item> { |
| loop { |
| if self.start >= self.end { |
| return None; |
| } |
| |
| let exclusive_sum = |
| self.group_end.checked_sub(1).map(|i| self.sums[i as usize]).unwrap_or_default(); |
| let inclusive_sum = self.sums[self.group_end as usize]; |
| |
| if exclusive_sum == inclusive_sum { |
| self.group_end = self.group_end.saturating_sub(1); |
| continue; |
| } |
| |
| let result = (self.group_end, self.end as u32 - 1 - exclusive_sum); |
| |
| self.end -= 1; |
| |
| if self.end == exclusive_sum { |
| self.group_end = self.group_end.saturating_sub(1); |
| } |
| |
| return Some(result); |
| } |
| } |
| } |
| |
| impl ExactSizeIterator for GroupedIter<'_> { |
| fn len(&self) -> usize { |
| (self.end - self.start) as usize |
| } |
| } |
| |
| struct GroupedIterProducer<'s> { |
| inner: GroupedIter<'s>, |
| } |
| |
| impl<'s> Producer for GroupedIterProducer<'s> { |
| type Item = (u32, u32); |
| |
| type IntoIter = GroupedIter<'s>; |
| |
| #[inline] |
| fn into_iter(self) -> Self::IntoIter { |
| self.inner |
| } |
| |
| #[inline] |
| fn split_at(self, index: usize) -> (Self, Self) { |
| let index = index as u32 + self.inner.start; |
| |
| let mid = match self.inner.sums.binary_search(&(index as u32)) { |
| Ok(mid) => mid + 1, |
| Err(mid) => mid, |
| } as u32; |
| |
| ( |
| Self { |
| inner: GroupedIter { |
| sums: self.inner.sums, |
| group_start: self.inner.group_start, |
| group_end: mid as u32, |
| start: self.inner.start, |
| end: index as u32, |
| }, |
| }, |
| Self { |
| inner: GroupedIter { |
| sums: self.inner.sums, |
| group_start: mid as u32, |
| group_end: self.inner.group_end, |
| start: index as u32, |
| end: self.inner.end, |
| }, |
| }, |
| ) |
| } |
| } |
| |
| impl<'s> GroupedIter<'s> { |
| #[inline] |
| pub fn new(sums: &'s [u32]) -> Self { |
| Self { |
| sums, |
| group_start: 0, |
| group_end: sums.len().saturating_sub(1) as u32, |
| start: 0, |
| end: sums.last().copied().unwrap_or_default(), |
| } |
| } |
| } |
| |
| impl<'s> IntoParallelIterator for GroupedIter<'s> { |
| type Iter = GroupedParIter<'s>; |
| |
| type Item = (u32, u32); |
| |
| #[inline] |
| fn into_par_iter(self) -> Self::Iter { |
| GroupedParIter { iter: self } |
| } |
| } |
| |
| pub struct GroupedParIter<'s> { |
| iter: GroupedIter<'s>, |
| } |
| |
| impl ParallelIterator for GroupedParIter<'_> { |
| type Item = (u32, u32); |
| |
| #[inline] |
| fn drive_unindexed<C>(self, consumer: C) -> C::Result |
| where |
| C: UnindexedConsumer<Self::Item>, |
| { |
| bridge(self, consumer) |
| } |
| |
| #[inline] |
| fn opt_len(&self) -> Option<usize> { |
| Some(self.iter.len()) |
| } |
| } |
| |
| impl IndexedParallelIterator for GroupedParIter<'_> { |
| #[inline] |
| fn len(&self) -> usize { |
| self.iter.len() |
| } |
| |
| #[inline] |
| fn drive<C: Consumer<Self::Item>>(self, consumer: C) -> C::Result { |
| bridge(self, consumer) |
| } |
| |
| #[inline] |
| fn with_producer<CB: ProducerCallback<Self::Item>>(self, callback: CB) -> CB::Output { |
| callback.callback(GroupedIterProducer { inner: self.iter }) |
| } |
| } |
| |
| #[cfg(test)] |
| mod tests { |
| use super::*; |
| |
| #[test] |
| fn empty_iter() { |
| let sums = &[]; |
| |
| assert_eq!(GroupedIter::new(sums).collect::<Vec<_>>(), []); |
| } |
| |
| #[test] |
| fn local_iter() { |
| let sums = &[2, 5, 9, 15]; |
| let iter = GroupedIter::new(sums); |
| |
| assert_eq!( |
| iter.collect::<Vec<_>>(), |
| [ |
| (0, 0), |
| (0, 1), |
| (1, 0), |
| (1, 1), |
| (1, 2), |
| (2, 0), |
| (2, 1), |
| (2, 2), |
| (2, 3), |
| (3, 0), |
| (3, 1), |
| (3, 2), |
| (3, 3), |
| (3, 4), |
| (3, 5), |
| ] |
| ); |
| } |
| |
| #[test] |
| fn local_iter_rev() { |
| let sums = &[2, 5, 9, 15]; |
| let iter = GroupedIter::new(sums); |
| |
| assert_eq!( |
| iter.rev().collect::<Vec<_>>(), |
| [ |
| (3, 5), |
| (3, 4), |
| (3, 3), |
| (3, 2), |
| (3, 1), |
| (3, 0), |
| (2, 3), |
| (2, 2), |
| (2, 1), |
| (2, 0), |
| (1, 2), |
| (1, 1), |
| (1, 0), |
| (0, 1), |
| (0, 0), |
| ] |
| ); |
| } |
| |
| #[test] |
| fn both_ends() { |
| let sums = &[2, 5, 9, 15]; |
| let mut iter = GroupedIter::new(sums); |
| |
| assert_eq!(iter.len(), 15); |
| |
| assert_eq!(iter.next(), Some((0, 0))); |
| assert_eq!(iter.next_back(), Some((3, 5))); |
| assert_eq!(iter.next(), Some((0, 1))); |
| assert_eq!(iter.next_back(), Some((3, 4))); |
| assert_eq!(iter.next(), Some((1, 0))); |
| assert_eq!(iter.next_back(), Some((3, 3))); |
| assert_eq!(iter.next(), Some((1, 1))); |
| |
| assert_eq!(iter.len(), 8); |
| |
| assert_eq!(iter.next_back(), Some((3, 2))); |
| assert_eq!(iter.next(), Some((1, 2))); |
| assert_eq!(iter.next_back(), Some((3, 1))); |
| assert_eq!(iter.next(), Some((2, 0))); |
| assert_eq!(iter.next_back(), Some((3, 0))); |
| assert_eq!(iter.next(), Some((2, 1))); |
| assert_eq!(iter.next_back(), Some((2, 3))); |
| assert_eq!(iter.next(), Some((2, 2))); |
| assert_eq!(iter.next(), None); |
| assert_eq!(iter.next_back(), None); |
| |
| assert_eq!(iter.len(), 0); |
| } |
| |
| #[test] |
| fn empty_groups() { |
| let sums = &[2, 2, 5, 5]; |
| let iter = GroupedIter::new(sums); |
| |
| assert_eq!(iter.collect::<Vec<_>>(), [(0, 0), (0, 1), (2, 0), (2, 1), (2, 2),]); |
| } |
| |
| #[test] |
| fn empty_groups_rev() { |
| let sums = &[2, 2, 5, 5]; |
| let iter = GroupedIter::new(sums); |
| |
| assert_eq!(iter.rev().collect::<Vec<_>>(), [(2, 2), (2, 1), (2, 0), (0, 1), (0, 0),]); |
| } |
| |
| #[test] |
| fn par_iter() { |
| let sums = &[2, 5, 9, 15]; |
| let iter = GroupedIter::new(sums); |
| |
| assert_eq!( |
| iter.into_par_iter().collect::<Vec<_>>(), |
| [ |
| (0, 0), |
| (0, 1), |
| (1, 0), |
| (1, 1), |
| (1, 2), |
| (2, 0), |
| (2, 1), |
| (2, 2), |
| (2, 3), |
| (3, 0), |
| (3, 1), |
| (3, 2), |
| (3, 3), |
| (3, 4), |
| (3, 5), |
| ] |
| ); |
| } |
| |
| #[test] |
| fn par_iter2() { |
| let sums = &[3, 6, 10, 11]; |
| let iter = GroupedIter::new(sums); |
| |
| assert_eq!( |
| iter.into_par_iter().collect::<Vec<_>>(), |
| [ |
| (0, 0), |
| (0, 1), |
| (0, 2), |
| (1, 0), |
| (1, 1), |
| (1, 2), |
| (2, 0), |
| (2, 1), |
| (2, 2), |
| (2, 3), |
| (3, 0), |
| ] |
| ); |
| } |
| } |