From 9a3f001b879679aa9c9383268852d71f9aee51cb Mon Sep 17 00:00:00 2001 From: Eirik Tsarpalis Date: Wed, 17 Feb 2021 04:37:10 +0000 Subject: [PATCH] PriorityQueue: Apply Comparer.Default optimization (#48346) * PriorityQueue: inline Comparer.Default.Compare calls * address feedback * Update src/libraries/System.Collections/src/System/Collections/Generic/PriorityQueue.cs Co-authored-by: Stephen Toub * address feedback * Update src/libraries/System.Collections/src/System/Collections/Generic/PriorityQueue.cs Co-authored-by: Stephen Toub Co-authored-by: Stephen Toub --- .../System/Collections/Generic/PriorityQueue.cs | 249 ++++++++++++++++----- 1 file changed, 193 insertions(+), 56 deletions(-) diff --git a/src/libraries/System.Collections/src/System/Collections/Generic/PriorityQueue.cs b/src/libraries/System.Collections/src/System/Collections/Generic/PriorityQueue.cs index ede3739..a99ed5d 100644 --- a/src/libraries/System.Collections/src/System/Collections/Generic/PriorityQueue.cs +++ b/src/libraries/System.Collections/src/System/Collections/Generic/PriorityQueue.cs @@ -21,6 +21,11 @@ namespace System.Collections.Generic private (TElement Element, TPriority Priority)[] _nodes; /// + /// Custom comparer used to order the heap. + /// + private IComparer? _comparer; + + /// /// Lazily-initialized collection used to expose the contents of the queue. /// private UnorderedItemsCollection? _unorderedItems; @@ -60,8 +65,9 @@ namespace System.Collections.Generic /// Creates an empty priority queue. /// public PriorityQueue() - : this(initialCapacity: 0, comparer: null) { + _nodes = Array.Empty<(TElement, TPriority)>(); + _comparer = InitializeComparer(null); } /// @@ -76,8 +82,9 @@ namespace System.Collections.Generic /// Creates an empty priority queue with the specified priority comparer. /// public PriorityQueue(IComparer? comparer) - : this(initialCapacity: 0, comparer) { + _nodes = Array.Empty<(TElement, TPriority)>(); + _comparer = InitializeComparer(comparer); } /// @@ -92,11 +99,8 @@ namespace System.Collections.Generic nameof(initialCapacity), initialCapacity, SR.ArgumentOutOfRange_NeedNonNegNum); } - _nodes = (initialCapacity == 0) - ? Array.Empty<(TElement, TPriority)>() - : new (TElement, TPriority)[initialCapacity]; - - Comparer = comparer ?? Comparer.Default; + _nodes = new (TElement, TPriority)[initialCapacity]; + _comparer = InitializeComparer(comparer); } /// @@ -119,7 +123,7 @@ namespace System.Collections.Generic } _nodes = EnumerableHelpers.ToArray(items, out _size); - Comparer = comparer ?? Comparer.Default; + _comparer = InitializeComparer(comparer); if (_size > 1) { @@ -135,7 +139,7 @@ namespace System.Collections.Generic /// /// Gets the priority comparer of the priority queue. /// - public IComparer Comparer { get; } + public IComparer Comparer => _comparer ?? Comparer.Default; /// /// Gets a collection that enumerates the elements of the queue. @@ -157,7 +161,14 @@ namespace System.Collections.Generic // Restore the heap order int lastNodeIndex = GetLastNodeIndex(); - MoveUp((element, priority), lastNodeIndex); + if (_comparer == null) + { + MoveUpDefaultComparer((element, priority), lastNodeIndex); + } + else + { + MoveUpCustomComparer((element, priority), lastNodeIndex); + } } /// @@ -186,7 +197,7 @@ namespace System.Collections.Generic } TElement element = _nodes[RootIndex].Element; - Remove(RootIndex); + RemoveRootNode(); return element; } @@ -201,7 +212,7 @@ namespace System.Collections.Generic if (_size != 0) { (element, priority) = _nodes[RootIndex]; - Remove(RootIndex); + RemoveRootNode(); return true; } @@ -237,15 +248,32 @@ namespace System.Collections.Generic if (_size != 0) { (TElement Element, TPriority Priority) root = _nodes[RootIndex]; - if (Comparer.Compare(priority, root.Priority) > 0) + + if (_comparer == null) { - (TElement Element, TPriority Priority) newRoot = (element, priority); - _nodes[RootIndex] = newRoot; + if (Comparer.Default.Compare(priority, root.Priority) > 0) + { + (TElement Element, TPriority Priority) newRoot = (element, priority); + _nodes[RootIndex] = newRoot; - MoveDown(newRoot, RootIndex); - _version++; + MoveDownDefaultComparer(newRoot, RootIndex); + _version++; - return root.Element; + return root.Element; + } + } + else + { + if (_comparer.Compare(priority, root.Priority) > 0) + { + (TElement Element, TPriority Priority) newRoot = (element, priority); + _nodes[RootIndex] = newRoot; + + MoveDownCustomComparer(newRoot, RootIndex); + _version++; + + return root.Element; + } } } @@ -411,9 +439,9 @@ namespace System.Collections.Generic } /// - /// Removes the node at the specified index. + /// Removes the node from the root of the heap /// - private void Remove(int indexOfNodeToRemove) + private void RemoveRootNode() { // The idea is to replace the specified node by the very last // node and shorten the array by one. @@ -424,34 +452,17 @@ namespace System.Collections.Generic { _nodes[lastNodeIndex] = default; } + _size--; _version++; - // In case we wanted to remove the node that was the last one, - // we are done. - - if (indexOfNodeToRemove == lastNodeIndex) - { - return; - } - - // Our last node was erased from the array and needs to be - // inserted again. Of course, we will overwrite the node we - // wanted to remove. After that operation, we will need - // to restore the heap property (in general). - - (TElement Element, TPriority Priority) nodeToRemove = _nodes[indexOfNodeToRemove]; - - int relation = Comparer.Compare(lastNode.Priority, nodeToRemove.Priority); - _nodes[indexOfNodeToRemove] = lastNode; - - if (relation < 0) + if (_comparer == null) { - MoveUp(lastNode, indexOfNodeToRemove); + MoveDownDefaultComparer(lastNode, RootIndex); } else { - MoveDown(lastNode, indexOfNodeToRemove); + MoveDownCustomComparer(lastNode, RootIndex); } } @@ -483,28 +494,42 @@ namespace System.Collections.Generic int lastNodeIndex = GetLastNodeIndex(); int lastParentWithChildren = GetParentIndex(lastNodeIndex); - for (int index = lastParentWithChildren; index >= 0; --index) + if (_comparer == null) { - MoveDown(_nodes[index], index); + for (int index = lastParentWithChildren; index >= 0; --index) + { + MoveDownDefaultComparer(_nodes[index], index); + } + } + else + { + for (int index = lastParentWithChildren; index >= 0; --index) + { + MoveDownCustomComparer(_nodes[index], index); + } } } /// /// Moves a node up in the tree to restore heap order. /// - private void MoveUp((TElement Element, TPriority Priority) node, int nodeIndex) + private void MoveUpDefaultComparer((TElement Element, TPriority Priority) node, int nodeIndex) { // Instead of swapping items all the way to the root, we will perform // a similar optimization as in the insertion sort. + Debug.Assert(_comparer is null); + + (TElement Element, TPriority Priority)[] nodes = _nodes; + while (nodeIndex > 0) { int parentIndex = GetParentIndex(nodeIndex); - (TElement Element, TPriority Priority) parent = _nodes[parentIndex]; + (TElement Element, TPriority Priority) parent = nodes[parentIndex]; - if (Comparer.Compare(node.Priority, parent.Priority) < 0) + if (Comparer.Default.Compare(node.Priority, parent.Priority) < 0) { - _nodes[nodeIndex] = parent; + nodes[nodeIndex] = parent; nodeIndex = parentIndex; } else @@ -513,31 +538,118 @@ namespace System.Collections.Generic } } - _nodes[nodeIndex] = node; + nodes[nodeIndex] = node; + } + + /// + /// Moves a node up in the tree to restore heap order. + /// + private void MoveUpCustomComparer((TElement Element, TPriority Priority) node, int nodeIndex) + { + // Instead of swapping items all the way to the root, we will perform + // a similar optimization as in the insertion sort. + + Debug.Assert(_comparer is not null); + + IComparer comparer = _comparer; + (TElement Element, TPriority Priority)[] nodes = _nodes; + + while (nodeIndex > 0) + { + int parentIndex = GetParentIndex(nodeIndex); + (TElement Element, TPriority Priority) parent = nodes[parentIndex]; + + if (comparer.Compare(node.Priority, parent.Priority) < 0) + { + nodes[nodeIndex] = parent; + nodeIndex = parentIndex; + } + else + { + break; + } + } + + nodes[nodeIndex] = node; + } + + /// + /// Moves a node down in the tree to restore heap order. + /// + private void MoveDownDefaultComparer((TElement Element, TPriority Priority) node, int nodeIndex) + { + // The node to move down will not actually be swapped every time. + // Rather, values on the affected path will be moved up, thus leaving a free spot + // for this value to drop in. Similar optimization as in the insertion sort. + + Debug.Assert(_comparer is null); + + (TElement Element, TPriority Priority)[] nodes = _nodes; + int size = _size; + + int i; + while ((i = GetFirstChildIndex(nodeIndex)) < size) + { + // Check if the current node (pointed by 'nodeIndex') should really be extracted + // first, or maybe one of its children should be extracted earlier. + (TElement Element, TPriority Priority) topChild = nodes[i]; + int childrenIndexesLimit = Math.Min(i + Arity, size); + int topChildIndex = i; + + while (++i < childrenIndexesLimit) + { + (TElement Element, TPriority Priority) child = nodes[i]; + if (Comparer.Default.Compare(child.Priority, topChild.Priority) < 0) + { + topChild = child; + topChildIndex = i; + } + } + + // In case no child needs to be extracted earlier than the current node, + // there is nothing more to do - the right spot was found. + if (Comparer.Default.Compare(node.Priority, topChild.Priority) <= 0) + { + break; + } + + // Move the top child up by one node and now investigate the + // node that was considered to be the top child (recursive). + nodes[nodeIndex] = topChild; + nodeIndex = topChildIndex; + } + + nodes[nodeIndex] = node; } /// /// Moves a node down in the tree to restore heap order. /// - private void MoveDown((TElement Element, TPriority Priority) node, int nodeIndex) + private void MoveDownCustomComparer((TElement Element, TPriority Priority) node, int nodeIndex) { // The node to move down will not actually be swapped every time. // Rather, values on the affected path will be moved up, thus leaving a free spot // for this value to drop in. Similar optimization as in the insertion sort. + Debug.Assert(_comparer is not null); + + IComparer comparer = _comparer; + (TElement Element, TPriority Priority)[] nodes = _nodes; + int size = _size; + int i; - while ((i = GetFirstChildIndex(nodeIndex)) < _size) + while ((i = GetFirstChildIndex(nodeIndex)) < size) { // Check if the current node (pointed by 'nodeIndex') should really be extracted // first, or maybe one of its children should be extracted earlier. - (TElement Element, TPriority Priority) topChild = _nodes[i]; - int childrenIndexesLimit = Math.Min(i + Arity, _size); + (TElement Element, TPriority Priority) topChild = nodes[i]; + int childrenIndexesLimit = Math.Min(i + Arity, size); int topChildIndex = i; while (++i < childrenIndexesLimit) { - (TElement Element, TPriority Priority) child = _nodes[i]; - if (Comparer.Compare(child.Priority, topChild.Priority) < 0) + (TElement Element, TPriority Priority) child = nodes[i]; + if (comparer.Compare(child.Priority, topChild.Priority) < 0) { topChild = child; topChildIndex = i; @@ -546,14 +658,14 @@ namespace System.Collections.Generic // In case no child needs to be extracted earlier than the current node, // there is nothing more to do - the right spot was found. - if (Comparer.Compare(node.Priority, topChild.Priority) <= 0) + if (comparer.Compare(node.Priority, topChild.Priority) <= 0) { break; } // Move the top child up by one node and now investigate the // node that was considered to be the top child (recursive). - _nodes[nodeIndex] = topChild; + nodes[nodeIndex] = topChild; nodeIndex = topChildIndex; } @@ -561,6 +673,31 @@ namespace System.Collections.Generic } /// + /// Initializes the custom comparer to be used internally by the heap. + /// + private static IComparer? InitializeComparer(IComparer? comparer) + { + if (typeof(TPriority).IsValueType) + { + if (comparer == Comparer.Default) + { + // if the user manually specifies the default comparer, + // revert to using the optimized path. + return null; + } + + return comparer; + } + else + { + // Currently the JIT doesn't optimize direct Comparer.Default.Compare + // calls for reference types, so we want to cache the comparer instance instead. + // TODO https://github.com/dotnet/runtime/issues/10050: Update if this changes in the future. + return comparer ?? Comparer.Default; + } + } + + /// /// Represents the contents of a without ordering. /// public sealed class UnorderedItemsCollection : IReadOnlyCollection<(TElement Element, TPriority Priority)>, ICollection -- 2.7.4