diff --git a/src/iter_index.rs b/src/iter_index.rs new file mode 100644 index 000000000..f92c818a0 --- /dev/null +++ b/src/iter_index.rs @@ -0,0 +1,110 @@ +use core::ops::{ Range, RangeTo, RangeFrom, RangeFull, RangeInclusive, RangeToInclusive }; +use core::iter::{Skip, Take}; +use crate::Itertools; + +mod private_iter_index { + use core::ops; + + pub trait Sealed {} + + impl Sealed for ops::Range {} + impl Sealed for ops::RangeInclusive {} + impl Sealed for ops::RangeTo {} + impl Sealed for ops::RangeToInclusive {} + impl Sealed for ops::RangeFrom {} + impl Sealed for ops::RangeFull {} +} + +/// Used by the ``range`` function to know which iterator +/// to turn different ranges into. +pub trait IteratorIndex : private_iter_index::Sealed + where I: Iterator +{ + /// The type that [`get`] or [`Itertools::get`] + /// returns when called with this type of index. + type Output: Iterator; + + /// Returns an iterator(or value) in the specified range. + /// + /// Prefer calling [`get`] or [`Itertools::get`] instead + /// of calling this directly. + fn index(self, from: I) -> Self::Output; +} + +impl IteratorIndex for Range + where I: Iterator +{ + type Output = Take>; + + fn index(self, iter: I) -> Self::Output { + iter.skip(self.start) + .take(self.end.saturating_sub(self.start)) + } +} + +impl IteratorIndex for RangeInclusive + where I: Iterator +{ + type Output = Take>; + + fn index(self, iter: I) -> Self::Output { + debug_assert!(!self.is_empty(), "The given `RangeInclusive` is exhausted. The result of indexing with an exhausted `RangeInclusive` is unspecified."); + iter.skip(*self.start()) + .take( + (1 + *self.end()) + .saturating_sub(*self.start()) + ) + } +} + +impl IteratorIndex for RangeTo + where I: Iterator +{ + type Output = Take; + + fn index(self, iter: I) -> Self::Output { + iter.take(self.end) + } +} + +impl IteratorIndex for RangeToInclusive + where I: Iterator +{ + type Output = Take; + + fn index(self, iter: I) -> Self::Output { + iter.take(self.end + 1) + } +} + +impl IteratorIndex for RangeFrom + where I: Iterator +{ + type Output = Skip; + + fn index(self, iter: I) -> Self::Output { + iter.skip(self.start) + } +} + +impl IteratorIndex for RangeFull + where I: Iterator +{ + type Output = I; + + fn index(self, iter: I) -> Self::Output { + iter + } +} + +/// Returns an element of the iterator or an iterator +/// over a subsection of the iterator. +/// +/// See [`Itertools::get`] for more information. +pub fn get(iter: I, index: R) + -> R::Output + where I: IntoIterator, + R: IteratorIndex +{ + index.index(iter.into_iter()) +} diff --git a/src/lib.rs b/src/lib.rs index ff0bf7640..60b2aa451 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -162,6 +162,7 @@ pub mod structs { /// Traits helpful for using certain `Itertools` methods in generic contexts. pub mod traits { pub use crate::tuple_impl::HomogeneousTuple; + pub use crate::iter_index::IteratorIndex; } #[allow(deprecated)] @@ -176,6 +177,7 @@ pub use crate::minmax::MinMaxResult; pub use crate::peeking_take_while::PeekingNext; pub use crate::process_results_impl::process_results; pub use crate::repeatn::repeat_n; +pub use crate::iter_index::get; #[allow(deprecated)] pub use crate::sources::{repeat_call, unfold, iterate}; pub use crate::with_position::Position; @@ -196,6 +198,7 @@ mod combinations; mod combinations_with_replacement; mod exactly_one_err; mod diff; +mod iter_index; mod flatten_ok; mod format; #[cfg(feature = "use_std")] @@ -485,6 +488,52 @@ pub trait Itertools : Iterator { intersperse::intersperse(self, element) } + /// Returns an element at a specific location, or 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). + /// + /// It's a generalisation of [`take`], [`skip`] and [`nth`], and uses these + /// under the hood. + /// + /// # 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(..).copied().collect(); + /// assert_eq!(range, vec); + /// ``` + /// + /// # Unspecified Behavior + /// The result of indexing with an exhausted [`core::ops::RangeInclusive`] is unspecified. + /// + /// [`take`]: https://doc.rust-lang.org/std/iter/trait.Iterator.html#method.take + /// [`skip`]: https://doc.rust-lang.org/std/iter/trait.Iterator.html#method.skip + /// [`nth`]: https://doc.rust-lang.org/std/iter/trait.Iterator.html#method.nth + fn get(self, index: R) -> R::Output + where R: iter_index::IteratorIndex, + Self: Sized + { + iter_index::get(self, index) + } + /// An iterator adaptor to insert a particular value created by a function /// between each element of the adapted iterator. ///