Fixed heapqueue.del for last elem (#5363)

1 files changed, 31 insertions, 10 deletions

diff --git a/lib/pure/collections/heapqueue.nim b/lib/pure/collections/heapqueue.nim
index abe20e556..f86ba1d3f 100644
--- a/lib/pure/collections/heapqueue.nim
+++ b/lib/pure/collections/heapqueue.nim
@@ -77,8 +77,10 @@ proc pop*[T](heap: var HeapQueue[T]): T =
 proc del*[T](heap: var HeapQueue[T], index: int) =
   ## Removes element at `index`, maintaining the heap invariant.
   swap(seq[T](heap)[^1], seq[T](heap)[index])
-  seq[T](heap).setLen(heap.len - 1)
-  heap.siftup(index)
+  let newLen = heap.len - 1
+  seq[T](heap).setLen(newLen)
+  if index < newLen:
+    heap.siftup(index)
 
 proc replace*[T](heap: var HeapQueue[T], item: T): T =
   ## Pop and return the current smallest value, and add the new item.
@@ -101,16 +103,19 @@ proc pushpop*[T](heap: var HeapQueue[T], item: T): T =
   return item
 
 when isMainModule:
+  proc toSortedSeq[T](h: HeapQueue[T]): seq[T] =
+    var tmp = h
+    result = @[]
+    while tmp.len > 0:
+      result.add(pop(tmp))
+
   block: # Simple sanity test
     var heap = newHeapQueue[int]()
     let data = [1, 3, 5, 7, 9, 2, 4, 6, 8, 0]
     for item in data:
       push(heap, item)
     doAssert(heap[0] == 0)
-    var sort = newSeq[int]()
-    while heap.len > 0:
-      sort.add(pop(heap))
-    doAssert(sort == @[0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
+    doAssert(heap.toSortedSeq == @[0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
 
   block: # Test del
     var heap = newHeapQueue[int]()
@@ -121,11 +126,27 @@ when isMainModule:
     doAssert(heap[0] == 1)
 
     heap.del(seq[int](heap).find(7))
+    doAssert(heap.toSortedSeq == @[1, 2, 3, 4, 5, 6, 8, 9])
+
     heap.del(seq[int](heap).find(5))
+    doAssert(heap.toSortedSeq == @[1, 2, 3, 4, 6, 8, 9])
+
     heap.del(seq[int](heap).find(6))
+    doAssert(heap.toSortedSeq == @[1, 2, 3, 4, 8, 9])
+
     heap.del(seq[int](heap).find(2))
+    doAssert(heap.toSortedSeq == @[1, 3, 4, 8, 9])
+
+  block: # Test del last
+    var heap = newHeapQueue[int]()
+    let data = [1, 2, 3]
+    for item in data: push(heap, item)
+
+    heap.del(2)
+    doAssert(heap.toSortedSeq == @[1, 2])
 
-    var sort = newSeq[int]()
-    while heap.len > 0:
-      sort.add(pop(heap))
-    doAssert(sort == @[1, 3, 4, 8, 9])
+    heap.del(1)
+    doAssert(heap.toSortedSeq == @[1])
+
+    heap.del(0)
+    doAssert(heap.toSortedSeq == @[])