knnGraphRangeQuery.h

/*
 This Source Code Form is subject to the terms of the Mozilla Public
 License, v. 2.0. If a copy of the MPL was not distributed with this
 file, You can obtain one at http://mozilla.org/MPL/2.0/.
*/
 
#pragma once
 
#include "../../query.h"
#include "../Iterator/knnGraphRangeIterator.h"
 
namespace Ponca
{
    template <typename Traits>
    class StaticKnnGraphBase;
 
    template <typename Traits>
    class KnnGraphRangeQuery : public RangeIndexQuery<typename Traits::IndexType, typename Traits::DataPoint::Scalar>
    {
    protected:
        using QueryType = RangeIndexQuery<typename Traits::IndexType, typename Traits::DataPoint::Scalar>;
        friend class KnnGraphRangeIterator<Traits>; // This type must be equal to KnnGraphRangeQuery::Iterator
 
    public:
        using DataPoint  = typename Traits::DataPoint;
        using IndexType  = typename Traits::IndexType;
        using Scalar     = typename DataPoint::Scalar;
        using VectorType = typename DataPoint::VectorType;
        using Iterator   = KnnGraphRangeIterator<Traits>;
        using Self       = KnnGraphRangeQuery<Traits>;
 
    public:
        PONCA_MULTIARCH inline KnnGraphRangeQuery(const StaticKnnGraphBase<Traits>* graph, Scalar radius, int index)
            : QueryType(radius, index), m_graph(graph), m_flag(), m_stack()
        {
        }
 
        PONCA_MULTIARCH inline Self& operator()(int index, Scalar radius)
        {
            return QueryType::template operator()<Self>(index, radius);
        }
 
        PONCA_MULTIARCH inline Self& operator()(int index) { return QueryType::template operator()<Self>(index); }
 
        PONCA_MULTIARCH inline Iterator begin()
        {
            QueryType::reset();
            Iterator it(this);
            this->initialize(it);
            this->advance(it);
            return it;
        }
 
        PONCA_MULTIARCH inline Iterator end() { return Iterator(this, static_cast<int>(m_graph->size())); }
 
    protected:
        PONCA_MULTIARCH inline void initialize(Iterator& iterator)
        {
            m_flag.clear();
            m_flag.insert(QueryType::input());
 
            PONCA_DEBUG_ASSERT(m_stack.empty());
            m_stack.push(QueryType::input());
 
            iterator.m_index = -1;
        }
 
        PONCA_MULTIARCH inline void advance(Iterator& iterator)
        {
            const auto& points = m_graph->points();
            const auto& point  = points[QueryType::input()].pos();
 
            if (iterator == end())
                return;
 
            // Search the next neighbor in range
            while (!m_stack.empty())
            {
                int idx_current = m_stack.top();
                m_stack.pop();
 
                PONCA_DEBUG_ASSERT((point - points[idx_current].pos()).squaredNorm() < QueryType::squaredRadius());
 
                iterator.m_index = idx_current;
 
                for (int idx_nei : m_graph->kNearestNeighbors(idx_current))
                {
                    PONCA_DEBUG_ASSERT(idx_nei >= 0);
                    // Add into the search stack only if :
                    // A - The point is within range
                    Scalar d  = (point - points[idx_nei].pos()).squaredNorm();
                    Scalar th = QueryType::descentDistanceThreshold();
                    if (d >= th)
                        continue;
 
                    // B - The point is not already visited
                    auto [iteratorToInsertedValue, wasInserted] = m_flag.insert(idx_nei);
                    if (!wasInserted)
                    {
                        // Check that the Set isn't full (in case if it's a limited capacity set)
                        PONCA_ASSERT(iteratorToInsertedValue != m_flag.end());
                        continue;
                    }
 
                    m_stack.push(idx_nei);
                }
 
                // Query is not included in returned set
                if (iterator.m_index != QueryType::input())
                    return; // Next neighbor was found
            }
            iterator = end(); // If no neighbor was found, we've reached the end
        }
 
    protected:
        const StaticKnnGraphBase<Traits>* m_graph{nullptr};
        typename Traits::KnnGraphRangeSet m_flag;    
        typename Traits::KnnGraphRangeStack m_stack; 
    };
 
} // namespace Ponca