index_sequence.hh

#pragma once
 
#include <algorithm>
#include <cassert>
#include <cstddef>
#include <iterator>
#include <utility>
 
#include <clingo/util/small_vector.hh>
 
namespace CppClingo::Util {
 
 
template <class T> class index_sequence {
  public:
    class iterator {
      public:
        friend class index_sequence;
        using iterator_category = std::forward_iterator_tag;
        using value_type = T;
        using difference_type = std::ptrdiff_t;
        using reference = T;
        using pointer = void;
 
        iterator() = default;
 
        auto operator*() const -> reference { return it_->second + static_cast<T>(idx_ - start_); }
 
        auto operator++() -> iterator & {
            ++idx_;
            if (idx_ == it_->first) {
                start_ = it_->first;
                ++it_;
            }
            return *this;
        }
 
        auto operator++(int) -> iterator {
            auto tmp = *this;
            ++(*this);
            return tmp;
        }
 
        auto operator==(const iterator &other) const -> bool { return idx_ == other.idx_; }
 
        auto operator!=(const iterator &other) const -> bool { return !(*this == other); }
 
      private:
        using interval_iterator = small_vector<std::pair<size_t, T>, 1>::const_iterator;
 
        iterator(interval_iterator it, size_t idx, size_t start) : it_{it}, idx_{idx}, start_{start} {}
 
        interval_iterator it_;
        size_t idx_ = 0;
        size_t start_ = 0;
    };
 
    void add(T value) {
        if (values_.empty()) {
            values_.emplace_back(1, value);
        } else {
            auto &[r, s] = values_.back();
            auto l = values_.size() > 1 ? values_[values_.size() - 2].first : size_t{0};
            if (s + static_cast<T>(r - l) == value) {
                ++r;
            } else {
                values_.emplace_back(size() + 1, value);
            }
        }
    }
 
    void assign(T s, T e) {
        last_ = 0;
        values_.clear();
        if (s < e) {
            values_.emplace_back(static_cast<size_t>(e - s), s);
        }
    }
 
    [[nodiscard]] auto operator[](size_t index) const -> T {
        assert(index < size() && last_ < values_.size());
        auto l = last_ > 0 ? values_[last_ - 1].first : size_t{0};
        auto r = values_[last_].first;
        assert(l < r);
        auto ib = index < l ? values_.begin() : values_.begin() + last_;
        auto ie = index < r ? values_.begin() + last_ + 1 : values_.end();
        auto it = std::upper_bound(ib, ie, index, [](auto const &a, auto const &b) { return a < b.first; });
        last_ = static_cast<size_t>(std::distance(values_.begin(), it));
        auto start = last_ > 0 ? values_[last_ - 1].first : size_t{0};
        return it->second + static_cast<T>(index - start);
    }
    [[nodiscard]] auto find(T x) const -> size_t {
        // NOTE: has linear complexity in the worst case:
        // - locality could be improved
        // - reverse mapping could be stored
        size_t l = 0;
        for (auto const &[r, s] : values_) {
            auto e = s + static_cast<T>(r - l);
            if (s <= x && x < e) {
                return l + static_cast<size_t>(x - s);
            }
            l = r;
        }
        return size();
    }
    [[nodiscard]] auto size() const -> size_t { return empty() ? 0 : values_.back().first; }
    [[nodiscard]] auto empty() const -> bool { return values_.empty(); }
    [[nodiscard]] auto begin() const -> iterator { return {values_.cbegin(), 0, 0}; }
    [[nodiscard]] auto end() const -> iterator { return {values_.cend(), size(), size()}; }
 
  private:
    small_vector<std::pair<size_t, T>, 1> values_;
    size_t mutable last_ = 0;
};
 
 
} // namespace CppClingo::Util