Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add .get(range) #891

Merged
merged 10 commits into from
May 14, 2024
116 changes: 116 additions & 0 deletions src/iter_index.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,116 @@
use core::iter::{Skip, Take};
use core::ops::{Range, RangeFrom, RangeFull, RangeInclusive, RangeTo, RangeToInclusive};

#[cfg(doc)]
use crate::Itertools;

mod private_iter_index {
use core::ops;

pub trait Sealed {}

impl Sealed for ops::Range<usize> {}
impl Sealed for ops::RangeInclusive<usize> {}
impl Sealed for ops::RangeTo<usize> {}
impl Sealed for ops::RangeToInclusive<usize> {}
impl Sealed for ops::RangeFrom<usize> {}
impl Sealed for ops::RangeFull {}
}

/// Used by [`Itertools::get`] to know which iterator
/// to turn different ranges into.
pub trait IteratorIndex<I>: private_iter_index::Sealed
where
I: Iterator,
{
/// The type returned for this type of index.
type Output: Iterator<Item = I::Item>;

/// Returns an adapted iterator for the current index.
///
/// Prefer calling [`Itertools::get`] instead
/// of calling this directly.
fn index(self, from: I) -> Self::Output;
}

impl<I> IteratorIndex<I> for Range<usize>
where
I: Iterator,
{
type Output = Skip<Take<I>>;

fn index(self, iter: I) -> Self::Output {
iter.take(self.end).skip(self.start)
}
}

impl<I> IteratorIndex<I> for RangeInclusive<usize>
where
I: Iterator,
{
type Output = Take<Skip<I>>;

fn index(self, iter: I) -> Self::Output {
// end - start + 1 without overflowing if possible
let length = if *self.end() == usize::MAX {
assert_ne!(*self.start(), 0);
self.end() - self.start() + 1
} else {
(self.end() + 1).saturating_sub(*self.start())
};
iter.skip(*self.start()).take(length)
}
}

impl<I> IteratorIndex<I> for RangeTo<usize>
where
I: Iterator,
{
type Output = Take<I>;

fn index(self, iter: I) -> Self::Output {
iter.take(self.end)
}
}

impl<I> IteratorIndex<I> for RangeToInclusive<usize>
where
I: Iterator,
{
type Output = Take<I>;

fn index(self, iter: I) -> Self::Output {
assert_ne!(self.end, usize::MAX);
iter.take(self.end + 1)
}
}

impl<I> IteratorIndex<I> for RangeFrom<usize>
where
I: Iterator,
{
type Output = Skip<I>;

fn index(self, iter: I) -> Self::Output {
iter.skip(self.start)
}
}

impl<I> IteratorIndex<I> for RangeFull
where
I: Iterator,
{
type Output = I;

fn index(self, iter: I) -> Self::Output {
iter
}
}

