CODEWITHSUNDEEP
rust
v6sm3z9s
v6sm3z9s

Reputation: 51

Is there a convenient way to iterate over a range of elements of an iterator?

I'm search for a function like this:

fn get_range<I: Iterator<Item = Item>, Item>(iterator: I, range: impl RangeBounds<usize>) -> impl Iterator<Item = Item> {
    todo!()
}

let vec = vec!['a', 'b', 'c', 'd', 'e'];

for c in get_range(vec.iter(), 1..=3) {
    println!("{c}");
} // returns 'b', 'c', and 'd' separated by newlines.

I have the feeling something like this should exist somewhere, but I haven't found anything.

If necessary I would implement it with a combination of Iterator::skip() and Iterator::take(), but I would prefer not having to do that.

Upvotes: 3

Views: 1060

Answers (2)

isaactfa
isaactfa

Reputation: 6651

I put together a solution that I think covers most if not all edge cases:

use std::ops::{Bound, RangeBounds};

pub fn get_range<I: Iterator>(
    iterator: I,
    r: impl RangeBounds<usize>,
) -> impl Iterator<Item = I::Item> {
    use Bound::*;
    let start = r.start_bound();
    let end = r.end_bound();
    let start = match start {
        Unbounded => 0,
        Included(s) => *s,
        Excluded(s) => s.saturating_add(1),
    };
    let end = match end {
        Unbounded => usize::MAX,
        Included(&e) => e,
        Excluded(e) => e.saturating_sub(1),
    };
    let take = (end.checked_sub(start).map(|t| t.saturating_add(1))).unwrap_or(0);
    iterator.skip(start).take(take)
}

You can check out the suite of tests at the playground.

Upvotes: 3

v6sm3z9s
v6sm3z9s

Reputation: 51

get_range could be implemented like:

fn get_range<I: Iterator<Item = Item>, Item>(
    iterator: I,
    range: impl RangeBounds<usize> + Debug,
) -> impl Iterator<Item = Item> {
    let start_bound = match range.start_bound() {
        std::ops::Bound::Included(&num) => num,
        std::ops::Bound::Excluded(&num) => num + 1,
        std::ops::Bound::Unbounded => 0,
    };

    let mut end_bound = match range.end_bound() {
        std::ops::Bound::Included(&num) => Some(num + 1),
        std::ops::Bound::Excluded(&num) => Some(num),
        std::ops::Bound::Unbounded => None,
    };

    iterator
        .take_while(move |_| {
            if let Some(num) = &mut end_bound {
                if *num == 0 {
                    false
                } else {
                    *num -= 1;
                    true
                }
            } else {
                true
            }
        })
        .skip(start_bound)
}

Here is a playground link for this.

Upvotes: 2

Related Questions