NativePriorityQueue.cs

using System;
using System.Collections.Generic;
using System.Runtime.InteropServices;
using Unity.Burst;
using Unity.Collections;
using Unity.Collections.LowLevel.Unsafe;
using Unity.Jobs;
// ReSharper disable InconsistentNaming

namespace Utilities
{
    [NativeContainer]
    public unsafe struct NativePriorityQueue<T> : IDisposable where T : unmanaged
    {
        [StructLayout(LayoutKind.Sequential)]
        public struct Node : IEquatable<Node>
        {
            public T Value;
            public int Priority;
            internal int Index;

            public bool Equals(Node other) =>
                Index == other.Index && Priority == other.Priority &&
                EqualityComparer<T>.Default.Equals(Value, other.Value);

            public override string ToString() => $"Node [{Value}, {Priority}]";
        }

        [NativeDisableUnsafePtrRestriction] private UnsafeList<Node>* _nodes;

        Node LastUnsafe => ReadUnsafe(_nodes->Length - 1);

        public NativePriorityQueue(int initialCapacity, Allocator allocator)
        {
            _nodes = UnsafeList<Node>.Create(initialCapacity, allocator);
            _nodes->Add(default);
        }

        public void SetCapacity(int capacity) => _nodes->SetCapacity(capacity);

        public int Length => LengthUnsafe;

        public bool IsEmpty => Length == 0;

        int LengthUnsafe => _nodes->Length - 1;

        public int Capacity => CapacityUnsafe;

        int CapacityUnsafe => _nodes->Capacity;
        
        public Node this[int index] => ReadUnsafe(index + 1);

        public bool Contains(Node node)
        {
            var listNode = ReadUnsafe(node.Index);
            return node.Equals(listNode);
        }

        public void Enqueue(T value, int priority)
        {
            var node = new Node
            {
                Value = value,
                Priority = priority,
                Index = LengthUnsafe + 1
            };

            // Add the node to the end of the list
            _nodes->Add(node);

            CascadeUpUnsafe(node);
        }

        public T Dequeue()
        {
            Node returnMe = ReadUnsafe(1);
             
            //If the node is already the last node, we can remove it immediately
            if (LengthUnsafe == 1)
            {
                _nodes->RemoveAtSwapBack(1);
                return returnMe.Value;
            }

            var formerLast = RemoveAtSwapBackUnsafe(1);

            // Bubble down the swapped node, which was prevously the final node and is now in the first position
            CascadeDownUnsafe(formerLast);

            return returnMe.Value;
        }

        public void Clear()
        {
            _nodes->Clear();
            _nodes->Add(default);
        }
        
        public void Dispose()
        {
            UnsafeList<Node>.Destroy(_nodes);
            _nodes = null;
        }

        public JobHandle Dispose(JobHandle inputDeps)
        {

            var jobHandle = new NativePriorityQueueDisposeJob 
            { 
                Data = new NativePriorityQueueDispose { m_ListData = _nodes } 
            }.Schedule(inputDeps);
            
            _nodes = null;

            return jobHandle;
        }

        private Node ReadUnsafe(int index) => UnsafeUtility.ReadArrayElement<Node>(_nodes->Ptr, index);

        private void WriteUnsafe(int index, Node value) => UnsafeUtility.WriteArrayElement(_nodes->Ptr, index, value);

        private void CascadeUpUnsafe(Node node)
        {
            //aka Heapify-up
            int parent;
            if (node.Index > 1)
            {
                parent = node.Index >> 1;
                Node parentNode = ReadUnsafe(parent);

                if (HasHigherOrEqualPriority(parentNode, node))
                    return;

                //Node has lower priority value, so move parent down the heap to make room
                WriteUnsafe(node.Index, parentNode);
                parentNode.Index = node.Index;
                node.Index = parent;
            }
            else
            {
                return;
            }
            while (parent > 1)
            {
                parent >>= 1;
                Node parentNode = ReadUnsafe(parent);
                if (HasHigherOrEqualPriority(parentNode, node))
                    break;

                //Node has lower priority value, so move parent down the heap to make room
                WriteUnsafe(node.Index, parentNode);
                parentNode.Index = node.Index;
                node.Index = parent;
            }
            WriteUnsafe(node.Index, node);
        }

