kdTreeKNearestQueries.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 "kdTreeQuery.h"
#include "../../query.h"
#include "../Iterator/kdTreeKNearestIterator.h"
 
namespace Ponca
{

    template <typename Traits, template <typename, typename, int> typename IteratorType, typename QueryType>
    class KdTreeKNearestQueryBase : public KdTreeQuery<Traits>, public QueryType
    {
    public:
        using DataPoint      = typename Traits::DataPoint;
        using IndexType      = typename Traits::IndexType;
        using Scalar         = typename DataPoint::Scalar;
        using VectorType     = typename DataPoint::VectorType;
        using QueryAccelType = KdTreeQuery<Traits>;
        using Iterator = IteratorType<typename Traits::IndexType, typename Traits::DataPoint, Traits::MAX_KNN_SIZE>;
        using Self     = KdTreeKNearestQueryBase<Traits, IteratorType, QueryType>;
 
        PONCA_MULTIARCH inline KdTreeKNearestQueryBase(const StaticKdTreeBase<Traits>* kdtree, IndexType k,
                                                       typename QueryType::InputType input)
            : KdTreeQuery<Traits>(kdtree), QueryType(k, input)
        {
        }

        PONCA_MULTIARCH inline Self& operator()(typename QueryType::InputType input, IndexType k)
        {
            return QueryType::template operator()<Self>(input, k);
        }

        PONCA_MULTIARCH inline Self& operator()(typename QueryType::InputType input)
        {
            return QueryType::template operator()<Self>(input);
        }

        PONCA_MULTIARCH inline Iterator begin()
        {
            QueryAccelType::reset();
            QueryType::reset();
            this->search();
            return Iterator(QueryType::m_queue.begin());
        }

        PONCA_MULTIARCH inline Iterator end() { return Iterator(QueryType::m_queue.end()); }
 
    protected:
        PONCA_MULTIARCH inline void search()
        {
            KdTreeQuery<Traits>::searchInternal(
                QueryType::template getInputPosition<VectorType>(QueryAccelType::m_kdtree->points()),
                [](IndexType, IndexType) {}, [this]() { return QueryType::descentDistanceThreshold(); },
                [this](IndexType idx) { return QueryType::skipIndexFunctor(idx); },
                [this](IndexType idx, IndexType, Scalar d) {
                    QueryType::m_queue.push({idx, d});
                    return false;
                });
        }
    };

    template <typename Traits>
    using KdTreeKNearestIndexQuery = KdTreeKNearestQueryBase<
        Traits, KdTreeKNearestIterator,
        KNearestIndexQuery<typename Traits::IndexType, typename Traits::DataPoint::Scalar, Traits::MAX_KNN_SIZE>>;
    template <typename Traits>
    using KdTreeKNearestPointQuery = KdTreeKNearestQueryBase<
        Traits, KdTreeKNearestIterator,
        KNearestPointQuery<typename Traits::IndexType, typename Traits::DataPoint, Traits::MAX_KNN_SIZE>>;
} // namespace Ponca