BLI_atomic_disjoint_set.hh

Go to the documentation of this file.

/* SPDX-FileCopyrightText: 2023 Blender Authors
 *
 * SPDX-License-Identifier: GPL-2.0-or-later */

 
#pragma once
 
#include <atomic>
 
#include "BLI_array.hh"
 
namespace blender {

class AtomicDisjointSet {
 private:
  /* Can generally used relaxed memory order with this algorithm. */
  static constexpr auto relaxed = std::memory_order_relaxed;
 
  struct Item {
    int parent;
    int rank;
  };

  mutable Array<std::atomic<Item>> items_;
 
 public:
  AtomicDisjointSet(const int size);

  void join(int x, int y)
  {
    while (true) {
      x = this->find_root(x);
      y = this->find_root(y);
 
      if (x == y) {
        /* They are in the same set already. */
        return;
      }
 
      int x_rank = items_[x].load(relaxed).rank;
      int y_rank = items_[y].load(relaxed).rank;
 
      if (
          /* Implement union by rank heuristic. */
          x_rank > y_rank
          /* If the rank is the same, make a consistent decision. */
          || (x_rank == y_rank && x < y))
      {
        std::swap(x_rank, y_rank);
        std::swap(x, y);
      }
 
      /* Update parent of item x. */
      Item x_item_old = {x, x_rank};
      const Item x_item_new{y, x_rank};
      if (!items_[x].compare_exchange_strong(x_item_old, x_item_new, relaxed)) {
        /* Another thread has updated item x, start again. */
        continue;
      }
 
      if (x_rank == y_rank) {
        /* Increase rank of item y. This may fail when another thread has updated item y in the
         * meantime. That may lead to worse behavior with the union by rank heurist, but seems to
         * be ok in practice. */
        Item y_item_old{y, y_rank};
        const Item y_item_new{y, y_rank + 1};
        items_[y].compare_exchange_weak(y_item_old, y_item_new, relaxed);
      }
 
      return;
    }
  }

  bool in_same_set(int x, int y) const
  {
    while (true) {
      x = this->find_root(x);
      y = this->find_root(y);
      if (x == y) {
        return true;
      }
      if (items_[x].load(relaxed).parent == x) {
        return false;
      }
    }
  }

  int find_root(int x) const
  {
    while (true) {
      const Item item = items_[x].load(relaxed);
      if (x == item.parent) {
        return x;
      }
      const int new_parent = items_[item.parent].load(relaxed).parent;
      if (item.parent != new_parent) {
        /* This halves the path for faster future lookups. That fail but that does not change
         * correctness. */
        Item expected = item;
        const Item desired{new_parent, item.rank};
        items_[x].compare_exchange_weak(expected, desired, relaxed);
      }
      x = new_parent;
    }
  }

  bool is_root(const int x) const
  {
    const Item item = items_[x].load(relaxed);
    return item.parent == x;
  }

  int calc_reduced_ids(MutableSpan<int> result) const;

  int count_sets() const;
};
 
}  // namespace blender