com.unity.demoteam.digital-human/Runtime/MeshTools/KdTree3.cs

using System;
using UnityEngine;
using UnityEngine.Profiling;
using Unity.Collections;
using Unity.Collections.LowLevel.Unsafe;
using Unity.Jobs;
using Unity.Burst;

[Serializable]
public class KdTree3
{
    public struct Node
    {
        public int point;
        public int stepL;
        public int stepR;
    }

    public struct Point3
    {
        public float x;// note: 'x' MUST be first field
        public float y;
        public float z;
        public int index;
    }

    public int size;
    public Node[] nodes;
    public Point3[] points;

    public KdTree3(Vector3[] pointCloud, int pointCount)
    {
        BuildFrom(pointCloud, pointCount);
    }

    //----------------------
    // kd-tree construction

    public void BuildFrom(Vector3[] pointBuffer, int pointCount)
    {
        ArrayUtils.ResizeCheckedIfLessThan(ref this.nodes, pointCount);
        ArrayUtils.ResizeCheckedIfLessThan(ref this.points, pointCount);

        Profiler.BeginSample("kd-prep");
        for (int i = 0; i != pointCount; i++)
        {
            this.points[i].x = pointBuffer[i].x;
            this.points[i].y = pointBuffer[i].y;
            this.points[i].z = pointBuffer[i].z;
            this.points[i].index = i;
        }
        Profiler.EndSample();

        Profiler.BeginSample("kd-build");
        if (pointCount > 0)
        {
            unsafe
            {
                int numThreads = SystemInfo.processorCount;
                //Debug.Log("numThreads = " + numThreads);

                int schedLeaves = Mathf.NextPowerOfTwo(numThreads) * 2;
                int schedNodes = BalancedBinaryTreeInfo.NodeCountFromLeafCount(schedLeaves);
                int schedDepth = BalancedBinaryTreeInfo.DepthFromLeafCount(schedLeaves) - 1;

                //Debug.Log("numThreads=" + numThreads + ", schedLeaves=" + schedLeaves + ", schedNodes=" + schedNodes + ", schedDepth=" + schedDepth);

                var jobArray = (JobHandle*)UnsafeUtility.Malloc(sizeof(JobHandle) * schedNodes, 1, Allocator.Temp);
                var jobSched = jobArray;

                fixed (Node* __node = this.nodes)
                fixed (Point3* __points = this.points)
                {
                    ScheduleNode(ref jobSched, null, __node, __points, 0, pointCount, 0, schedDepth);

                    JobHandle.ScheduleBatchedJobs();

                    long jobCount = jobSched - jobArray;
                    for (long i = 0; i != jobCount; i++)
                    {
                        jobArray[i].Complete();
                    }
                }

                UnsafeUtility.Free(jobArray, Allocator.Temp);
            }
        }
        Profiler.EndSample();

        size = pointCount;
    }

    unsafe static void ScheduleNode(ref JobHandle* jobSched, JobHandle* jobDepends, Node* node, Point3* points, int offset, int length, int depth, int depthSingleThreaded)
    {
        // pick median
        int median = length >> 1;

        // calc offsets
        var offsetL = offset;
        var offsetR = offset + median + 1;
        var lengthL = median;
        var lengthR = length - median - 1;
        var stepL = lengthL > 0 ? 1 : 0;
        var stepR = lengthR > 0 ? 1 + lengthL : 0;

        // make job
        var job = new BuildNodeJob()
        {
            node = node,
            points = points,
            offset = offset,
            length = length,
            depth = depth,
            leaf = (depth >= depthSingleThreaded),
        };

        var jobHandle = jobSched++;
        if (jobDepends != null)
            *jobHandle = job.Schedule(*jobDepends);
        else
            *jobHandle = job.Schedule();

        if (depth == 0)
            JobHandle.ScheduleBatchedJobs();

        if (job.leaf)
            return;

        // schedule subtrees
        if (lengthL > 0) ScheduleNode(ref jobSched, jobHandle, node + stepL, points, offsetL, lengthL, depth + 1, depthSingleThreaded);
        if (lengthR > 0) ScheduleNode(ref jobSched, jobHandle, node + stepR, points, offsetR, lengthR, depth + 1, depthSingleThreaded);
    }

    [BurstCompile]
    unsafe struct BuildNodeJob : IJob
    {
        [NativeDisableUnsafePtrRestriction] public Node* node;
        [NativeDisableUnsafePtrRestriction] public Point3* points;// shared

        public int offset;
        public int length;
        public int depth;
        public bool leaf;

        public void Execute()
        {
            if (leaf)
                BuildNode(node, points, offset, length, depth);
            else
                BuildNodeDeferSubtrees(node, points, offset, length, depth);
        }
    }

    unsafe static void BuildNode(Node* node, Point3* points, int offset, int length, int depth)
    {
        // pick median
        int median = length >> 1;

        // pick splitting axis
        int axis = depth % 3;

        // split points by median
        SelectByAxis(median, points + offset, length, axis);

        // calc offsets
        var offsetL = offset;
        var offsetR = offset + median + 1;
        var lengthL = median;
        var lengthR = length - median - 1;
        var stepL = lengthL > 0 ? 1 : 0;
        var stepR = lengthR > 0 ? 1 + lengthL : 0;

        // make node
        node->point = offset + median;
        node->stepL = stepL;
        node->stepR = stepR;

        // build subtrees
        if (lengthL > 0) BuildNode(node + stepL, points, offsetL, lengthL, depth + 1);
        if (lengthR > 0) BuildNode(node + stepR, points, offsetR, lengthR, depth + 1);
    }

