Mutating a Vec while iterating through it and keeping order

code review
1

Hello, I found myself needing to mutate a Vec while also iterating through it, and keeping the order of it after I'm finished. I found a solution that is currently working, but I would appreciate any feedback on the code if there is a more idiomatic safe way to approach this.

To showcase the solution I came up with I've abstracted the design away from my project into a simplified example. I have four enum variants that help the control flow of this example:

enum Choice {
    Purify,  // Needs to be changed.
    Changed, // The result of being changed.
    Remove,
    Keep,
}

I then create a Vec of these enums inside the variable x:

let mut x: Vec<Choice> = vec![
    Choice::Purify,
    Choice::Remove,
    Choice::Keep,
    Choice::Keep,
    Choice::Remove,
];

The goal is to iterate through x and change Purify into Changed, and remove all Remove variants while maintaining the original order of the vector. In this very simplified example there is only one remove type, but in my project that is determined per-iteration.

As x is defined above the goal would be:

[
    Changed,
    Keep,
    Keep,
]

I can achieve this list by what I call "plucking." I create a while loop to iterate through the index of the vector and remove the first index element, since it is removed from the vector I can then mutate the vector using retain() on it to remove or keep elements. Afterwards I add the plucked element back to the end of the vector by using push(). Note, when removing an element from the vector I subtract it from the while loop's end. This gives it a "treadmill" like effect because all of the elements are removed and re-added in the same order they were in originally.

let mut i = 0;
let mut end = x.len();

while i < end {
    // Pluck the first element always
    // as this moves like a treadmill.
    let mut plucked = x.remove(0);
    if matches!(plucked, Choice::Purify) { plucked = Choice::Changed; }
    
    x.retain(|choice| {
        // Remove anything that matches our logic.
        if matches!(choice, Choice::Remove) {
            end = end - 1; // Since removed, decrease end.
            false
        } else {
            true
        }
    });
    
    // Add the plucked back to the end.
    x.push(plucked);
    
    // Increase the index for looping.
    i += 1;
}

You can view the results of this here on the Rust Playground.

Is there a more idiomatic safe way to handle this? Or in general a way to improve this code? Thank you! :slight_smile:

2

Just use retain_mut

#[derive(Debug)]
 enum Choice {
     Purify,  // Needs to be changed
     Changed, // Was changed
     Remove,
     Keep,
 }
 
 fn main() {
     let mut x: Vec<Choice> = vec![
         Choice::Purify,
         Choice::Remove,
         Choice::Keep,
         Choice::Keep,
         Choice::Remove,
     ];
 
     x.retain_mut(|choice| {
         match choice {
             Choice::Remove => return false,
             Choice::Purify => *choice = Choice::Changed,
             _ => (),
         }
         true
     });
 
     dbg!(&x);
 }
3

Thank you for your response!

Yes I've already attempted this and it doesn't fit what I'm needing to do in my project. I didn't do so well in explaining through my example but I'm not just removing elements from my project's vector like I am with x in the example; Per-iteration in the while loop for the element I've plucked I'm using the retain() on the vector to remove certain elements that fit the plucked element (custom logic).

I may need to post an overall better example but consider this version.

enum Fruits {
    Apple(Health),
    Orange(Health),
    Banana(Health),
    Basket(u8),
}

enum Health {
    Fresh,
    Spoiled,
}

Say that I have a vector of fruits in a certain order that contain fresh and spoiled. I need to target the spoiled ones and turn them into a basket, and then iterate through the vector and remove the fresh ones that match the same spoiled type that I turned into a basket. Every fruit I remove from the original vector I increase the basket count. I then add the basket back to the end of the list.

Here is an original vector of fruits:

let mut x: Vec<Fruits> = vec![
    Fruits::Apple(Health::Spoiled),
    Fruits::Apple(Health::Fresh),
    Fruits::Orange(Health::Spoiled),
    Fruits::Orange(Health::Fresh),
    Fruits::Orange(Health::Fresh),
    Fruits::Banana(Health::Spoiled),
    Fruits::Banana(Health::Fresh),
    Fruits::Banana(Health::Fresh),
    Fruits::Banana(Health::Fresh),
];