pub fn get<I, R>(iter: I, index: R) -> R::Output
Philippe-Cholet marked this conversation as resolved.
Show resolved Hide resolved
where
I: IntoIterator,
R: IteratorIndex<I::IntoIter>,
{
index.index(iter.into_iter())
}
52 changes: 52 additions & 0 deletions src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -146,6 +146,7 @@ pub mod structs {

/// Traits helpful for using certain `Itertools` methods in generic contexts.
pub mod traits {
pub use crate::iter_index::IteratorIndex;
pub use crate::tuple_impl::HomogeneousTuple;
}

Expand Down Expand Up @@ -194,6 +195,7 @@ mod groupbylazy;
#[cfg(feature = "use_std")]
mod grouping_map;
mod intersperse;
mod iter_index;
#[cfg(feature = "use_alloc")]
mod k_smallest;
#[cfg(feature = "use_alloc")]
Expand Down Expand Up @@ -504,6 +506,56 @@ pub trait Itertools: Iterator {
intersperse::intersperse_with(self, element)
}

/// Returns an iterator over a subsection of the iterator.
///
/// Works similarly to [`slice::get`](https://doc.rust-lang.org/std/primitive.slice.html#method.get).
///
/// **Panics** for ranges `..=usize::MAX` and `0..=usize::MAX`.
///
/// It's a generalisation of [`Iterator::take`] and [`Iterator::skip`],
/// and uses these under the hood.
/// Therefore, the resulting iterator is:
/// - [`ExactSizeIterator`] if the adapted iterator is [`ExactSizeIterator`].
/// - [`DoubleEndedIterator`] if the adapted iterator is [`DoubleEndedIterator`] and [`ExactSizeIterator`].
///
/// # Unspecified Behavior
/// The result of indexing with an exhausted [`core::ops::RangeInclusive`] is unspecified.
///
/// # Examples
///
/// ```
/// use itertools::Itertools;
///
/// let vec = vec![3, 1, 4, 1, 5];
///
/// let mut range: Vec<_> =
/// vec.iter().get(1..=3).copied().collect();
/// assert_eq!(&range, &[1, 4, 1]);
///
/// // It works with other types of ranges, too
/// range = vec.iter().get(..2).copied().collect();
/// assert_eq!(&range, &[3, 1]);
///
/// range = vec.iter().get(0..1).copied().collect();
/// assert_eq!(&range, &[3]);
///
/// range = vec.iter().get(2..).copied().collect();
/// assert_eq!(&range, &[4, 1, 5]);
///
/// range = vec.iter().get(..=2).copied().collect();
/// assert_eq!(&range, &[3, 1, 4]);
///
/// range = vec.iter().get(..).copied().collect();
/// assert_eq!(range, vec);
/// ```
fn get<R>(self, index: R) -> R::Output
where
Self: Sized,
R: traits::IteratorIndex<Self>,
{
iter_index::get(self, index)
}

/// Create an iterator which iterates over both this and the specified
/// iterator simultaneously, yielding pairs of two optional elements.
///
Expand Down
8 changes: 8 additions & 0 deletions tests/laziness.rs
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,14 @@ must_use_tests! {
intersperse_with {
let _ = Panicking.intersperse_with(|| 0);
}
get {
let _ = Panicking.get(1..4);
let _ = Panicking.get(1..=4);
let _ = Panicking.get(1..);
let _ = Panicking.get(..4);
let _ = Panicking.get(..=4);
let _ = Panicking.get(..);
}
zip_longest {
let _ = Panicking.zip_longest(Panicking);
}
Expand Down
35 changes: 35 additions & 0 deletions tests/test_core.rs
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,41 @@ use crate::it::Itertools;
use core::iter;
use itertools as it;

#[allow(dead_code)]
fn get_esi_then_esi<I: ExactSizeIterator + Clone>(it: I) {
fn is_esi(_: impl ExactSizeIterator) {}
is_esi(it.clone().get(1..4));
is_esi(it.clone().get(1..=4));
is_esi(it.clone().get(1..));
is_esi(it.clone().get(..4));
is_esi(it.clone().get(..=4));
is_esi(it.get(..));
}

#[allow(dead_code)]
fn get_dei_esi_then_dei_esi<I: DoubleEndedIterator + ExactSizeIterator + Clone>(it: I) {
fn is_dei_esi(_: impl DoubleEndedIterator + ExactSizeIterator) {}
is_dei_esi(it.clone().get(1..4));
is_dei_esi(it.clone().get(1..=4));
is_dei_esi(it.clone().get(1..));
is_dei_esi(it.clone().get(..4));
is_dei_esi(it.clone().get(..=4));
is_dei_esi(it.get(..));
}

#[test]
fn get_1_max() {
let mut it = (0..5).get(1..=usize::MAX);
assert_eq!(it.next(), Some(1));
assert_eq!(it.next_back(), Some(4));
}

#[test]
#[should_panic]
fn get_full_range_inclusive() {
let _it = (0..5).get(0..=usize::MAX);
}

#[test]
fn product0() {
let mut prod = iproduct!();
Expand Down