    unsafe static void BuildNodeDeferSubtrees(Node* node, Point3* points, int offset, int length, int depth)
    {
        // pick median
        int median = length >> 1;

        // pick splitting axis
        int axis = depth % 3;

        // split points by median
        SelectByAxis(median, points + offset, length, axis);

        // calc offsets
        var lengthL = median;
        var lengthR = length - median - 1;
        var stepL = lengthL > 0 ? 1 : 0;
        var stepR = lengthR > 0 ? 1 + lengthL : 0;

        // make node
        node->point = offset + median;
        node->stepL = stepL;
        node->stepR = stepR;
    }

    unsafe static void SelectByAxis(int nth, Point3* points, int length, int axis)
    {
        // note: this function adapted from public domain variants listed in KdTreeUtils.cs
        const int strideLsh = 2;

        var i = 0;
        var j = length - 1;
        var v = &points->x + axis;

        while (i < j)
        {
            var r = i;
            var w = j;

            float k = v[((r + w) >> 1) << strideLsh];

            while (r < w)
            {
                if (v[r << strideLsh] >= k)
                {
                    Point3 pw = points[w];
                    points[w] = points[r];
                    points[r] = pw;
                    w--;
                }
                else
                {
                    r++;
                }
            }

            if (v[r << strideLsh] > k)
                r--;

            if (nth <= r)
                j = r;
            else
                i = r + 1;
        }
    }

    //-----------------
    // kd-tree queries

    public int FindNearest(ref Vector3 target)
    {
        var bestDist = float.PositiveInfinity;
        var bestNode = -1;

        Profiler.BeginSample("kd-nearest");
        FindNearest(ref bestDist, ref bestNode, 0, 0, ref target);
        Profiler.EndSample();

        if (bestNode != -1)
            return points[nodes[bestNode].point].index;
        else
            return -1;
    }

    public bool FindNearest(ref float bestDist, ref int bestNode, ref Vector3 target)
    {
        var __bestDist = float.PositiveInfinity;
        var __bestNode = -1;

        Profiler.BeginSample("kd-nearest");
        FindNearest(ref __bestDist, ref __bestNode, 0, 0, ref target);
        Profiler.EndSample();

        if (__bestNode != -1)
        {
            bestDist = __bestDist;
            bestNode = points[nodes[__bestNode].point].index;
            return true;
        }
        else
        {
            return false;
        }
    }

    unsafe void FindNearest(ref float bestDist, ref int bestNode, int node, int depth, ref Vector3 target)
    {
        // update best index
        int point = nodes[node].point;
        Vector3 r;
        r.x = target.x - points[point].x;
        r.y = target.y - points[point].y;
        r.z = target.z - points[point].z;

        var dist = r.x * r.x + r.y * r.y + r.z * r.z;
        if (dist < bestDist)
        {
            bestDist = dist;
            bestNode = node;
        }

        // pick search axis
        int axis = depth % 3;

        // pick near, far
        var delta = *(&r.x + axis);// avoid calling operator[]
        int stepN = delta < 0.0f ? nodes[node].stepL : nodes[node].stepR;
        int stepF = delta < 0.0f ? nodes[node].stepR : nodes[node].stepL;

        // search near
        if (stepN != 0)
        {
            FindNearest(ref bestDist, ref bestNode, node + stepN, depth + 1, ref target);
        }

        // search far
        if (stepF != 0 && delta * delta < bestDist)
        {
            FindNearest(ref bestDist, ref bestNode, node + stepF, depth + 1, ref target);
        }
    }

    public int RaycastApprox(ref Vector3 origin, ref Vector3 direction, int maxIterations = 100)
    {
        var position = origin;
        var numIterations = 0;

        var bestDist = float.PositiveInfinity;
        var bestNode = -1;

        while (numIterations++ < maxIterations)
        {
            var stepDist = float.PositiveInfinity;
            var stepNode = -1;

            FindNearest(ref stepDist, ref stepNode, ref position);

            if (stepDist < bestDist)
            {
                bestDist = stepDist;
                bestNode = stepNode;

                if (bestDist < float.Epsilon)
                {
                    break;// crude termination criteria
                }
            }

            position = position + Mathf.Sqrt(stepDist) * direction;
        }

        return bestNode;
    }
}

public static class BalancedBinaryTreeInfo
{
    public static int LeafCountFromDepth(int depth)
    {
        int leafCount = (depth > 0) ? 1 : 0;
        while (--depth > 0) leafCount = (leafCount << 1);
        return leafCount;
    }

    public static int NodeCountFromDepth(int depth)
    {
        int nodeCount = 0;
        while (--depth >= 0) nodeCount = (nodeCount << 1) | 1;
        return nodeCount;
    }

    public static int NodeCountFromLeafCount(int leafCount)
    {
        Debug.Assert(Mathf.IsPowerOfTwo(leafCount));
        int nodeCount = leafCount;
        while ((leafCount >>= 1) > 0) nodeCount += leafCount;
        return nodeCount;
    }

    public static int DepthFromLeafCount(int leafCount)
    {
        Debug.Assert(Mathf.IsPowerOfTwo(leafCount));
        int depth = 0;
        while (leafCount > 0) { leafCount >>= 1; depth++; }
        return depth;
    }
}