SpatialSearch.h

Go to the documentation of this file.

 
#ifndef PBAT_GEOMETRY_SPATIALSEARCH_H
#define PBAT_GEOMETRY_SPATIALSEARCH_H
 
#include "pbat/Aliases.h"
#include "pbat/HostDevice.h"
#include "pbat/common/ConstexprFor.h"
#include "pbat/common/Heap.h"
#include "pbat/common/NAryTreeTraversal.h"
#include "pbat/common/Queue.h"
#include "pbat/common/Stack.h"
 
#include <algorithm>
#include <array>
#include <limits>
#include <queue>
#include <type_traits>
 
namespace pbat::geometry {
template <
    class FChild,
    class FIsLeaf,
    class FLeafSize,
    class FLeafObject,
    class FNodeOverlap,
    class FObjectOverlap,
    class FOnFound,
    auto N           = 2,
    class TIndex     = Index,
    auto kStackDepth = 64>
PBAT_HOST_DEVICE bool Overlaps(
    FChild fChild,
    FIsLeaf fIsLeaf,
    FLeafSize fLeafSize,
    FLeafObject fLeafObject,
    FNodeOverlap fNodeOverlaps,
    FObjectOverlap fObjectOverlaps,
    FOnFound fOnFound,
    TIndex root = 0);

template <
    class FChildLhs,
    class FIsLeafLhs,
    class FLeafSizeLhs,
    class FLeafObjectLhs,
    class FChildRhs,
    class FIsLeafRhs,
    class FLeafSizeRhs,
    class FLeafObjectRhs,
    class FNodesOverlap,
    class FObjectsOverlap,
    class FOnFound,
    auto NLhs        = 2,
    auto NRhs        = 2,
    class TIndex     = Index,
    auto kStackDepth = 128>
PBAT_HOST_DEVICE bool Overlaps(
    FChildLhs fChildLhs,
    FIsLeafLhs fIsLeafLhs,
    FLeafSizeLhs fLeafSizeLhs,
    FLeafObjectLhs fLeafObjectLhs,
    FChildRhs fChildRhs,
    FIsLeafRhs fIsLeafRhs,
    FLeafSizeRhs fLeafSizeRhs,
    FLeafObjectRhs fLeafObjectRhs,
    FNodesOverlap fNodesOverlap,
    FObjectsOverlap fObjectsOverlap,
    FOnFound fOnFound,
    TIndex rootLhs = 0,
    TIndex rootRhs = 0);

template <
    class FChild,
    class FIsLeaf,
    class FLeafSize,
    class FLeafObject,
    class FNodesOverlap,
    class FObjectsOverlap,
    class FOnFound,
    auto N           = 2,
    class TIndex     = Index,
    auto kStackDepth = 128>
PBAT_HOST_DEVICE bool SelfOverlaps(
    FChild fChild,
    FIsLeaf fIsLeaf,
    FLeafSize fLeafSize,
    FLeafObject fLeafObject,
    FNodesOverlap fNodesOverlap,
    FObjectsOverlap fObjectsOverlap,
    FOnFound fOnFound,
    TIndex root = 0);

template <
    class FChild,
    class FIsLeaf,
    class FLeafSize,
    class FLeafObject,
    class FDistanceLowerBound,
    class FDistance,
    class FOnFound,
    auto N           = 2,
    class TScalar    = Scalar,
    class TIndex     = Index,
    auto kStackDepth = 64,
    auto kQueueSize  = 8>
PBAT_HOST_DEVICE bool DfsAllNearestNeighbours(
    FChild fChild,
    FIsLeaf fIsLeaf,
    FLeafSize fLeafSize,
    FLeafObject fLeafObject,
    FDistanceLowerBound fLower,
    FDistance fDistance,
    FOnFound fOnFound,
    bool bUseBestFirstSearch = false,
    TScalar fUpper           = std::numeric_limits<TScalar>::max(),
    TScalar eps              = TScalar(0),
    TIndex root              = 0);

template <
    class FChild,
    class FIsLeaf,
    class FLeafSize,
    class FLeafObject,
    class FDistanceLowerBound,
    class FDistance,
    class FOnFound,
    auto N             = 2,
    class TScalar      = Scalar,
    class TIndex       = Index,
    auto kHeapCapacity = N * 1024>
PBAT_HOST_DEVICE bool NearestToFurthestNeighbours(
    FChild fChild,
    FIsLeaf fIsLeaf,
    FLeafSize fLeafSize,
    FLeafObject fLeafObject,
    FDistanceLowerBound fLower,
    FDistance fDistance,
    FOnFound fOnFound,
    TScalar fUpper = std::numeric_limits<TScalar>::max(),
    TIndex root    = 0);

template <
    class FChild,
    class FIsLeaf,
    class FLeafSize,
    class FLeafObject,
    class FDistanceLowerBound,
    class FDistance,
    class FOnFound,
    auto N             = 2,
    class TScalar      = Scalar,
    class TIndex       = Index,
    auto kHeapCapacity = N * 1024>
PBAT_HOST_DEVICE bool AllNearestNeighbours(
    FChild fChild,
    FIsLeaf fIsLeaf,
    FLeafSize fLeafSize,
    FLeafObject fLeafObject,
    FDistanceLowerBound fLower,
    FDistance fDistance,
    FOnFound fOnFound,
    TScalar fUpper = std::numeric_limits<TScalar>::max(),
    TScalar eps    = TScalar(0),
    TIndex root    = 0);

template <
    class FChild,
    class FIsLeaf,
    class FLeafSize,
    class FLeafObject,
    class FDistanceLowerBound,
    class FDistance,
    class FOnFound,
    auto N             = 2,
    class TScalar      = Scalar,
    class TIndex       = Index,
    auto kHeapCapacity = N * 1024>
PBAT_HOST_DEVICE bool KNearestNeighbours(
    FChild fChild,
    FIsLeaf fIsLeaf,
    FLeafSize fLeafSize,
    FLeafObject fLeafObject,
    FDistanceLowerBound fLower,
    FDistance fDistance,
    FOnFound fOnFound,
    TIndex K,
    TScalar fUpper = std::numeric_limits<TScalar>::max(),
    TIndex root    = 0);
 
template <
    class FChild,
    class FIsLeaf,
    class FLeafSize,
    class FLeafObject,
    class FNodeOverlap,
    class FObjectOverlap,
    class FOnFound,
    auto N,
    class TIndex,
    auto kStackDepth>
bool Overlaps(
    FChild fChild,
    FIsLeaf fIsLeaf,
    FLeafSize fLeafSize,
    FLeafObject fLeafObject,
    FNodeOverlap fNodeOverlaps,
    FObjectOverlap fObjectOverlaps,
    FOnFound fOnFound,
    TIndex root)
{
    Index k{0};
    auto fVisit = [&] PBAT_HOST_DEVICE(TIndex node) {
        if (not fNodeOverlaps(node))
            return false;
        if (fIsLeaf(node))
        {
            TIndex const nLeafObjects = fLeafSize(node);
            for (TIndex i = 0; i < nLeafObjects; ++i)
            {
                TIndex o = fLeafObject(node, i);
                if (fObjectOverlaps(o))
                    fOnFound(o, k++);
            }
        }
        return true;
    };
    // NOTE:
    // Instead of using a pre-order traversal, we visit children of the currently visited node
    // simultaneously. This is because the branch and bound tree might store each node's
    // children in contiguous memory, or at least with some notion of locality. Thus, the fLower
    // function will benefit from many cache hits.
    // using FVisit = decltype(fVisit);
    // common::TraverseNAryTreePseudoPreOrder<FVisit, FChild, TIndex, N, kStackDepth>(
    //     fVisit,
    //     fChild,
    //     root);
    common::Stack<TIndex, kStackDepth> stack{};
    if (not fVisit(root))
        return true;
    stack.Push(root);
    while (not stack.IsEmpty())
    {
        if (stack.IsFull())
            return false;
        TIndex const node = stack.Pop();
        common::ForRange<0, N>([&]<auto i> PBAT_HOST_DEVICE() {
            TIndex const child = fChild.template operator()<i>(node);
            if (child >= 0)
                if (fVisit(child) and not stack.IsFull())
                    stack.Push(child);
        });
    }
    return true;
}
 
template <
    class FChildLhs,
    class FIsLeafLhs,
    class FLeafSizeLhs,
    class FLeafObjectLhs,
    class FChildRhs,
    class FIsLeafRhs,
    class FLeafSizeRhs,
    class FLeafObjectRhs,
    class FNodesOverlap,
    class FObjectsOverlap,
    class FOnFound,
    auto NLhs,
    auto NRhs,
    class TIndex,
    auto kStackDepth>
PBAT_HOST_DEVICE bool Overlaps(
    FChildLhs fChildLhs,
    FIsLeafLhs fIsLeafLhs,
    FLeafSizeLhs fLeafSizeLhs,
    FLeafObjectLhs fLeafObjectLhs,
    FChildRhs fChildRhs,
    FIsLeafRhs fIsLeafRhs,
    FLeafSizeRhs fLeafSizeRhs,
    FLeafObjectRhs fLeafObjectRhs,
    FNodesOverlap fNodesOverlap,
    FObjectsOverlap fObjectsOverlap,
    FOnFound fOnFound,
    TIndex rootLhs,
    TIndex rootRhs)
{
#include "pbat/warning/Push.h"
#include "pbat/warning/SignConversion.h"
    struct Pair
    {
        TIndex nLhs;
        TIndex nRhs;
        std::int16_t levelLhs;
        std::int16_t levelRhs;
    };
    common::Stack<Pair, kStackDepth> stack{};
    auto const fRecurseRight = [&] PBAT_HOST_DEVICE(Pair const& p) {
        common::ForRange<0, NRhs>([&]<auto i> PBAT_HOST_DEVICE() {
            TIndex const childRhs = fChildRhs.template operator()<i>(p.nRhs);
            if (childRhs >= 0 and not stack.IsFull())
                stack.Push(
                    {p.nLhs, childRhs, p.levelLhs, static_cast<std::int16_t>(p.levelRhs + 1)});
        });
    };
    auto const fRecurseLeft = [&] PBAT_HOST_DEVICE(Pair const& p) {
        common::ForRange<0, NLhs>([&]<auto i> PBAT_HOST_DEVICE() {
            TIndex const childLhs = fChildLhs.template operator()<i>(p.nLhs);
            if (childLhs >= 0 and not stack.IsFull())
                stack.Push(
                    {childLhs, p.nRhs, static_cast<std::int16_t>(p.levelLhs + 1), p.levelRhs});
        });
    };
    // Traverse top-down
    TIndex k{0};
    stack.Push({rootLhs, rootRhs, std::int16_t{0}, std::int16_t{0}});
    while (not stack.IsEmpty())
    {
        if (stack.IsFull())
            return false;
        Pair const p = stack.Pop();
        if (not fNodesOverlap(p.nLhs, p.nRhs))
            continue;
        bool const bIsLhsLeaf = fIsLeafLhs(p.nLhs);
        bool const bIsRhsLeaf = fIsLeafRhs(p.nRhs);
        if (bIsLhsLeaf and bIsRhsLeaf)
        {
            TIndex const nLeafObjectsLhs = fLeafSizeLhs(p.nLhs);
            TIndex const nLeafObjectsRhs = fLeafSizeRhs(p.nRhs);
            for (TIndex i = 0; i < nLeafObjectsLhs; ++i)
            {
                TIndex const oLhs = fLeafObjectLhs(p.nLhs, i);
                for (TIndex j = 0; j < nLeafObjectsRhs; ++j)
                {
                    TIndex const oRhs = fLeafObjectRhs(p.nRhs, j);
                    if (fObjectsOverlap(oLhs, oRhs))
                        fOnFound(oLhs, oRhs, k++);
                }
            }
        }
        else if (bIsLhsLeaf and not bIsRhsLeaf)
        {
            fRecurseRight(p);
        }
        else if (not bIsLhsLeaf and bIsRhsLeaf)
        {
            fRecurseLeft(p);
        }
        else
        {
            if (p.levelLhs > p.levelRhs)
            {
                fRecurseRight(p);
            }
            else
            {
                fRecurseLeft(p);
            }
        }
    }
    return true;
#include "pbat/warning/Pop.h"
};
 
template <
    class FChild,
    class FIsLeaf,
    class FLeafSize,
    class FLeafObject,
    class FNodesOverlap,
    class FObjectsOverlap,
    class FOnFound,
    auto N,
    class TIndex,
    auto kStackDepth>
PBAT_HOST_DEVICE bool SelfOverlaps(
    FChild fChild,
    FIsLeaf fIsLeaf,
    FLeafSize fLeafSize,
    FLeafObject fLeafObject,
    FNodesOverlap fNodesOverlap,
    FObjectsOverlap fObjectsOverlap,
    FOnFound fOnFound,
    TIndex root)
{
#include "pbat/warning/Push.h"
#include "pbat/warning/SignConversion.h"
    struct Pair
    {
        TIndex nLhs;
        TIndex nRhs;
        std::int16_t levelLhs;
        std::int16_t levelRhs;
    };
    common::Stack<Pair, kStackDepth> stack{};
    auto const fRecurseRight = [&] PBAT_HOST_DEVICE(Pair const& p) {
        common::ForRange<0, N>([&]<auto i> PBAT_HOST_DEVICE() {
            TIndex const childRhs = fChild.template operator()<i>(p.nRhs);
            if (childRhs >= 0 and not stack.IsFull())
                stack.Push(
                    {p.nLhs, childRhs, p.levelLhs, static_cast<std::int16_t>(p.levelRhs + 1)});
        });
    };
    auto const fRecurseLeft = [&] PBAT_HOST_DEVICE(Pair const& p) {
        common::ForRange<0, N>([&]<auto i> PBAT_HOST_DEVICE() {
            TIndex const childLhs = fChild.template operator()<i>(p.nLhs);
            if (childLhs >= 0 and not stack.IsFull())
                stack.Push(
                    {childLhs, p.nRhs, static_cast<std::int16_t>(p.levelLhs + 1), p.levelRhs});
        });
    };
    // For any given nodes (n,a) where a is an ancestor of n, it is guaranteed that
    // (n,a) are overlapping, since n is embedded in a. We need to find a way to avoid
    // adding pairs (n,a) to the stack. A pair (ni,nj) should only exist or be tested
    // if neither ni nor nj is an ancestor of the other.
    TIndex k{0};
    stack.Push({root, root, std::int16_t{0}, std::int16_t{0}});
    while (not stack.IsEmpty())
    {
        if (stack.IsFull())
            return false;
        Pair const p = stack.Pop();
        // This is a node visitor, not a pair to check for overlap, add children to stack
        if (p.nLhs == p.nRhs)
        {
            TIndex const n = p.nLhs;
            if (fIsLeaf(n))
            {
                TIndex const nLeafObjects = fLeafSize(n);
                for (TIndex i = 0; i < nLeafObjects; ++i)
                {
                    for (TIndex j = i + 1; j < nLeafObjects; ++j)
                    {
                        TIndex const oLhs = fLeafObject(n, i);
                        TIndex const oRhs = fLeafObject(n, j);
                        if (fObjectsOverlap(oLhs, oRhs))
                            fOnFound(oLhs, oRhs, k++);
                    }
                }
            }
            else
            {
                // Collect children
                std::array<TIndex, N> children;
                common::ForRange<0, N>([&]<auto i> PBAT_HOST_DEVICE() {
                    children[i] = fChild.template operator()<i>(n);
                });
                // Add node visitors first, so that we don't traverse the whole tree in one go
                std::int16_t const nextLevel = p.levelLhs + std::int16_t{1};
                common::ForRange<0, N>([&]<auto i> PBAT_HOST_DEVICE() {
                    if (stack.IsFull())
                        return;
                    if (children[i] < 0)
                        return;
                    stack.Push({children[i], children[i], nextLevel, nextLevel});
                });
                // Add all distinct child pairs to the stack
                common::ForRange<0, N>([&]<auto i> PBAT_HOST_DEVICE() {
                    if (children[i] < 0)
                        return;
                    common::ForRange<i + 1, N>([&]<auto j> PBAT_HOST_DEVICE() {
                        if (stack.IsFull())
                            return;
                        if (children[j] < 0)
                            return;
                        stack.Push({children[i], children[j], nextLevel, nextLevel});
                    });
                });
            }
        }
        else
        {
            if (not fNodesOverlap(p.nLhs, p.nRhs))
                continue;
 
            bool const bIsLeftLeaf  = fIsLeaf(p.nLhs);
            bool const bIsRightLeaf = fIsLeaf(p.nRhs);
            if (bIsLeftLeaf and bIsRightLeaf)
            {
                TIndex const nLeafObjectsLhs = fLeafSize(p.nLhs);
                TIndex const nLeafObjectsRhs = fLeafSize(p.nRhs);
                for (TIndex i = 0; i < nLeafObjectsLhs; ++i)
                {
                    TIndex const oLhs = fLeafObject(p.nLhs, i);
                    for (TIndex j = 0; j < nLeafObjectsRhs; ++j)
                    {
                        TIndex const oRhs = fLeafObject(p.nRhs, j);
                        if (fObjectsOverlap(oLhs, oRhs))
                            fOnFound(oLhs, oRhs, k++);
                    }
                }
            }
            else if (bIsLeftLeaf and not bIsRightLeaf)
            {
                fRecurseRight(p);
            }
            else if (not bIsLeftLeaf and bIsRightLeaf)
            {
                fRecurseLeft(p);
            }
            else
            {
                if (p.levelLhs > p.levelRhs)
                    fRecurseRight(p);
                else
                    fRecurseLeft(p);
            }
        }
    }
    return true;
#include "pbat/warning/Pop.h"
}
 
template <
    class FChild,
    class FIsLeaf,
    class FLeafSize,
    class FLeafObject,
    class FDistanceLowerBound,
    class FDistance,
    class FOnFound,
    auto N,
    class TScalar,
    class TIndex,
    auto kStackDepth,
    auto kQueueSize>
bool DfsAllNearestNeighbours(
    FChild fChild,
    FIsLeaf fIsLeaf,
    FLeafSize fLeafSize,
    FLeafObject fLeafObject,
    FDistanceLowerBound fLower,
    FDistance fDistance,
    FOnFound fOnFound,
    bool bUseBestFirstSearch,
    TScalar fUpper,
    TScalar eps,
    TIndex root)
{
    common::Queue<TIndex, kQueueSize> nn{};
    auto fDoVisit = [&] PBAT_HOST_DEVICE(TIndex node, TScalar lower) {
        if (lower > fUpper)
            return false;
        if (fIsLeaf(node))
        {
            TIndex const nLeafObjects = fLeafSize(node);
            for (TIndex i = 0; i < nLeafObjects; ++i)
            {
                TIndex o         = fLeafObject(node, i);
                TScalar const d  = fDistance(o);
                TScalar const lo = fUpper - eps;
                TScalar const hi = fUpper + eps;
                if (d < lo)
                {
                    nn.Clear();
                    nn.Push(o);
                    fUpper = d;
                }
                else if (d <= hi and not nn.IsFull())
                {
                    nn.Push(o);
                }
            }
        }
        return true;
    };
    auto fVisit = [&] PBAT_HOST_DEVICE(TIndex node) {
        return fDoVisit(node, fLower(node));
    };
    // For the nearest neighbour search, it might be best not to blindly run a pre-order
    // traversal of the branch and bound tree. If requested, let's optimistically run a best-first
    // search. At each node, we will sort the children by their lower bounds and visit them in that
    // order.
    common::Stack<TIndex, kStackDepth> stack{};
    if (not fVisit(root))
        return true;
    stack.Push(root);
    if (bUseBestFirstSearch)
    {
        std::array<TIndex, N> ordering{};
        std::array<TIndex, N> children{};
        std::array<TScalar, N> lowers{};
        while (not stack.IsEmpty())
        {
            if (stack.IsFull())
                return false;
            TIndex const node = stack.Pop();
            common::ForRange<0, N>([&]<auto i> PBAT_HOST_DEVICE() {
                TIndex const child   = fChild.template operator()<i>(node);
                ordering[i]          = i;
                children[i]          = child;
                bool const bHasChild = child >= 0;
                lowers[i] = bHasChild ? fLower(child) : std::numeric_limits<TScalar>::max();
            });
            if constexpr (N == 2)
            {
                if (lowers[0] > lowers[1])
                    std::swap(ordering[0], ordering[1]);
            }
            else
            {
                std::sort(ordering.begin(), ordering.end(), [&](TIndex i, TIndex j) {
                    return lowers[i] < lowers[j];
                });
            }
#include "pbat/warning/Push.h"
#include "pbat/warning/SignConversion.h"
            // NOTE: I hope this loop gets unrolled by the compiler!!
            for (int j = N - 1; j >= 0; --j)
            {
                auto const i = ordering[j];
                if (children[i] >= 0)
                {
                    if (fDoVisit(children[i], lowers[i]) and not stack.IsFull())
                        stack.Push(children[i]);
                }
                else
                    break; // The ordering is such that all non-children are at the end, so exit as
                // soon as we encounter non-child.
            }
#include "pbat/warning/Pop.h"
        }
    }
    else
    {
        // NOTE:
        // Instead of using a pre-order traversal, we visit children of the currently visited node
        // simultaneously. This is because the branch and bound tree might store each node's
        // children in contiguous memory, or at least with some notion of locality. Thus, the fLower
        // function will benefit from many cache hits.
        // common::TraverseNAryTreePseudoPreOrder<FVisit, FChild, TIndex, N, kStackDepth>(
        //     fVisit,
        //     fChild,
        //     root);
        while (not stack.IsEmpty())
        {
            if (stack.IsFull())
                return false;
            TIndex const node = stack.Pop();
            common::ForRange<0, N>([&]<auto i> PBAT_HOST_DEVICE() {
                TIndex const child = fChild.template operator()<i>(node);
                if (child >= 0)
                    if (fVisit(child) and not stack.IsFull())
                        stack.Push(child);
            });
        }
    }
    TIndex k{0};
    while (not nn.IsEmpty())
    {
        TIndex o = nn.Top();
        nn.Pop();
        fOnFound(o, fUpper, k++);
    }
    return true;
}
 
template <
    class FChild,
    class FIsLeaf,
    class FLeafSize,
    class FLeafObject,
    class FDistanceLowerBound,
    class FDistance,
    class FOnFound,
    auto N,
    class TScalar,
    class TIndex,
    auto kHeapCapacity>
bool NearestToFurthestNeighbours(
    FChild fChild,
    FIsLeaf fIsLeaf,
    FLeafSize fLeafSize,
    FLeafObject fLeafObject,
    FDistanceLowerBound fLower,
    FDistance fDistance,
    FOnFound fOnFound,
    TScalar fUpper,
    TIndex root)
{
    enum class EQueueItem { Node, Object };
    struct QueueItem
    {
        EQueueItem type; // Indicates if this QueueItem holds a primitive or a volume
        TIndex idx;      // Index of the primitive, if this QueueItem holds a primitive, or index
        // of the node, if this QueueItem holds a volume (recall that node_idx =
        // bv_idx + 1)
        TScalar d; // Distance from this QueueItem to p
    };
    auto fMakeNodeQueueItem = [&](TIndex node) {
        return QueueItem{EQueueItem::Node, node, fLower(node)};
    };
    auto fMakeObjectQueueItem = [&](TIndex object) {
        return QueueItem{EQueueItem::Object, object, fDistance(object)};
    };
    auto fGreater = [](QueueItem const& q1, QueueItem const& q2) {
        return q1.d > q2.d;
    };
    using Comparator = decltype(fGreater);
    using MinHeap    = common::Heap<QueueItem, Comparator, kHeapCapacity>;
    MinHeap heap{fGreater};
    heap.Push(fMakeNodeQueueItem(root));
    TIndex k{0};
    while (not heap.IsEmpty())
    {
        if (heap.IsFull())
            return false;
 
        QueueItem const q = heap.Pop();
        if (q.d > fUpper)
            break;
        if (q.type == EQueueItem::Node)
        {
            TIndex const node = q.idx;
            if (fIsLeaf(node))
            {
                TIndex const nLeafObjects = fLeafSize(node);
                for (TIndex i = 0; i < nLeafObjects; ++i)
                {
                    if (not heap.IsFull())
                    {
                        TIndex const o = fLeafObject(node, i);
                        heap.Push(fMakeObjectQueueItem(o));
                    }
                }
            }
            else
            {
                common::ForRange<0, N>([&]<auto i> PBAT_HOST_DEVICE() {
                    TIndex const child = fChild.template operator()<i>(node);
                    if (child >= 0 and not heap.IsFull())
                        heap.Push(fMakeNodeQueueItem(child));
                });
            }
        }
        else
        {
            bool const bContinue = fOnFound(q.idx, q.d, k++);
            if (not bContinue)
                break;
        }
    }
    return true;
}
 
template <
    class FChild,
    class FIsLeaf,
    class FLeafSize,
    class FLeafObject,
    class FDistanceLowerBound,
    class FDistance,
    class FOnFound,
    auto N,
    class TScalar,
    class TIndex,
    auto kHeapCapacity>
PBAT_HOST_DEVICE bool AllNearestNeighbours(
    FChild fChild,
    FIsLeaf fIsLeaf,
    FLeafSize fLeafSize,
    FLeafObject fLeafObject,
    FDistanceLowerBound fLower,
    FDistance fDistance,
    FOnFound fOnFound,
    TScalar fUpper,
    TScalar eps,
    TIndex root)
{
    TScalar dmin;
    auto fOnFoundWrapper = [&](TIndex o, TScalar d, TIndex k) {
        if (k == 0)
            dmin = d;
        bool const bIsNearest = d <= dmin + eps;
        if (bIsNearest)
            fOnFound(o, d, k);
        return bIsNearest;
    };
    return NearestToFurthestNeighbours<
        FChild,
        FIsLeaf,
        FLeafSize,
        FLeafObject,
        FDistanceLowerBound,
        FDistance,
        decltype(fOnFoundWrapper),
        N,
        TScalar,
        TIndex,
        kHeapCapacity>(
        fChild,
        fIsLeaf,
        fLeafSize,
        fLeafObject,
        fLower,
        fDistance,
        fOnFoundWrapper,
        fUpper,
        root);
}
 
template <
    class FChild,
    class FIsLeaf,
    class FLeafSize,
    class FLeafObject,
    class FDistanceLowerBound,
    class FDistance,
    class FOnFound,
    auto N,
    class TScalar,
    class TIndex,
    auto kHeapCapacity>
bool KNearestNeighbours(
    FChild fChild,
    FIsLeaf fIsLeaf,
    FLeafSize fLeafSize,
    FLeafObject fLeafObject,
    FDistanceLowerBound fLower,
    FDistance fDistance,
    FOnFound fOnFound,
    TIndex K,
    TScalar fUpper,
    TIndex root)
{
    auto fOnFoundWrapper = [&](TIndex o, TScalar d, TIndex k) {
        fOnFound(o, d, k);
        return (k + 1) < K;
    };
    return NearestToFurthestNeighbours<
        FChild,
        FIsLeaf,
        FLeafSize,
        FLeafObject,
        FDistanceLowerBound,
        FDistance,
        decltype(fOnFoundWrapper),
        N,
        TScalar,
        TIndex,
        kHeapCapacity>(
        fChild,
        fIsLeaf,
        fLeafSize,
        fLeafObject,
        fLower,
        fDistance,
        fOnFoundWrapper,
        fUpper,
        root);
}
} // namespace pbat::geometry
 
#endif // PBAT_GEOMETRY_SPATIALSEARCH_H