        private void CascadeDownUnsafe(Node node)
        {
            //aka Heapify-down
            int finalQueueIndex = node.Index;
            int childLeftIndex = 2 * finalQueueIndex;

            // If leaf node, we're done
            if (childLeftIndex > LengthUnsafe)
                return;

            // Check if the left-child is higher-priority than the current node
            int childRightIndex = childLeftIndex + 1;
            Node childLeft = ReadUnsafe(childLeftIndex);

            if (HasHigherPriority(childLeft, node))
            {
                // Check if there is a right child. If not, swap and finish.
                if (childRightIndex > LengthUnsafe)
                {
                    node.Index = childLeftIndex;
                    childLeft.Index = finalQueueIndex;
                    WriteUnsafe(finalQueueIndex, childLeft);
                    WriteUnsafe(childLeftIndex, node);
                    return;
                }
                // Check if the left-child is higher-priority than the right-child
                Node childRight = ReadUnsafe(childRightIndex);
                if (HasHigherPriority(childLeft, childRight))
                {
                    // left is highest, move it up and continue
                    childLeft.Index = finalQueueIndex;
                    WriteUnsafe(finalQueueIndex, childLeft);
                    finalQueueIndex = childLeftIndex;
                }
                else
                {
                    // right is even higher, move it up and continue
                    childRight.Index = finalQueueIndex;
                    WriteUnsafe(finalQueueIndex, childRight);
                    finalQueueIndex = childRightIndex;
                }
            }
            // Not swapping with left-child, does right-child exist?
            else if (childRightIndex > LengthUnsafe)
            {
                return;
            }
            else
            {
                // Check if the right-child is higher-priority than the current node
                Node childRight = ReadUnsafe(childRightIndex);
                if (HasHigherPriority(childRight, node))
                {
                    childRight.Index = finalQueueIndex;
                    WriteUnsafe(finalQueueIndex, childRight);
                    finalQueueIndex = childRightIndex;
                }
                // Neither child is higher-priority than current, so finish and stop.
                else
                {
                    return;
                }
            }

            while (true)
            {
                childLeftIndex = 2 * finalQueueIndex;

                // If leaf node, we're done
                if (childLeftIndex > LengthUnsafe)
                {
                    node.Index = finalQueueIndex;
                    WriteUnsafe(finalQueueIndex, node);
                    break;
                }

                // Check if the left-child is higher-priority than the current node
                childRightIndex = childLeftIndex + 1;
                childLeft = ReadUnsafe(childLeftIndex);
                if (HasHigherPriority(childLeft, node))
                {
                    // Check if there is a right child. If not, swap and finish.
                    if (childRightIndex > LengthUnsafe)
                    {
                        node.Index = childLeftIndex;
                        childLeft.Index = finalQueueIndex;
                        WriteUnsafe(finalQueueIndex, childLeft);
                        WriteUnsafe(childLeftIndex, node);
                        break;
                    }
                    // Check if the left-child is higher-priority than the right-child
                    Node childRight = ReadUnsafe(childRightIndex);
                    if (HasHigherPriority(childLeft, childRight))
                    {
                        // left is highest, move it up and continue
                        childLeft.Index = finalQueueIndex;
                        WriteUnsafe(finalQueueIndex, childLeft);
                        finalQueueIndex = childLeftIndex;
                    }
                    else
                    {
                        // right is even higher, move it up and continue
                        childRight.Index = finalQueueIndex;
                        WriteUnsafe(finalQueueIndex, childRight);
                        finalQueueIndex = childRightIndex;
                    }
                }
                // Not swapping with left-child, does right-child exist?
                else if (childRightIndex > LengthUnsafe)
                {
                    node.Index = finalQueueIndex;
                    WriteUnsafe(finalQueueIndex, node);
                    break;
                }
                else
                {
                    // Check if the right-child is higher-priority than the current node
                    Node childRight = ReadUnsafe(childRightIndex);
                    if (HasHigherPriority(childRight, node))
                    {
                        childRight.Index = finalQueueIndex;
                        WriteUnsafe(finalQueueIndex, childRight);
                        finalQueueIndex = childRightIndex;
                    }
                    // Neither child is higher-priority than current, so finish and stop.
                    else
                    {
                        node.Index = finalQueueIndex;
                        WriteUnsafe(finalQueueIndex, node);
                        break;
                    }
                }
            }
        }

        private bool HasHigherPriority(Node higher, Node lower) => higher.Priority < lower.Priority;

        private bool HasHigherOrEqualPriority(Node higher, Node lower) => higher.Priority <= lower.Priority;

        private Node RemoveAtSwapBackUnsafe(int index)
        {
            // Swap the last node with the node at the given index
            var formerLast = LastUnsafe;
            formerLast.Index = index;
            WriteUnsafe(index, formerLast);
            _nodes->RemoveAtSwapBack(LengthUnsafe);
            return formerLast;
        }
        
        [NativeContainer]
        internal struct NativePriorityQueueDispose
        {
            [NativeDisableUnsafePtrRestriction]
            public UnsafeList<Node>* m_ListData;

            public void Dispose()
            {
                UnsafeList<Node>.Destroy(m_ListData);
            }
        }

        [BurstCompile]
        internal struct NativePriorityQueueDisposeJob : IJob
        {
            internal NativePriorityQueueDispose Data;

            public void Execute()
            {
                Data.Dispose();
            }
        }
    }
}