void PushToHeap(int nodeIndex, float3 tempClosestPoint, float3 queryPosition, ref KnnQueryTemp temp)
{
float sqrDist = math.lengthsq(tempClosestPoint - queryPosition);
KdQueryNode queryNode = new KdQueryNode {
NodeIndex = nodeIndex,
TempClosestPoint = tempClosestPoint,
Distance = sqrDist
};
temp.MinHeap.PushObj(queryNode, sqrDist);
}
internal unsafe void KNearest(float3 queryPosition, NativeSlice <int> result, ref KnnQueryTemp temp)
{
int k = result.Length;
temp.Heap.Clear();
var points = Points;
var permutation = m_permutation;
var rootNode = RootNode;
var nodePtr = m_nodes.GetUnsafePtr();
// Biggest Smallest Squared Radius
float bssr = float.PositiveInfinity;
float3 rootClosestPoint = rootNode.Bounds.ClosestPoint(queryPosition);
PushToHeap(m_rootNodeIndex[0], rootClosestPoint, queryPosition, ref temp);
// searching
while (temp.MinHeap.Count > 0)
{
KdQueryNode queryNode = temp.MinHeap.PopObj();
if (queryNode.Distance > bssr)
{
continue;
}
ref KdNode node = ref UnsafeUtilityEx.ArrayElementAsRef <KdNode>(nodePtr, queryNode.NodeIndex);
if (!node.Leaf)
{
int partitionAxis = node.PartitionAxis;
float partitionCoord = node.PartitionCoordinate;
float3 tempClosestPoint = queryNode.TempClosestPoint;
if (tempClosestPoint[partitionAxis] - partitionCoord < 0)
{
// we already know we are on the side of negative bound/node,
// so we don't need to test for distance
// push to stack for later querying
PushToHeap(node.NegativeChildIndex, tempClosestPoint, queryPosition, ref temp);
// project the tempClosestPoint to other bound
tempClosestPoint[partitionAxis] = partitionCoord;
if (node.Count != 0)
{
PushToHeap(node.PositiveChildIndex, tempClosestPoint, queryPosition, ref temp);
}
}
else
{
// we already know we are on the side of positive bound/node,
// so we don't need to test for distance
// push to stack for later querying
PushToHeap(node.PositiveChildIndex, tempClosestPoint, queryPosition, ref temp);
// project the tempClosestPoint to other bound
tempClosestPoint[partitionAxis] = partitionCoord;
if (node.Count != 0)
{
PushToHeap(node.NegativeChildIndex, tempClosestPoint, queryPosition, ref temp);
}
}
}
else
{
for (int i = node.Start; i < node.End; i++)
{
int index = permutation[i];
float sqrDist = math.lengthsq(points[index] - queryPosition);
if (sqrDist <= bssr)
{
temp.Heap.PushObj(index, sqrDist);
if (temp.Heap.Count == k)
{
bssr = temp.Heap.HeadValue;
}
}
}
}
}