Is there an idiomatic way to partition a Rust iterator into multiple iterators by Enum variant?

I am struggling to find a clean/idiomatic approach to partition an iterator into a fixed number of sub parts based on Enum variants.

Example

enum State {
   A(A),
   B(B),
   C(C),
}

let states = vec![A(...), C(...), ...]; // Some iterable structure containing State

I wish to construct three iterators which process the As, Bs, and Cs separately. Here's an approach I tried that failed:

Attempt

Use filter_map.

let As = states.iter().filter_map(|s| match s {
    State::A(a) => Some(a)
    _ => None
});
let Bs = states.iter().filter_map(|s| match s {
    State::B(b) => Some(b)
    _ => None
})
let Cs = ...

I see three issues with this approach:

Firstly, this will not work if mutable references are required, as taking iter_mut on both would result in the compiler believing states is mutably borrowed multiple times.

Secondly, the compiler cannot check if all variants are handled, which does not feel very idiomatic.

Finally, we consume the iterator three times. If we just had an iterator over the source data (instead of the actual array in my example), this would not be possible.

Answer

Solution

I didn't consider that we could simply pass the pointers around to construct what I wanted without needing to consume the iterator multiple times.

Create three vectors to store the references.

let mut As = Vec::new();
let mut Bs = Vec::new();
let mut Cs = Vec::new();
states.iter_mut().for_each(|(state)| match state {
   State::A(a) => As.push(a),
   State::B(b) => Bs.push(b),
   State::C(c) => Cs.push(c),
})