cvmmap-streamer/lib/mcap/include/mcap/intervaltree.hpp

// Adapted from <https://github.com/ekg/intervaltree/blob/master/IntervalTree.h>

#pragma once

#include <algorithm>
#include <cassert>
#include <iostream>
#include <limits>
#include <memory>
#include <vector>

namespace mcap::internal {

template <class Scalar, typename Value>
class Interval {
public:
  Scalar start;
  Scalar stop;
  Value value;
  Interval(const Scalar& s, const Scalar& e, const Value& v)
      : start(std::min(s, e))
      , stop(std::max(s, e))
      , value(v) {}
};

template <class Scalar, typename Value>
Value intervalStart(const Interval<Scalar, Value>& i) {
  return i.start;
}

template <class Scalar, typename Value>
Value intervalStop(const Interval<Scalar, Value>& i) {
  return i.stop;
}

template <class Scalar, typename Value>
std::ostream& operator<<(std::ostream& out, const Interval<Scalar, Value>& i) {
  out << "Interval(" << i.start << ", " << i.stop << "): " << i.value;
  return out;
}

template <class Scalar, class Value>
class IntervalTree {
public:
  using interval = Interval<Scalar, Value>;
  using interval_vector = std::vector<interval>;

  struct IntervalStartCmp {
    bool operator()(const interval& a, const interval& b) {
      return a.start < b.start;
    }
  };

  struct IntervalStopCmp {
    bool operator()(const interval& a, const interval& b) {
      return a.stop < b.stop;
    }
  };

  IntervalTree()
      : left(nullptr)
      , right(nullptr)
      , center(Scalar(0)) {}

  ~IntervalTree() = default;

  std::unique_ptr<IntervalTree> clone() const {
    return std::unique_ptr<IntervalTree>(new IntervalTree(*this));
  }

  IntervalTree(const IntervalTree& other)
      : intervals(other.intervals)
      , left(other.left ? other.left->clone() : nullptr)
      , right(other.right ? other.right->clone() : nullptr)
      , center(other.center) {}

  IntervalTree& operator=(IntervalTree&&) = default;
  IntervalTree(IntervalTree&&) = default;

  IntervalTree& operator=(const IntervalTree& other) {
    center = other.center;
    intervals = other.intervals;
    left = other.left ? other.left->clone() : nullptr;
    right = other.right ? other.right->clone() : nullptr;
    return *this;
  }

  IntervalTree(interval_vector&& ivals, std::size_t depth = 16, std::size_t minbucket = 64,
               std::size_t maxbucket = 512, Scalar leftextent = 0, Scalar rightextent = 0)
      : left(nullptr)
      , right(nullptr) {
    --depth;
    const auto minmaxStop = std::minmax_element(ivals.begin(), ivals.end(), IntervalStopCmp());
    const auto minmaxStart = std::minmax_element(ivals.begin(), ivals.end(), IntervalStartCmp());
    if (!ivals.empty()) {
      center = (minmaxStart.first->start + minmaxStop.second->stop) / 2;
    }
    if (leftextent == 0 && rightextent == 0) {
      // sort intervals by start
      std::sort(ivals.begin(), ivals.end(), IntervalStartCmp());
    } else {
      assert(std::is_sorted(ivals.begin(), ivals.end(), IntervalStartCmp()));
    }
    if (depth == 0 || (ivals.size() < minbucket && ivals.size() < maxbucket)) {
      std::sort(ivals.begin(), ivals.end(), IntervalStartCmp());
      intervals = std::move(ivals);
      assert(is_valid().first);
      return;
    } else {
      Scalar leftp = 0;
      Scalar rightp = 0;

      if (leftextent || rightextent) {
        leftp = leftextent;
        rightp = rightextent;
      } else {
        leftp = ivals.front().start;
        rightp = std::max_element(ivals.begin(), ivals.end(), IntervalStopCmp())->stop;
      }

      interval_vector lefts;
      interval_vector rights;

      for (typename interval_vector::const_iterator i = ivals.begin(); i != ivals.end(); ++i) {
        const interval& cur = *i;
        if (cur.stop < center) {
          lefts.push_back(cur);
        } else if (cur.start > center) {
          rights.push_back(cur);
        } else {
          assert(cur.start <= center);
          assert(center <= cur.stop);
          intervals.push_back(cur);
        }
      }

      if (!lefts.empty()) {
        left.reset(new IntervalTree(std::move(lefts), depth, minbucket, maxbucket, leftp, center));
      }
      if (!rights.empty()) {
        right.reset(
          new IntervalTree(std::move(rights), depth, minbucket, maxbucket, center, rightp));
      }
    }
    assert(is_valid().first);
  }

