follow#16107 make heapqueue support user-defined cmp function

lib/pure/collections/heapqueue.nim

@@ -52,16 +52,32 @@
 ]##
 import std/private/since
 
-type HeapQueue*[T] = object
-  ## A heap queue, commonly known as a priority queue.
-  data: seq[T]
+type
+  Cmp*[T] = proc (x, y: T): bool
+  HeapQueue*[T] = object
+    ## A heap queue, commonly known as a priority queue.
+    data: seq[T]
+    cmp: Cmp[T]
+
+proc minHeapCmp*[T](x, y: T): bool =
+  result = (x < y)
+
+proc maxHeapCmp*[T](x, y: T): bool =
+  result = (x > y)
 
 proc initHeapQueue*[T](): HeapQueue[T] =
   ## Creates a new empty heap.
   ##
   ## See also:
   ## * `toHeapQueue proc <#toHeapQueue,openArray[T]>`_
-  discard
+  result.cmp = minHeapCmp
+
+proc initHeapQueue*[T](cmp: Cmp[T]): HeapQueue[T] =
+  ## Creates a new empty heap.
+  ##
+  ## See also:
+  ## * `toHeapQueue proc <#toHeapQueue,openArray[T]>`_
+  result.cmp = cmp
 
 proc len*[T](heap: HeapQueue[T]): int {.inline.} =
   ## Returns the number of elements of `heap`.
@@ -71,9 +87,6 @@ proc `[]`*[T](heap: HeapQueue[T], i: Natural): T {.inline.} =
   ## Accesses the i-th element of `heap`.
   heap.data[i]
 
-proc heapCmp[T](x, y: T): bool {.inline.} =
-  return (x < y)
-
 proc siftdown[T](heap: var HeapQueue[T], startpos, p: int) =
   ## 'heap' is a heap at all indices >= startpos, except possibly for `pos`. `pos`
   ## is the index of a leaf with a possibly out-of-order value. Restores the
@@ -82,14 +95,20 @@ proc siftdown[T](heap: var HeapQueue[T], startpos, p: int) =
   var newitem = heap[pos]
   # Follow the path to the root, moving parents down until finding a place
   # newitem fits.
+  let cmp =
+    if heap.cmp != nil:
+      heap.cmp
+    else:
+      minHeapCmp
   while pos > startpos:
     let parentpos = (pos - 1) shr 1
     let parent = heap[parentpos]
-    if heapCmp(newitem, parent):
+    if cmp(newitem, parent):
       heap.data[pos] = parent
       pos = parentpos
     else:
       break
+
   heap.data[pos] = newitem
 
 proc siftup[T](heap: var HeapQueue[T], p: int) =
@@ -99,10 +118,16 @@ proc siftup[T](heap: var HeapQueue[T], p: int) =
   let newitem = heap[pos]
   # Bubble up the smaller child until hitting a leaf.
   var childpos = 2*pos + 1 # leftmost child position
+  let cmp =
+    if heap.cmp != nil:
+      heap.cmp
+    else:
+      minHeapCmp
   while childpos < endpos:
     # Set childpos to index of smaller child.
     let rightpos = childpos + 1
-    if rightpos < endpos and not heapCmp(heap[childpos], heap[rightpos]):
+
+    if rightpos < endpos and not cmp(heap[childpos], heap[rightpos]):
       childpos = rightpos
     # Move the smaller child up.
     heap.data[pos] = heap[childpos]
@@ -174,7 +199,12 @@ proc replace*[T](heap: var HeapQueue[T], item: T): T =
 proc pushpop*[T](heap: var HeapQueue[T], item: T): T =
   ## Fast version of a push followed by a pop.
   result = item
-  if heap.len > 0 and heapCmp(heap.data[0], item):
+  let cmp =
+    if heap.cmp != nil:
+      heap.cmp
+    else:
+      minHeapCmp
+  if heap.len > 0 and cmp(heap.data[0], item):
     swap(result, heap.data[0])
     siftup(heap, 0)