index_set.hpp

// SPDX-FileCopyrightText: 2017 - 2025 The Ginkgo authors
//
// SPDX-License-Identifier: BSD-3-Clause
 
#ifndef GKO_PUBLIC_CORE_BASE_INDEX_SET_HPP_
#define GKO_PUBLIC_CORE_BASE_INDEX_SET_HPP_
 
 
#include <algorithm>
#include <initializer_list>
#include <mutex>
#include <vector>
 
#include <ginkgo/core/base/array.hpp>
#include <ginkgo/core/base/exception_helpers.hpp>
#include <ginkgo/core/base/executor.hpp>
#include <ginkgo/core/base/types.hpp>
#include <ginkgo/core/base/utils.hpp>
 
 
namespace gko {
 
 
template <typename IndexType = int32>
class index_set {
public:
    using index_type = IndexType;
 
    explicit index_set(std::shared_ptr<const Executor> exec) noexcept
        : exec_(std::move(exec)),
          index_space_size_{0},
          num_stored_indices_{0},
          subsets_begin_{array<index_type>(exec_)},
          subsets_end_{array<index_type>(exec_)},
          superset_cumulative_indices_{array<index_type>(exec_)}
    {}
 
    explicit index_set(std::shared_ptr<const gko::Executor> exec,
                       std::initializer_list<IndexType> init_list,
                       const bool is_sorted = false)
        : exec_(std::move(exec)),
          index_space_size_(init_list.size() > 0
                                ? *(std::max_element(std::begin(init_list),
                                                     std::end(init_list))) +
                                      1
                                : 0),
          num_stored_indices_{static_cast<IndexType>(init_list.size())}
    {
        GKO_ASSERT(index_space_size_ > 0);
        this->populate_subsets(
            array<IndexType>(this->get_executor(), init_list), is_sorted);
    }
 
    explicit index_set(std::shared_ptr<const gko::Executor> exec,
                       const index_type size,
                       const gko::array<index_type>& indices,
                       const bool is_sorted = false)
        : exec_(std::move(exec)), index_space_size_(size)
    {
        GKO_ASSERT(index_space_size_ >= indices.get_size());
        this->populate_subsets(indices, is_sorted);
    }
 
    index_set(std::shared_ptr<const Executor> exec, const index_set& other)
        : index_set(exec)
    {
        *this = other;
    }
 
    index_set(const index_set& other) : index_set(other.get_executor(), other)
    {}
 
    index_set(std::shared_ptr<const Executor> exec, index_set&& other)
        : index_set(exec)
    {
        *this = std::move(other);
    }
 
    index_set(index_set&& other)
        : index_set(other.get_executor(), std::move(other))
    {}
 
    index_set& operator=(const index_set& other)
    {
        if (&other == this) {
            return *this;
        }
        this->index_space_size_ = other.index_space_size_;
        this->num_stored_indices_ = other.num_stored_indices_;
        this->subsets_begin_ = other.subsets_begin_;
        this->subsets_end_ = other.subsets_end_;
        this->superset_cumulative_indices_ = other.superset_cumulative_indices_;
 
        return *this;
    }
 
    index_set& operator=(index_set&& other)
    {
        if (&other == this) {
            return *this;
        }
        this->index_space_size_ = std::exchange(other.index_space_size_, 0);
        this->num_stored_indices_ = std::exchange(other.num_stored_indices_, 0);
        this->subsets_begin_ = std::move(other.subsets_begin_);
        this->subsets_end_ = std::move(other.subsets_end_);
        this->superset_cumulative_indices_ =
            std::move(other.superset_cumulative_indices_);
 
        return *this;
    }
 
    void clear() noexcept
    {
        this->index_space_size_ = 0;
        this->num_stored_indices_ = 0;
        this->subsets_begin_.clear();
        this->subsets_end_.clear();
        this->superset_cumulative_indices_.clear();
    }
 
    std::shared_ptr<const Executor> get_executor() const { return this->exec_; }
 
    index_type get_size() const { return this->index_space_size_; }
 
    bool is_contiguous() const { return (this->get_num_subsets() <= 1); }
 
    index_type get_num_elems() const { return this->num_stored_indices_; };
 
    index_type get_global_index(index_type local_index) const;
 
    index_type get_local_index(index_type global_index) const;
 
    array<index_type> map_local_to_global(
        const array<index_type>& local_indices,
        const bool is_sorted = false) const;
 
    array<index_type> map_global_to_local(
        const array<index_type>& global_indices,
        const bool is_sorted = false) const;
 
    array<index_type> to_global_indices() const;
 
    array<bool> contains(const array<index_type>& global_indices,
                         const bool is_sorted = false) const;
 
    bool contains(const index_type global_index) const;
 
    index_type get_num_subsets() const
    {
        return this->subsets_begin_.get_size();
    }
 
    const index_type* get_subsets_begin() const
    {
        return this->subsets_begin_.get_const_data();
    }
 
    const index_type* get_subsets_end() const
    {
        return this->subsets_end_.get_const_data();
    }
 
    const index_type* get_superset_indices() const
    {
        return this->superset_cumulative_indices_.get_const_data();
    }
 
private:
    void populate_subsets(const gko::array<index_type>& indices,
                          const bool is_sorted);
 
    std::shared_ptr<const Executor> exec_;
    index_type index_space_size_;
    index_type num_stored_indices_;
    gko::array<index_type> subsets_begin_;
    gko::array<index_type> subsets_end_;
    gko::array<index_type> superset_cumulative_indices_;
};
 
 
}  // namespace gko
 
 
#endif  // GKO_PUBLIC_CORE_BASE_INDEX_SET_HPP_