multi_index.hpp

Go to the documentation of this file.

// SPDX-FileCopyrightText: 2025 VTT Technical Research Centre of Finland Ltd
// SPDX-License-Identifier: AGPL-3.0-or-later
 
#ifndef PFC_MULTI_INDEX_HPP
#define PFC_MULTI_INDEX_HPP
 
#include <array>
#include <iostream>
#include <vector>
 
namespace pfc {
 
template <size_t D> class MultiIndex {
private:
  const std::array<int, D> m_begin; 
  const std::array<int, D> m_size;  
  const std::array<int, D> m_end; 
  const size_t m_linear_begin;    
  const size_t
      m_linear_end; 
  const size_t m_linear_size; 
 
  std::array<int, D> calculate_end(std::array<int, D> begin,
                                   std::array<int, D> size) {
    std::array<int, D> end;
    for (size_t i = 0; i < D; i++) end[i] = begin[i] + size[i] - 1;
    return end;
  }
 
public:
  MultiIndex(std::array<int, D> size,
             std::array<int, D> offset = std::array<int, D>())
      : m_begin(offset), m_size(size), m_end(calculate_end(m_begin, m_size)),
        m_linear_begin(to_linear(m_begin)), m_linear_end(to_linear(m_end)),
        m_linear_size(m_linear_end - m_linear_begin + 1) {}
 
  const std::array<int, D> &get_begin() const { return m_begin; }
  const std::array<int, D> &get_size() const { return m_size; }
  const std::array<int, D> &get_end() const { return m_end; }
  const size_t &get_linear_begin() const { return m_linear_begin; }
  const size_t &get_linear_size() const { return m_linear_size; }
  const size_t &get_linear_end() const { return m_linear_end; }
 
  size_t to_linear(const std::array<int, D> &indices) const {
    size_t linear_index = indices[0] - m_begin[0];
    size_t multiplier = 1;
    for (size_t i = 1; i < D; ++i) {
      multiplier *= m_size[i - 1];
      linear_index += (indices[i] - m_begin[i]) * multiplier;
    }
    return linear_index;
  }
 
  std::array<int, D> to_multi(size_t idx) const {
    std::array<int, D> indices;
    for (size_t i = 0; i < D; ++i) {
      indices[i] = (idx % m_size[i]) + m_begin[i];
      idx /= m_size[i];
    }
    return indices;
  }
 
  bool inbounds(const std::array<int, D> &indices) const {
    for (size_t i = 0; i < D; ++i) {
      if (indices[i] < m_begin[i] || indices[i] > m_end[i]) {
        return false;
      }
    }
    return true;
  }
 
  friend std::ostream &operator<<(std::ostream &os, const MultiIndex<D> &index) {
    os << "MultiIndex set with " << D << " dimensions (indices from {";
    for (size_t i = 0; i < D - 1; i++) os << index.m_begin[i] << ",";
    os << index.m_begin[D - 1] << "} to {";
    for (size_t i = 0; i < D - 1; i++) os << index.m_end[i] << ",";
    os << index.m_end[D - 1] << "}, size {";
    for (size_t i = 0; i < D - 1; i++) os << index.m_size[i] << ",";
    os << index.m_size[D - 1] << "}, linear size " << index.m_linear_size << ")";
    return os;
  }
 
  class Iterator {
  private:
    std::array<int, D> m_indices;    
    const MultiIndex &m_multi_index; 
    size_t m_linear_index;           
 
  public:
    Iterator(std::array<int, D> indices, const MultiIndex &multi_index)
        : m_indices(indices), m_multi_index(multi_index) {
      m_linear_index = m_multi_index.to_linear(indices);
    }
 
    Iterator(int linear_index, const MultiIndex &multi_index)
        : m_multi_index(multi_index), m_linear_index(linear_index) {
      m_indices = m_multi_index.to_multi(linear_index);
    }
 
    std::array<int, D> &operator*() { return m_indices; }
 
    Iterator &operator++() {
      m_linear_index++;
      for (size_t i = 0; i < D; i++) {
        m_indices[i]++;
        if (m_indices[i] > m_multi_index.m_end[i]) {
          m_indices[i] = m_multi_index.m_begin[i];
        } else {
          break;
        }
      }
      return *this;
    }
 
    Iterator operator++(int) {
      Iterator copy(*this);
      ++(*this);
      return copy;
    }
 
    Iterator operator-(int n) const {
      return Iterator(m_linear_index - n, m_multi_index);
    }
 
    bool operator!=(const Iterator &other) const {
      return m_linear_index != other.m_linear_index;
    }
 
    bool operator==(const Iterator &other) const {
      return m_linear_index == other.m_linear_index;
    }
 
    operator size_t() const { return m_linear_index; }
 
    operator std::array<int, D>() const { return m_indices; }
 
    size_t get_linear_index() const { return m_linear_index; }
 
    friend std::ostream &operator<<(std::ostream &os, const Iterator &it) {
      os << "{";
      for (size_t i = 0; i < D - 1; i++) os << it.m_indices[i] << ", ";
      os << it.m_indices[D - 1] << "}";
      return os;
    }
  };
 
  friend class Iterator;
 
  Iterator begin() { return Iterator(m_linear_begin, *this); }
 
  Iterator end() { return Iterator(m_linear_end + 1, *this); }
 
  Iterator begin() const { return Iterator(m_linear_begin, *this); }
 
  Iterator end() const { return Iterator(m_linear_end + 1, *this); }
 
  Iterator from(std::array<int, D> from) { return Iterator(from, *this); }
 
  Iterator from(size_t from) { return Iterator(from, *this); }
};
 
} // namespace pfc
 
#endif