PartitionedVector.h

Go to the documentation of this file.

/* Copyright (C) 2023-2025 FIAS Frankfurt Institute for Advanced Studies, Frankfurt / Main
   SPDX-License-Identifier: GPL-3.0-only
   Authors: Felix Weiglhofer [committer] */
#ifndef CBM_ALGO_BASE_PARTITIONED_VECTOR_H
#define CBM_ALGO_BASE_PARTITIONED_VECTOR_H
 
#include "PODAllocator.h"
 
#include <boost/serialization/access.hpp>
#include <boost/serialization/vector.hpp>
 
#include <gsl/span>
#include <vector>
 
namespace cbm
{
  template<typename T>
  class PartitionedSpan;

  template<typename T, class Allocator = std::allocator<T>>
  class PartitionedVector {
 
   public:
    using Container_t = std::vector<T, Allocator>;  //< Underlying container type

    PartitionedVector() : fData(), fOffsets({0}), fAddresses() { EnsureDimensions(); }

    PartitionedVector(Container_t&& data, gsl::span<const size_t> sizes, gsl::span<const uint32_t> addresses)
      : fData(std::move(data))
      , fOffsets()
      , fAddresses(addresses.begin(), addresses.end())
    {
      ComputeOffsets(sizes);
      EnsureDimensions();
    }

    template<typename OtherAllocator>
    PartitionedVector(const PartitionedVector<T, OtherAllocator>& other)
      : fData(other.Data().begin(), other.Data().end())
      , fOffsets(other.Offsets())
      , fAddresses(other.Addresses())
    {
      // TODO: this check is overkill? We already know that the dimensions are correct,
      // since they were already checked in the other vector
      EnsureDimensions();
    }

    PartitionedVector(const PartitionedVector<T, Allocator>& other)
      : fData(other.Data().begin(), other.Data().end())
      , fOffsets(other.Offsets())
      , fAddresses(other.Addresses())
    {
      EnsureDimensions();
    }

    PartitionedVector(PartitionedVector<T, Allocator>&& other)
      : fData(std::move(other.fData))
      , fOffsets(std::move(other.fOffsets))
      , fAddresses(std::move(other.fAddresses))
    {
      EnsureDimensions();
      other.fOffsets = {0};
    }
 
    template<typename U>
    PartitionedVector(PartitionedSpan<U> other)
      : fData(other.Data().begin(), other.Data().end())
      , fOffsets(other.Offsets().begin(), other.Offsets().end())
      , fAddresses(other.Addresses().begin(), other.Addresses().end())
    {
      EnsureDimensions();
    }

    PartitionedVector& operator=(const PartitionedVector<T, Allocator>& other)
    {
      if (this != &other) {
        fData      = other.fData;
        fOffsets   = other.fOffsets;
        fAddresses = other.fAddresses;
      }
      return *this;
    }

    PartitionedVector& operator=(PartitionedVector<T, Allocator>&& other)
    {
      if (this != &other) {
        fData          = std::move(other.fData);
        fOffsets       = std::move(other.fOffsets);
        fAddresses     = std::move(other.fAddresses);
        other.fOffsets = {0};
      }
      return *this;
    }

    gsl::span<T> operator[](size_t i)
    {
      EnsureBounds(i);
      return UnsafePartitionSpan(i);
    }

    gsl::span<const T> operator[](size_t i) const
    {
      EnsureBounds(i);
      return UnsafePartitionSpan(i);
    }

    uint32_t Address(size_t i) const
    {
      EnsureBounds(i);
      return fAddresses[i];
    }

    void Clear()
    {
      fData.clear();
      fOffsets.clear();
      fOffsets = {0};
      fAddresses.clear();
    }

    std::pair<gsl::span<T>, uint32_t> Partition(size_t i)
    {
      EnsureBounds(i);
      return std::pair<gsl::span<T>, uint32_t>(UnsafePartitionSpan(i), fAddresses[i]);
    }

    std::pair<gsl::span<const T>, uint32_t> Partition(size_t i) const
    {
      EnsureBounds(i);
      return std::pair<gsl::span<const T>, uint32_t>(UnsafePartitionSpan(i), fAddresses[i]);
    }

    size_t NPartitions() const { return fAddresses.size(); }

    size_t Size(size_t i) const
    {
      EnsureBounds(i);
      return UnsafeSize(i);
    }

    size_t NElements() const { return fData.size(); }

    size_t SizeBytes() const { return fData.size() * sizeof(T); }

    gsl::span<T> Data() { return fData; }

    gsl::span<const T> Data() const { return fData; }

    const std::vector<uint32_t>& Addresses() const { return fAddresses; }

    const std::vector<size_t>& Offsets() const { return fOffsets; }
 
 
   private:
    Container_t fData;                 
    std::vector<size_t> fOffsets;      
    std::vector<uint32_t> fAddresses;  
 
    void EnsureDimensions() const
    {
      if (fOffsets.size() - 1 != fAddresses.size()) {
        throw std::runtime_error("PartitionedVector: fOffsets.size() - 1 != fAddresses.size()");
      }
      if (fOffsets.front() != 0) {
        throw std::runtime_error("PartitionedVector: fOffsets.front() != 0");
      }
      if (fOffsets.back() != fData.size()) {
        throw std::runtime_error("PartitionedVector: fOffsets.back() != fData.size()");
      }
    }
 
    void EnsureBounds(size_t i) const
    {
      if (i >= fAddresses.size()) throw std::out_of_range("PartitionedVector: index out of bounds");
    }
 
    void ComputeOffsets(gsl::span<const size_t> sizes)
    {
      fOffsets.reserve(sizes.size() + 1);
      fOffsets.push_back(0);
      for (auto n : sizes) {
        fOffsets.push_back(fOffsets.back() + n);
      }
    }
 
    size_t UnsafeSize(size_t i) const { return fOffsets[i + 1] - fOffsets[i]; }
 
    gsl::span<T> UnsafePartitionSpan(size_t i) { return gsl::span<T>(fData.data() + fOffsets[i], UnsafeSize(i)); }
 
    gsl::span<const T> UnsafePartitionSpan(size_t i) const
    {
      return gsl::span<const T>(fData.data() + fOffsets[i], UnsafeSize(i));
    }
 
   private:  // serialization
    friend class boost::serialization::access;
 
    template<class Archive>
    void serialize(Archive& ar, unsigned int /*version*/)
    {
      ar& fData;
      ar& fOffsets;
      ar& fAddresses;
    }
  };
 
  template<typename T>
  using PartitionedPODVector = PartitionedVector<T, PODAllocator<T>>;
 
}  // namespace cbm
 
#endif