I am trying to write an iterator which conditionally uses elements in a separate iterator. In my example, the separate iterator should increment the sum variable. Once another condition is met *n == 4, the iterator should stop checking the condition and assume rest of elements are increments for the sum variable. I have the following working example:
fn conditional(n: &i64) -> bool {
// a lot of code here which is omitted for brevity
n % 2 == 0
}
fn main() {
let buf = vec![1,2,3,4,5,6];
let mut sum = 0;
let mut iter = buf.iter();
while let Some(n) = iter.next() {
if conditional(n) {
sum += n;
}
if *n == 4 {
// end of file - assume rest of elements are `conditional`
break;
}
};
// rest of elements [5,6]
for n in iter {
sum += n;
}
println!("sum (2+4+5+6): {:?}", sum);
}
output:
sum (2+4+5+6): 17
I would rather write the same thing with a single iterator using something like flat_map:
fn conditional(n: &i64) -> bool {
// a lot of code here which is omitted for brevity
n % 2 == 0
}
fn main() {
let buf = vec![1,2,3,4,5,6];
let mut sum = 0;
let mut terminate = false;
buf.iter().flat_map(|n| {
if *n == 4 {
// hard terminate here - return Some(n) for rest of iterator [5,6]
terminate = true;
return Some(n);
}
if terminate {
return Some(n);
}
if conditional(n) {
return Some(n);
}
None // odd
})
.for_each(|n| {
sum += n;
});
println!("sum (2+4+5+6): {:?}", sum);
}
output:
sum (2+4+5+6): 17
Is there a way to write this in a more concise manner? I want to short-circuit the iterator once the *n == 4 condition is reached.
>Solution :
There are many ways to solve this.
Here are a couple:
fn conditional(n: &i64) -> bool {
// a lot of code here which is omitted for brevity
n % 2 == 0
}
fn main() {
let buf = vec![1, 2, 3, 4, 5, 6];
let sum = buf
.iter()
.fold((0, false), |(mut sum, mut terminate), value| {
if *value == 4 {
terminate = true;
}
if terminate || conditional(value) {
sum += *value;
}
(sum, terminate)
})
.0;
println!("sum (2+4+5+6): {:?}", sum);
}
sum (2+4+5+6): 17
Or using filter and a stateful closure:
fn conditional(n: &i64) -> bool {
// a lot of code here which is omitted for brevity
n % 2 == 0
}
fn main() {
let buf = vec![1, 2, 3, 4, 5, 6];
let sum: i64 = buf
.iter()
.filter({
let mut terminate = false;
move |&value| {
terminate || {
if *value == 4 {
terminate = true;
}
conditional(value)
}
}
})
.sum();
println!("sum (2+4+5+6): {:?}", sum);
}
sum (2+4+5+6): 17