nearest_iterator.hpp

 
#ifndef UFO_CONTAINER_DETAIL_TREE_MAP_NEAREST_ITERATOR_HPP
#define UFO_CONTAINER_DETAIL_TREE_MAP_NEAREST_ITERATOR_HPP
 
// UFO
#include <ufo/container/tree/index.hpp>
#include <ufo/geometry/distance.hpp>
#include <ufo/geometry/dynamic_geometry.hpp>
 
// STL
#include <cassert>
#include <cmath>
#include <cstddef>
#include <iterator>
#include <queue>
#include <type_traits>
 
namespace ufo
{
// Forward declare
template <std::size_t Dim, class T>
class TreeMap;
 
template <bool Const, std::size_t Dim, class T, class Geometry = DynamicGeometry>
class TreeMapNearestIterator
{
 private:
    //
    // Friends
    //
 
    template <bool, std::size_t, class, class>
    friend class TreeMapNearestIterator;
 
    friend class TreeMap<Dim, T>;
 
 private:
    static constexpr std::size_t const BF = TreeMap<Dim, T>::branchingFactor();
 
    using RawIterator =
        std::conditional_t<Const, typename TreeMap<Dim, T>::container_type::const_iterator,
                           typename TreeMap<Dim, T>::container_type::iterator>;
 
    using Point = typename TreeMap<Dim, T>::Point;
 
    struct S {
        float       dist_sq;
        TreeIndex   node;
        RawIterator it;
 
        S(float dist_sq, TreeIndex node, RawIterator it = {}) noexcept
            : dist_sq(dist_sq), node(node), it(it)
        {
        }
 
        bool operator>(S const& rhs) const noexcept { return dist_sq > rhs.dist_sq; }
    };
 
    using Queue = std::priority_queue<S, std::vector<S>, std::greater<S>>;
 
 public:
    //
    // Tags
    //
 
    using iterator_category = std::forward_iterator_tag;
    using difference_type   = std::ptrdiff_t;
    using value_type        = typename std::iterator_traits<RawIterator>::value_type;
    using reference         = typename std::iterator_traits<RawIterator>::reference;
    using pointer           = typename std::iterator_traits<RawIterator>::pointer;
 
    TreeMapNearestIterator() = default;
 
    TreeMapNearestIterator(TreeMapNearestIterator const&) = default;
 
    // From non-const to const or change of geometry type
    template <bool Const2, class Geometry2,
              std::enable_if_t<(Const && !Const2) || (Const == Const2 &&
                                                      !std::is_same_v<Geometry, Geometry2>),
                               bool> = true>
    TreeMapNearestIterator(TreeMapNearestIterator<Const2, Dim, T, Geometry2> const& other)
        : tm_(other.tm_), query_(other.query_), epsilon_sq_(other.epsilon_sq_)
    {
        auto queue = other.queue_;
        while (!queue.empty()) {
            queue_.emplace(queue.top().dist_sq, queue.top().node, queue.top().it);
            queue.pop();
        }
    }
 
    TreeMapNearestIterator& operator++()
    {
        queue_.pop();
        next();
        return *this;
    }
 
    TreeMapNearestIterator operator++(int)
    {
        TreeMapNearestIterator tmp(*this);
        ++*this;
        return tmp;
    }
 
    reference operator*() const { return *queue_.top().it; }
 
    pointer operator->() const { return &*queue_.top().it; }
 
    template <bool Const2, class Geometry2>
    bool operator==(TreeMapNearestIterator<Const2, Dim, T, Geometry2> const& other)
    {
        return queue_.empty() == other.queue_.empty() &&
               (queue_.empty() || queue_.top().it == other.queue_.top().it);
    }
 
    template <bool Const2, class Geometry2>
    bool operator!=(TreeMapNearestIterator<Const2, Dim, T, Geometry2> const& other)
    {
        return !(*this == other);
    }
 
 private:
    [[nodiscard]] bool returnable(TreeIndex node) const
    {
        return tm_->isPureLeaf(node) && !tm_->empty(node);
    }
 
    [[nodiscard]] bool returnable(S const& s) const { return RawIterator{} != s.it; }
 
    [[nodiscard]] bool traversable(TreeIndex node) const { return tm_->isParent(node); }
 
    void next()
    {
        while (!queue_.empty()) {
            auto cur = queue_.top();
            if (returnable(cur)) {
                return;
            }
 
            queue_.pop();
 
            if (returnable(cur.node)) {
                auto& v = tm_->values(cur.node);
                using std::begin;
                using std::end;
                for (auto it = begin(v), last = end(v); it != last; ++it) {
                    Point p       = it->first;
                    float dist_sq = distanceSquared(query_, p);
                    queue_.emplace(dist_sq, cur.node, it);
                }
                continue;
            }
 
            TreeIndex node = tm_->child(cur.node, 0);
            for (; BF > node.offset; ++node.offset) {
                if (!traversable(node) && !returnable(node)) {
                    continue;
                }
 
                float dist_sq = distanceSquared(query_, tm_->bounds(node)) + epsilon_sq_;
                queue_.emplace(dist_sq, node);
            }
        }
    }
 
    [[nodiscard]] RawIterator iterator() { return queue_.top().it; }
 
 private:
    TreeMapNearestIterator(TreeMap<Dim, T>* tm, TreeIndex node, Geometry const& query,
                           float epsilon = 0.0f)
        : tm_(tm), query_(query), epsilon_sq_(epsilon * epsilon)
    {
        if (traversable(node) || returnable(node)) {
            float dist_sq = distanceSquared(query_, tm_->bounds(node));
            queue_.emplace(dist_sq, node);
            next();
        }
    }
 
    // From const to non-const
    template <bool Const2, class Geometry2, std::enable_if_t<!Const && Const2, bool> = true>
    TreeMapNearestIterator(TreeMapNearestIterator<Const2, Dim, T, Geometry2> const& other)
        : tm_(other.tm), query_(other.query_), epsilon_sq_(other.epsilon_sq_)
    {
        auto queue = other.queue_;
        while (!queue.empty()) {
            auto const& cur = queue.top();
            // Remove const from it
            queue_.emplace(cur.dist_sq, cur.node, tm_->values(cur.node).erase(cur.it, cur.it));
            queue.pop();
        }
    }
 
 private:
    TreeMap<Dim, T>* tm_ = nullptr;
 
    Geometry query_;
    float    epsilon_sq_;
 
    Queue queue_;
};
}  // namespace ufo
 
#endif  // UFO_CONTAINER_DETAIL_TREE_MAP_NEAREST_ITERATOR_HPP