/// \ingroup base
/// \date December 2025.
///
/// In this file is implemented a concurrent disjoint-set (also called
/// union-find) data structure. It was implemented by Wenzel Jakob and comes
/// from the following repository: https://github.com/wjakob/dset
/// The code was slightly modified. Below can be found the license of the
/// original code.
///
/// \b Related \b publication \n
/// "Wait-free Parallel Algorithms for the Union-Find Problem" \n
/// Richard J. Anderson and Heather Woll \n
/// STOC 1991.

/*
zlib license

Copyright (c) 2015 Wenzel Jakob <wenzel@inf.ethz.ch>

This software is provided 'as-is', without any express or implied
warranty.  In no event will the authors be held liable for any damages
arising from the use of this software.

Permission is granted to anyone to use this software for any purpose,
including commercial applications, and to alter it and redistribute it
freely, subject to the following restrictions:

1. The origin of this software must not be misrepresented; you must not
claim that you wrote the original software. If you use this software
in a product, an acknowledgment in the product documentation would be
appreciated but is not required.
2. Altered source versions must be plainly marked as such, and must not be
misrepresented as being the original software.
3. This notice may not be removed or altered from any source distribution.
*/

#pragma once

#include <atomic>
#include <vector>

/**
 * Lock-free parallel disjoint set data structure (aka UNION-FIND)
 * with path compression and union by rank
 *
 * Supports concurrent find(), same() and unite() calls as described
 * in the paper
 *
 * "Wait-free Parallel Algorithms for the Union-Find Problem"
 * by Richard J. Anderson and Heather Woll
 *
 * In addition, this class supports optimistic locking (try_lock/unlock)
 * of disjoint sets and a combined unite+unlock operation.
 *
 * \author Wenzel Jakob
 */
class DisjointSets {
public:
  DisjointSets(uint32_t size) : mData(size) {
    for(uint32_t i = 0; i < size; ++i)
      mData[i] = (uint32_t)i;
  }

  uint32_t find(uint32_t id) const {
    while(id != parent(id)) {
      uint64_t value = mData[id];
      uint32_t new_parent = parent((uint32_t)value);
      uint64_t new_value = (value & 0xFFFFFFFF00000000ULL) | new_parent;
      /* Try to update parent (may fail, that's ok) */
      if(value != new_value)
        mData[id].compare_exchange_weak(value, new_value);
      id = new_parent;
    }
    return id;
  }

  bool same(uint32_t id1, uint32_t id2) const {
    for(;;) {
      id1 = find(id1);
      id2 = find(id2);
      if(id1 == id2)
        return true;
      if(parent(id1) == id1)
        return false;
    }
  }

  uint32_t unite(uint32_t id1, uint32_t id2) {
    for(;;) {
      id1 = find(id1);
      id2 = find(id2);

      if(id1 == id2)
        return id1;

      uint32_t r1 = rank(id1), r2 = rank(id2);

      if(r1 > r2 || (r1 == r2 && id1 < id2)) {
        std::swap(r1, r2);
        std::swap(id1, id2);
      }

      uint64_t oldEntry = ((uint64_t)r1 << 32) | id1;
      uint64_t newEntry = ((uint64_t)r1 << 32) | id2;

      if(!mData[id1].compare_exchange_strong(oldEntry, newEntry))
        continue;

      if(r1 == r2) {
        oldEntry = ((uint64_t)r2 << 32) | id2;
        newEntry = ((uint64_t)(r2 + 1) << 32) | id2;
        /* Try to update the rank (may fail, that's ok) */
        mData[id2].compare_exchange_weak(oldEntry, newEntry);
      }

      break;
    }
    return id2;
  }

  /**
   * Try to lock the a disjoint union identified by one
   * of its elements (this can occasionally fail when there
   * are concurrent operations). The parameter 'id' will be
   * updated to store the current representative ID of the
   * union
   */
  bool try_lock(uint32_t &id) {
    const uint64_t lock_flag = 1ULL << 63;
    id = find(id);
    uint64_t value = mData[id];
    if((value & lock_flag) || (uint32_t)value != id)
      return false;
// On IA32/x64, a PAUSE instruction is recommended for CAS busy loops
#if defined(__i386__) || defined(__amd64__)
    __asm__ __volatile__("pause\n");
#endif
    return mData[id].compare_exchange_strong(value, value | lock_flag);
  }

  void unlock(uint32_t id) {
    const uint64_t lock_flag = 1ULL << 63;
    mData[id] &= ~lock_flag;
  }

  /**
   * Return the representative index of the set that results from merging
   * locked disjoint sets 'id1' and 'id2'
   */
  uint32_t unite_index_locked(uint32_t id1, uint32_t id2) const {
    uint32_t r1 = rank(id1), r2 = rank(id2);
    return (r1 > r2 || (r1 == r2 && id1 < id2)) ? id1 : id2;
  }

  /**
   * Atomically unite two locked disjoint sets and unlock them. Assumes
   * that here are no other concurrent unite() involving the same sets
   */
  uint32_t unite_unlock(uint32_t id1, uint32_t id2) {
    uint32_t r1 = rank(id1), r2 = rank(id2);

    if(r1 > r2 || (r1 == r2 && id1 < id2)) {
      std::swap(r1, r2);
      std::swap(id1, id2);
    }

    mData[id1] = ((uint64_t)r1 << 32) | id2;
    mData[id2] = ((uint64_t)(r2 + ((r1 == r2) ? 1 : 0)) << 32) | id2;

    return id2;
  }

  uint32_t size() const {
    return (uint32_t)mData.size();
  }

  uint32_t rank(uint32_t id) const {
    return ((uint32_t)(mData[id] >> 32)) & 0x7FFFFFFFu;
  }

  uint32_t parent(uint32_t id) const {
    return (uint32_t)mData[id];
  }

  bool isRoot(uint32_t id) const {
    return id == parent(id);
  }

  mutable std::vector<std::atomic<uint64_t>> mData;
};