CODEWITHSUNDEEP
rust
Jonathan Woollett-light
Jonathan Woollett-light

Reputation: 3261

How to circumvent `take_while` skipping values?

In trying to chain std::iter::Iterator::take_while calls together I'm losing the last values of each call.

Is there a way to chain calls together like this without skipping values?

Code Playground link:

use std::fmt;

#[derive(Clone)]
struct Point {
    value: u8,
    xe: u8,
    xs: u8,
    y: u8,
}

impl fmt::Debug for Point {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}", self.value)
    }
}

fn main() {
    // All values following 5s within its x distance, partitioned by whether it is above or below.

    // Sorted by xs (x start) (xe = x end)
    #[rustfmt::skip]
    let vec:Vec<Point> = vec![
        Point { value: 4, xe: 1, xs: 1, y: 2 }, // 4
        Point { value: 3, xe: 3, xs: 2, y: 3 }, // 3
        Point { value: 5, xe: 7, xs: 4, y: 6 }, // ---- 5 -----
        Point { value: 3, xe: 5, xs: 5, y: 4 }, // 3
        Point { value: 6, xe: 6, xs: 6, y: 8 }, // 6
        Point { value: 2, xe: 8, xs: 8, y: 3 }, // 2
        Point { value: 8, xe: 10, xs: 9, y: 2 }, // 8
        Point { value: 5, xe: 15, xs: 10, y: 7 }, // ---- 5 -----
        Point { value: 2, xe: 12, xs: 11, y: 10 }, // 2
        Point { value: 7, xe: 13, xs: 13, y: 9 }, // 7
        Point { value: 4, xe: 14, xs: 14, y: 2 } // 4
    ];

    let mut iter = vec.iter();
    loop {
        let c: Vec<_> = iter
            .by_ref()
            .take_while(|x| x.value != 5)
            .cloned()
            .collect();

        println!("c: {:.?}", c);

        if let Some(var) = iter.next() {
            println!("var: {:.?}", var);

            let (a, b): (Vec<_>, Vec<_>) = iter
                .by_ref()
                .take_while(|x| x.xe < var.xe)
                .partition(|x| x.y > var.y);

            println!("a: {:.?}", a);
            println!("b: {:.?}", b);
        } else {
            break;
        }
    }
}

Output:

c: [4, 3]
var: 3
a: []
b: []
c: [2, 8]
var: 2
a: []
b: []
c: [4]

It should output:

c: [4, 3]
var: 5 
a: [3]
b: [6] 
c: [2, 8]
var: 5
a: [2, 7]
b: [4]

Using take_while with std::iter::Iterator::partition seemed a good way to make the code for this relatively clean.

In context the c, a and b values would be passed to functions whose results would be appended to a return value.

Upvotes: 2

Views: 590

Answers (1)

Stargateur
Stargateur

Reputation: 26767

Using next_if() and from_fn():

use std::iter::from_fn;

// ...

let mut iter = vec.iter().peekable();

// ...

let c: Vec<_> = from_fn(|| iter.next_if(|x| x.value != 5))
                    .cloned()
                    .collect();

// ...

let (a, b): (Vec<_>, Vec<_>) = from_fn(|| iter.next_if(|x| x.xe < var.xe))
                                   .partition(|x| x.y > var.y);

Using peeking_take_while() (better) or take_while_ref() from itertools, just replace the function.

Upvotes: 1

Related Questions