This is the outcome I expect to have (again crude example, I understand that the basket doesn't explain which fruit it contains but this should convey the idea):

&x = [
    Fruits::Basket(1),
    Fruits::Basket(2),
    Fruits::Basket(3),
]
4

I think your algorithm has quadratic complexity, because of the vec.remove(0) call each iteration. I would prefer a two-pass algorithm, first iterate through the vector and marks all elements to be removed, then iterate again to remove them.

this does mean I need a separate vec to store the marked flags/indexes, so it trades memory for performance.

5

Perhaps:

// I'm guessing at the exact logic.
//
// But I'm assuming everything is sorted:
// - All fruits adjacent
// - All spoiled for a given fruit before fresh
// - Baskets last
fn consolidate(x: &mut Vec<Fruits>) {
    let mut saw_spoiled = [false; 3];
    let mut count = [0; 3];
    x.retain_mut(|fruit| {
        let (idx, spoiled) = match fruit {
            Fruits::Apple(h) => (0, *h == Health::Spoiled),
            Fruits::Orange(h) => (1, *h == Health::Spoiled),
            Fruits::Banana(h) => (2, *h == Health::Spoiled),
            Fruits::Basket(_) => return true,
        };
        if spoiled {
            saw_spoiled[idx] = true;
            false
        } else if saw_spoiled[idx] {
            count[idx] += 1;
            false
        } else {
            true
        }
    });

    for (spoiled, count) in saw_spoiled.into_iter().zip(count) {
        if spoiled {
            x.push(Fruits::Basket(count));
        }
    }
}

You could do something similar without the arrays in two passes:

  • One using iter_mut to turn the first apple (etc) into a basket which is updated in place, turning anything fresh into spoiled
  • Then use retain_mut to yeet the spoiled

But tracking the state ("am I the same fruit?") is a bit more involved. Example.

If you used a VecDeque, you could pop from the front so long as the front was fruit, keeping track of which is the current fruit. It would be similar to updatng in place in some ways.

  • Front is same fruit? Pop and update current basket.
  • Front is new fruit? Push current basket to the back and start another.
  • Front is a basket? Push current basket (if any) to the back and you're done.
6

This looks like input and output should contain different types.

Either way, the easiest way to do this is to write to a second Vec rather than rearranging things in place in one Vec.

7

although I don't understand why you'd do something like this, if this is some form of tally or statistic, I prefer to create a new Vec for the result instead.

nevertheless, looking at this input/output example, I think this operation is a variant (in that the elements to keep is modified) of dedup. here's an example for this concept (in pseudo code):

let mut x: Vec<Fruits> = vec![
    Fruits::Apple(Health::Spoiled),
    Fruits::Apple(Health::Fresh),
    Fruits::Orange(Health::Spoiled),
    Fruits::Orange(Health::Fresh),
    Fruits::Orange(Health::Fresh),
    Fruits::Banana(Health::Spoiled),
    Fruits::Banana(Health::Fresh),
    Fruits::Banana(Health::Fresh),
    Fruits::Banana(Health::Fresh),
];

// keep track of recent spoiled
let mut last_spoiled = None;

// note the order of the elements to the `same_bucket` predicate callback:
// `a` comes *after* `b` in the original `Vec`, in other words, it is `a`, not `b`,
// that will be removed if the predicate returns `true`.
// `b` is the last one that gets "dedup"-ed, i.e. we returned `false` for it
// in a previous iteration, or it is the very first element of the `Vec`
x.dedup_by(|a, b| {
    // if b is spoiled, save it for comparison, and change it to a basket
    if b.is_spoiled() {
        last_spoiled = Some(mem::replace(b, Fruits::Basket(1)));
    }
    // now check if `a` is fresh and matches the last spoiled type
    if a.is_fresh() &&  last_spoiled.as_ref().is_some_and(|last| same_type(a, last)) {
        // `b` is the last "kept" element and must be a basket at this point
        // since `last_spoiled` is `Some`
        // increase the basket count number and return `true` to remove `a`
        let &mut Fruit::Basket(ref mut n) = b else { unreachable!() };
        n += 1;
        true
    } else {
        // otherwise, return `false` to keep `a` in the `Vec`, and it becomes `b` in
        // next iteration and be checked again, and changed to a basket if spoiled
        // check it here won't help for a Vec with single element. see the caveat below
        false
    }
});

// big CAVEAT for edge cases:
// because `dedup_by` operates on pairs of elements, we must check
// the *last* element after the dedup, like we did for `b` in the callback.
// alternatively, we can check the *first* element *before* the dedup.
// both will also cover the case when the Vec contains a single element.
x.last_mut().map(|fruit| if fruit.is_spoiled() { *fruit = Fruits::Basket(1) });

// check result
assert_eq!(&x, &[Fruits::Basket(1), Fruit::Basket(2), Fruits::Basket(3)]);
8

The key ingredients for successful list (sequence) processing is to take small steps toward your goal. Instead of creating a complex loop, it is much easier to run through the list(s) several times using simple loops.

While from a theoretical view you only need the mighty fold function (you can find a lot of study material in the functional language literature about fold), in practice you have two types of operations in (list) processing:

  1. iteration to gather some new data, summary or information,
  2. modifications, like sorting, removing, and modifying elements etc.

So for your artificial example:

These are just three simple passes/steps:

    // Step #1: gather data
    let mut spoiled_apple = 0;
    let mut spoiled_orange = 0;
    let mut spoiled_banana = 0;

    let mut fresh_apple = 0;
    let mut fresh_orange = 0;
    let mut fresh_banana = 0;

    for fruit in &x {
        match fruit {
            &Fruits::Apple(Health::Spoiled) => spoiled_apple += 1,
            &Fruits::Orange(Health::Spoiled) => spoiled_orange += 1,
            &Fruits::Banana(Health::Spoiled) => spoiled_banana += 1,

            &Fruits::Apple(Health::Fresh) => fresh_apple += 1,
            &Fruits::Orange(Health::Fresh) => fresh_orange += 1,
            &Fruits::Banana(Health::Fresh) => fresh_banana += 1,

            _ => (),
        }
    }

    // Step #2: filter list
    x.retain(|fruit| match fruit {
        &Fruits::Apple(_) if spoiled_apple > 0 => false,
        &Fruits::Orange(_) if spoiled_orange > 0 => false,
        &Fruits::Banana(_) if spoiled_banana > 0 => false,
        _ => true,
    });

    // Step #3: extend list
    x.extend([
        Fruits::Basket(fresh_apple),
        Fruits::Basket(fresh_orange),
        Fruits::Basket(fresh_banana),
    ]);

The key insight in combining simple list scans is: Because every step is O(n) multiple (fixed) such steps are still O(n).

For practical reasons that means, if the list is somewhat small and fits well into the CPU cache, multiple scans are pretty cheap. And if the list is large, multiple scans are often perfectly fine, because a linear scan of RAM is relatively fast compared to random access.