How to iterate through a set when members might be removed during iteration?

Question

This simple program is a minimal version of a problem I was having. I have an unordered set of pointers to objects, and while iterating through that set, some of the objects were supposed to be erased from the set. In my larger program, this was causing it to crash. In this smaller one, it's merely removing one element and then ending the loop.

In the example below, we have a game featuring Units. When units are idle, the game will make them perform an action. Idle Units are kept track of in an unordered_set named idle_units.

During the game's update loop, it would iterate through idle_units and have them performing an action. When a Unit's act() function is called, it is no longer idle, and therefore removed from idle_units

#include <vector>
#include <memory>
#include <unordered_set>
#include <unordered_map>
#include <cstdlib>
#include <iostream>

// forward declarations
struct Unit;
void set_not_idle(Unit* u_);

struct Unit {  // a simple object with unique identifier
    Unit() { static int id_ = 0; id = id_++; }
    void act() { set_not_idle(this); }
    int id;
};

// data
std::vector<std::unique_ptr<Unit>> unit_storage;
std::unordered_set<Unit*> unit_ptrs;
std::unordered_map<int, Unit*> unit_from_id;
std::unordered_set<Unit*> idle_units;

Unit* get_unit_ptr(Unit u) { return unit_from_id[u.id]; }

void set_idle(Unit* u_) {
    Unit* u = get_unit_ptr(*u_); // ensure we have pointer to unit in storage
    idle_units.insert(u);
}

void set_not_idle(Unit* u_) {
    Unit* u = get_unit_ptr(*u_); // ensure we have pointer to unit in storage
    idle_units.erase(u);
}

void add_unit(Unit u_) {
    unit_storage.push_back(std::make_unique<Unit>(u_));
    Unit* u_ptr = unit_storage.back().get();
    unit_ptrs.insert(u_ptr);
    unit_from_id[u_ptr->id] = u_ptr; // set map from id to pointer in storage
    set_idle(u_ptr); // units start as idle
}

void print() {
    std::cout << "Units in storage: ";
    for (auto a : unit_ptrs) {
        std::cout << a->id << " ";
    }
    std::cout << "  Idle units: ";
    for (auto it = idle_units.begin(); it != idle_units.end(); ++it) {
        std::cout << (*it)->id << " ";
    }
    std::cout << std::endl;
}

int main() {
    srand(25);
    std::vector<Unit> units;

    // randomly populate our unit_storage with 8 units
    for (int i = 0; i < 50; i++) units.push_back(Unit());
    for (int i = 0; i < 8; ) {
        int idx = rand() % units.size();
        if (!get_unit_ptr(units[idx])) {
            add_unit(units[idx]);
            i++;
        }
    }
    print();
    // get all idle units, and have them perform an action
    for (auto it = idle_units.begin(); it != idle_units.end(); ++it) {
        (*it)->act();
    }
    print();
    return 0;
}

This results in the following output:

Units in storage: 36 2 15 43 18 10 38 11   Idle units: 36 2 15 43 18 10 38 11 
Units in storage: 36 2 15 43 18 10 38 11   Idle units: 2 15 43 18 10 38 11

Whereas it should result in no Units remaining in idle_units. What is the most elegant solution to this?

While trying to solve this, I tried different methods of iteration, including a for (auto it : idle_units) loop, or by moving the iterator increment to the body of the loop, but neither of those solutions solved the problem.

[`erase`](https://en.cppreference.com/w/cpp/container/unordered_set/erase) returns an iterator to the element following the erased element. You have to update `it` inside of your loop after erasing an element. — 0x5453, Nov 15 '22 at 20:51

score 1 · Answer 1 · answered Nov 15 '22 at 21:42

The most elegant solution to this is to use an iterator-based loop, and to increment the iterator within the loop body. This ensures that the iterator is always valid, even if the element it points to is removed from the set.

for (auto it = idle_units.begin(); it != idle_units.end(); ) {
    (*it)->act();
    it = idle_units.erase(it);
}

This code will iterate through the set, calling act() on each element. act() will remove the element from the set, so the iterator will be invalidated. The erase() function returns a valid iterator pointing to the next element in the set, so we can simply assign that iterator back to it and continue the loop.

rturrado · Answer 2 · 2022-11-15T22:46:45.120

The loop you have:

calls void Unit::act() for every object of idle_units, which, in turn,
calls void set_no_idle(Unit*), and, finally,
removes the caller object object via idle_units.erase(u);.

For this particular case, since you are not returning anything neither from set_no_idle nor from act, you could get along with a minimal change: just post increment it within the loop, i.e. turn (*it)->act() into (*it++)->act(). [Demo]

for (auto it = idle_units.begin(); it != idle_units.end();) {
    (*it++)->act();
}

// Outputs:
//
//   Units in storage: 36 2 15 43 18 10 38 11   Idle units: 36 2 15 43 18 10 38 11 
//   Units in storage: 36 2 15 43 18 10 38 11   Idle units:

The post increment operation:

creates a copy of it, increments it, and returns a reference to the copy.
The copy is then used to call void Unit::act(), erasing the current idle_units object, while the loop continues with the next idle_units object.

Something equivalent to:

for (auto it = idle_units.begin(); it != idle_units.end();) {
    auto tmp{ it };
    ++tmp;
    (*it)->act();
    it = tmp;
}

How to iterate through a set when members might be removed during iteration?

2 Answers2