VariableMap.h

#pragma once
 
#include <algorithm>
#include <limits>
#include <optional>
#include <utility>
#include <vector>
 
#include "delpi/symbolic/Variable.h"
#include "delpi/util/concepts.h"
#include "delpi/util/exception.h"
 
namespace delpi {

template <class T>
class VariableMap {
  static_assert(std::is_copy_constructible_v<T>, "T must be copy constructible");
 
 public:
  using Id = Variable::Id;
  using key_type = Variable;
  using mapped_type = T;
  using value_type = std::pair<const Variable, T>;

  VariableMap() : min_id_{std::numeric_limits<Id>::max()} {}
  template <TypedIterable<value_type> Items>
  explicit VariableMap(const Items& items) : VariableMap{} {
    Insert(items);
  }

  [[nodiscard]] Id min_id() const { return min_id_; }
  [[nodiscard]] Id max_id() const { return vars_.size() + min_id_; }
  [[nodiscard]] std::size_t size() const {
    return std::ranges::count_if(vars_, [](const std::optional<mapped_type>& item) { return item.has_value(); });
  }
  [[nodiscard]] std::size_t capacity() const { return vars_.size(); }
  [[nodiscard]] bool empty() const { return vars_.empty(); }
  [[nodiscard]] std::vector<mapped_type> items() const {
    std::vector<mapped_type> items;
    items.reserve(vars_.size());
    for (const std::optional<mapped_type>& item : vars_) {
      if (item.has_value()) items.emplace_back(item.value());
    }
    return items;
  }

  [[nodiscard]] mapped_type At(const Variable var) const {
    if (!Contains(var)) throw DelpiOutOfRangeException("Variable not in the map");
    return vars_.at(var.id() - min_id_).value();
  }

  [[nodiscard]] bool Contains(const Variable& var) const { return Contains(var.id()); }
  bool Insert(const Variable& var, const mapped_type& value) { return Insert(var.id(), value); }
  template <TypedIterable<value_type> I>
  void Insert(const I& items) {
    const auto Comparator = [](const std::pair<Variable, T>& a, const std::pair<Variable, T>& b) {
      return a.first.less(b.first);
    };
    // If the map is empty, initialise it with min_id_ = min(vars.id) and size m = max(vars.id) - min_id_ + 1
    if (min_id_ == std::numeric_limits<Id>::max()) {
      min_id_ = std::ranges::min_element(items, Comparator)->first.id();
      vars_.resize(std::ranges::max_element(items, Comparator)->first.id() - min_id_ + 1, false);
      for (const value_type& item : items) vars_.at(item.first.id() - min_id_) = item.second;
      return;
    }
    // If the map is not empty, resize it to the new range and insert the variables
    // First compute the new min_id and max_id comparing with the current min_id and max_id
    // Then resize the vector to the new range m = new_max - new_min + 1
    const Id new_min = std::min(std::ranges::min_element(items, Comparator)->first.id(), min_id_);
    const Id new_max = std::max(std::ranges::max_element(items, Comparator)->first.id(), vars_.size() + min_id_ - 1);
    const std::size_t max_diff = new_max - (vars_.size() + min_id_ - 1);
    const std::size_t old_size = vars_.size();
    vars_.resize(new_max - new_min + 1);
    // If the lower bound of the new range is less than the current min_id, rotate the vector to the left by max_diff
    if (new_min < min_id_) {
      std::rotate(vars_.begin(), vars_.begin() + static_cast<std::int64_t>(old_size),
                  vars_.end() - static_cast<std::int64_t>(max_diff));
      min_id_ = new_min;
    }
    for (const value_type& item : items) vars_.at(item.first.id() - min_id_) = item.second;
  }

  bool Remove(const Variable& var) { return Remove(var.id()); }
  template <TypedIterable<Variable> V>
  void Remove(const V& vars) {
    for (const Variable& var : vars) Remove(var);
  }

  void Clear() {
    vars_.clear();
    min_id_ = std::numeric_limits<Id>::max();
  }
 
 private:
  [[nodiscard]] bool Contains(const Id id) const {
    return id >= min_id_ && id - min_id_ < vars_.size() && vars_.at(id - min_id_).has_value();
  }

  bool Insert(const Id id, const T& value) {
    if (min_id_ == std::numeric_limits<Id>::max()) {
      min_id_ = id;
      vars_.emplace_back(value);
      return true;
    }
    if (id < min_id_) {
      vars_.resize(vars_.size() + min_id_ - id, std::nullopt);
      std::ranges::rotate(vars_, vars_.begin() + static_cast<std::int64_t>(vars_.size() - (min_id_ - id)));
      min_id_ = id;
      vars_.front() = value;
      return true;
    }
    if (id >= min_id_ + vars_.size()) {
      vars_.resize(id - min_id_ + 1, std::nullopt);
      vars_.back() = value;
      return true;
    }
    vars_.at(id - min_id_) = value;
    return false;
  }

  bool Remove(const Id id) {
    if (!Contains(id)) return false;
    const bool was_present = vars_.at(id - min_id_).has_value();
    vars_.at(id - min_id_) = std::nullopt;
    return was_present;
  }
 
  Id min_id_;                           
  std::vector<std::optional<T>> vars_;  
};
 
template <class T>
std::ostream& operator<<(std::ostream& os, const VariableMap<T>& var_map);
 
extern template class VariableMap<int>;
 
}  // namespace delpi
 
#ifdef DELPI_INCLUDE_FMT
 
#include "delpi/util/logging.h"
 
OSTREAM_FORMATTER(delpi::VariableMap<int>)
 
#endif