Consume the remaining elements of a Zip

When zipping 2 iterators of different length, I would like to consume the remaining elements of the longest iterator.

let a = [1, 2, 3];
let b = [1, 2];

let mut sum = 0;

let mut zip = a.iter().zip(b.iter());
while let Some(tuple) = zip.next() {
   let (value_a, value_b) = tuple;
   sum += value_a + value_b;
}

sum += zip.rest().sum();

The Zip struct has a and b containing the 2 iterators but those fields are private.

Itertools::zip_longest handles that case. Alternatively you could do something like this.

EDIT: The code below looks right at first glance, but is wrong and will drop one value in the case of uneven iterators (see @alice's comment).

let a = [1, 2, 3];
let b = [1, 2];

let mut sum = 0;

let mut a_iter = a.iter();
let mut b_iter = b.iter();
let zip = (&mut a_iter).zip(&mut b_iter);
for (value_a, value_b) in zip {
   sum += value_a + value_b;
}
for value in a_iter {
    sum += value;
}
for value in b_iter {
    sum += value;
}
2 Likes

If you print the sum after running the code you posted, you will understand why zip doesn't let you get back the inner iterators. Doing so will lose the first extra message in the longer iterator, and to avoid it, you need some sort of peeking construct to not consume the message when checking if we are at the end.

It will print 6, not 9.

2 Likes

Good catch. I guess what I wrote is only useful if you're fine dropping one value, which probably isn't what OP wants.

This code gets the right solution, but you probably shouldn't use it :smile:. I would just use Itertools::zip_longest. Its cleaner.

let a = [1, 2, 3];
let b = [1, 2];

let mut sum = 0;

struct RewindableIter<I,T>{
    iter: I,
    last: Option<T>,
    rewound: bool
};

impl <I,T> RewindableIter<I,T> where I: Iterator<Item=T>, T:Copy {
    fn new(iter: I) -> Self {
        Self {
            iter,
            last: None,
            rewound: false
        }
    }
    
    fn rewind_one(&mut self) {
        if self.rewound {
            panic!("can only call rewind_one once")
        }
        self.rewound = true;
    }
}


impl <I,T> Iterator for RewindableIter<I,T> where I: Iterator<Item=T>, T:Copy {
    type Item = T;
    fn next(&mut self) -> Option<T> {
        if self.rewound {
            self.rewound = false;
            return self.last
        }
        if let Some(v) = self.iter.next() {
            self.last = Some(v);
            Some(v)
        } else {
            self.last = None;
            None
        }
    }
}

trait Rewindable<T> where Self: Iterator<Item=T> + Sized, T:Copy {
    fn rewindable(self) -> RewindableIter<Self,T> {
        RewindableIter::new(self)
    }
}

impl <I,T> Rewindable<T> for I where I: Iterator<Item=T> + Sized, T: Copy {}

let mut a_iter = a.iter().rewindable();
let mut b_iter = b.iter().rewindable();
let zip = (&mut a_iter).zip(&mut b_iter);
for (value_a, value_b) in zip {
   sum += value_a + value_b;
}
a_iter.rewind_one();
b_iter.rewind_one();
for value in a_iter {
    sum += value;
}
for value in b_iter {
    sum += value;
}
println!("{}",sum)

playground

Thank you @drewkett, zip_longest is exactly what I needed.