diff --git a/src/groupbylazy.rs b/src/groupbylazy.rs index af7512343..61d502261 100644 --- a/src/groupbylazy.rs +++ b/src/groupbylazy.rs @@ -2,6 +2,8 @@ use alloc::vec::{self, Vec}; use std::cell::{Cell, RefCell}; use std::fmt::Debug; +use crate::size_hint; + /// A trait to unify `FnMut` for `ChunkBy` with the chunk key in `IntoChunks` trait KeyFunction { type Key; @@ -637,11 +639,37 @@ where first: Some(elt), }) } + + fn size_hint(&self) -> (usize, Option) { + let inner = self.parent.inner.borrow(); + // chunk_size cannot be null as it is checked in `Itertools::chunks` + let chunk_size = inner.key.size; + debug_assert!(chunk_size != 0); + + if self.parent.index.get() == 0 { + // inner iter has not been used yet + size_hint::div_ceil_scalar(inner.iter.size_hint(), chunk_size) + } else { + // inner iter has already been used and its length is < original length + // so we need to return result of floor division because of `IntoChunks::step` + // impl which consumes only 1 element on first `next` and then `chunk_size` + // every nedxt time + size_hint::div_scalar(inner.iter.size_hint(), chunk_size) + } + } +} + +impl<'a, I> ExactSizeIterator for Chunks<'a, I> +where + I: ExactSizeIterator, + I::Item: 'a, +{ } /// An iterator for the elements in a single chunk. /// /// Iterator element type is `I::Item`. +#[must_use = "iterator adaptors are lazy and do nothing unless consumed"] #[derive(Debug)] pub struct Chunk<'a, I> where diff --git a/src/size_hint.rs b/src/size_hint.rs index 6cfead7f2..5657c6bd0 100644 --- a/src/size_hint.rs +++ b/src/size_hint.rs @@ -57,6 +57,28 @@ pub fn mul_scalar(sh: SizeHint, x: usize) -> SizeHint { (low, hi) } +/// Correct ceiling division by `x` with a `SizeHint`. +#[inline] +#[track_caller] +pub fn div_ceil_scalar(sh: SizeHint, x: usize) -> SizeHint { + let (low, hi) = sh; + let (dlow, dhi) = (low / x, hi.map(|hi| hi / x)); + let (rlow, rhi) = (low % x, hi.map(|hi| hi % x)); + + let low = if rlow > 0 { dlow + 1 } else { dlow }; + let hi = dhi + .and_then(|dhi| rhi.map(|rhi| (dhi, rhi))) + .map(|(dhi, rhi)| if rhi > 0 { dhi + 1 } else { dhi }); + (low, hi) +} + +/// Correct division by `x` with a `SizeHint`. +#[inline] +#[track_caller] +pub fn div_scalar((low, hi): SizeHint, x: usize) -> SizeHint { + (low.saturating_div(x), hi.map(|hi| hi.saturating_div(x))) +} + /// Return the maximum #[inline] pub fn max(a: SizeHint, b: SizeHint) -> SizeHint { @@ -92,3 +114,19 @@ fn mul_size_hints() { assert_eq!(mul((3, Some(4)), (usize::MAX, None)), (usize::MAX, None)); assert_eq!(mul((3, None), (0, Some(0))), (0, Some(0))); } + +#[test] +fn div_ceil_size_scalar() { + assert_eq!(div_ceil_scalar((4, Some(4)), 2), (2, Some(2))); + assert_eq!(div_ceil_scalar((3, Some(3)), 2), (2, Some(2))); + assert_eq!(div_ceil_scalar((4, Some(4)), usize::MAX), (1, Some(1))); + assert_eq!(div_ceil_scalar((4, None), 2), (2, None)); +} + +#[test] +fn div_size_scalar() { + assert_eq!(div_scalar((4, Some(4)), 2), (2, Some(2))); + assert_eq!(div_scalar((3, Some(3)), 2), (1, Some(1))); + assert_eq!(div_scalar((4, Some(4)), usize::MAX), (0, Some(0))); + assert_eq!(div_scalar((4, None), 2), (2, None)); +} diff --git a/tests/quick.rs b/tests/quick.rs index 0af73778a..b1924a42d 100644 --- a/tests/quick.rs +++ b/tests/quick.rs @@ -1116,6 +1116,24 @@ quickcheck! { } } +quickcheck! { + fn chunks_exact_size(a: Vec, size: u8) -> bool { + let mut size = size; + if size == 0 { + size += 1; + } + exact_size(a.into_iter().chunks(size as usize).into_iter()) + } + + fn chunks_correct_size_hint(a: Vec, size: u8) -> bool { + let mut size = size; + if size == 0 { + size += 1; + } + correct_size_hint(a.into_iter().chunks(size as usize).into_iter()) + } +} + // tuple iterators quickcheck! { fn equal_circular_tuple_windows_1(a: Vec) -> bool {