  // Call f on all intervals near the range [start, stop]:
  template <class UnaryFunction>
  void visit_near(const Scalar& start, const Scalar& stop, UnaryFunction f) const {
    if (!intervals.empty() && !(stop < intervals.front().start)) {
      for (auto& i : intervals) {
        f(i);
      }
    }
    if (left && start <= center) {
      left->visit_near(start, stop, f);
    }
    if (right && stop >= center) {
      right->visit_near(start, stop, f);
    }
  }

  // Call f on all intervals crossing pos
  template <class UnaryFunction>
  void visit_overlapping(const Scalar& pos, UnaryFunction f) const {
    visit_overlapping(pos, pos, f);
  }

  // Call f on all intervals overlapping [start, stop]
  template <class UnaryFunction>
  void visit_overlapping(const Scalar& start, const Scalar& stop, UnaryFunction f) const {
    auto filterF = [&](const interval& cur) {
      if (cur.stop >= start && cur.start <= stop) {
        // Only apply f if overlapping
        f(cur);
      }
    };
    visit_near(start, stop, filterF);
  }

  // Call f on all intervals contained within [start, stop]
  template <class UnaryFunction>
  void visit_contained(const Scalar& start, const Scalar& stop, UnaryFunction f) const {
    auto filterF = [&](const interval& cur) {
      if (start <= cur.start && cur.stop <= stop) {
        f(cur);
      }
    };
    visit_near(start, stop, filterF);
  }

  interval_vector find_overlapping(const Scalar& start, const Scalar& stop) const {
    interval_vector result;
    visit_overlapping(start, stop, [&](const interval& cur) {
      result.emplace_back(cur);
    });
    return result;
  }

  interval_vector find_contained(const Scalar& start, const Scalar& stop) const {
    interval_vector result;
    visit_contained(start, stop, [&](const interval& cur) {
      result.push_back(cur);
    });
    return result;
  }

  bool empty() const {
    if (left && !left->empty()) {
      return false;
    }
    if (!intervals.empty()) {
      return false;
    }
    if (right && !right->empty()) {
      return false;
    }
    return true;
  }

  template <class UnaryFunction>
  void visit_all(UnaryFunction f) const {
    if (left) {
      left->visit_all(f);
    }
    std::for_each(intervals.begin(), intervals.end(), f);
    if (right) {
      right->visit_all(f);
    }
  }

  std::pair<Scalar, Scalar> extent() const {
    struct Extent {
      std::pair<Scalar, Scalar> x{std::numeric_limits<Scalar>::max(),
                                  std::numeric_limits<Scalar>::min()};
      void operator()(const interval& cur) {
        x.first = std::min(x.first, cur.start);
        x.second = std::max(x.second, cur.stop);
      }
    };
    Extent extent;

    visit_all([&](const interval& cur) {
      extent(cur);
    });
    return extent.x;
  }

  // Check all constraints.
  // If first is false, second is invalid.
  std::pair<bool, std::pair<Scalar, Scalar>> is_valid() const {
    const auto minmaxStop =
      std::minmax_element(intervals.begin(), intervals.end(), IntervalStopCmp());
    const auto minmaxStart =
      std::minmax_element(intervals.begin(), intervals.end(), IntervalStartCmp());

    std::pair<bool, std::pair<Scalar, Scalar>> result = {
      true, {std::numeric_limits<Scalar>::max(), std::numeric_limits<Scalar>::min()}};
    if (!intervals.empty()) {
      result.second.first = std::min(result.second.first, minmaxStart.first->start);
      result.second.second = std::min(result.second.second, minmaxStop.second->stop);
    }
    if (left) {
      auto valid = left->is_valid();
      result.first &= valid.first;
      result.second.first = std::min(result.second.first, valid.second.first);
      result.second.second = std::min(result.second.second, valid.second.second);
      if (!result.first) {
        return result;
      }
      if (valid.second.second >= center) {
        result.first = false;
        return result;
      }
    }
    if (right) {
      auto valid = right->is_valid();
      result.first &= valid.first;
      result.second.first = std::min(result.second.first, valid.second.first);
      result.second.second = std::min(result.second.second, valid.second.second);
      if (!result.first) {
        return result;
      }
      if (valid.second.first <= center) {
        result.first = false;
        return result;
      }
    }
    if (!std::is_sorted(intervals.begin(), intervals.end(), IntervalStartCmp())) {
      result.first = false;
    }
    return result;
  }

private:
  interval_vector intervals;
  std::unique_ptr<IntervalTree> left;
  std::unique_ptr<IntervalTree> right;
  Scalar center;
};

}  // namespace mcap::internal