heap-remove! sometimes needs to heapify up

closes PR 14651

pkgs/data-pkgs/data-lib/data/heap.rkt

@@ -149,12 +149,29 @@
                   "empty heap: ~s" index)
            (error 'heap-remove-index!
                   "index out of bounds [0,~s]: ~s" (sub1 size) index)))
-     (vector-set! vec index (vector-ref vec (sub1 size)))
-     (vector-set! vec (sub1 size) #f)
-     (heapify-down <=? vec index (sub1 size))
+     (define sub1-size (sub1 size))
+     (vector-set! vec index (vector-ref vec sub1-size))
+     (vector-set! vec sub1-size #f)
+     (cond
+       [(= sub1-size index)
+        ;; easy to remove the right-most leaf
+        (void)]
+       [(= index 0)
+        ;; can only go down when at the root
+        (heapify-down <=? vec index sub1-size)]
+       [else
+        (define index-parent (vt-parent index))
+        (cond
+          ;; if we are in the right relationship with our parent,
+          ;; try to heapify down
+          [(<=? (vector-ref vec index-parent) (vector-ref vec index))
+           (heapify-down <=? vec index sub1-size)]
+          [else
+           ;; otherwise we need to heapify up
+           (heapify-up <=? vec index)])])
      (when (< MIN-SIZE size (quotient (vector-length vec) 4))
        (set-heap-vec! h (shrink-vector vec)))
-     (set-heap-count! h (sub1 size))]))
+     (set-heap-count! h sub1-size)]))
 
 (define (heap-get-index h v same?)
   (match h
@@ -250,3 +267,18 @@
  [in-heap/consume! (-> heap? sequence?)])
 
 (provide heap-sort!)
+
+(module+ test-util
+  (provide valid-heap?)
+  (define (valid-heap? a-heap)
+    (match a-heap
+      [(heap vec size <=?)
+       (let loop ([i 0]
+                  [parent -inf.0])
+         (cond
+           [(< i size)
+            (define this (vector-ref vec i))
+            (and (<=? parent this)
+                 (loop (vt-leftchild i) this)
+                 (loop (vt-rightchild i) this))]
+           [else #t]))])))

pkgs/data-pkgs/data-test/tests/data/heap.rkt

@@ -1,6 +1,7 @@
 #lang racket
 (require rackunit
-         data/heap)
+         data/heap
+         (submod data/heap test-util))
 
 (define (mkheap) (vector->heap <= (vector 6 2 4 10 8)))
 
@@ -108,3 +109,131 @@
   (test-equal? "heap-sort (random)"
     (begin (heap-sort! v <=) (vector->list v))
     (sort l <)))
+
+(define (heap-add!/v h v)
+  (heap-add! h v)
+  (unless (valid-heap? h) (error 'heap-add!/v "failed")))
+(define (heap-remove!/v h v)
+  (heap-remove! h v)
+  (unless (valid-heap? h) (error 'heap-remove!/v "post failed")))
+(define (heap-remove-min!/v h)
+  (heap-remove-min! h)
+  (unless (valid-heap? h) (error 'heap-remove-min!/v "post failed")))
+
+;; test case from PR14651
+(let ([h (make-heap <=)])
+  (heap-add!/v h 0)
+  (heap-add!/v h -5942)
+  (heap-add!/v h 8358)
+  (heap-add!/v h 569)
+  (heap-add!/v h 6723)
+  (heap-add!/v h -151)
+  (heap-add!/v h 6807)
+  (heap-add!/v h -1612)
+  (heap-remove-min!/v h)
+  (heap-add!/v h -1008)
+  (heap-add!/v h -7157)
+  (heap-add!/v h -1734)
+  (heap-add!/v h 6497)
+  (heap-add!/v h 1603)
+  (heap-add!/v h -7927)
+  (heap-remove!/v h -151)
+  (heap-add!/v h -349)
+  (heap-add!/v h -7570)
+  (heap-remove-min!/v h)
+  (heap-add!/v h 4008)
+  (heap-add!/v h 6101)
+  (heap-add!/v h -9013)
+  (heap-add!/v h -3447)
+  (heap-add!/v h -4294)
+  (heap-add!/v h 8187)
+  (heap-add!/v h 1465)
+  (heap-remove-min!/v h)
+  (heap-add!/v h -1598)
+  (heap-add!/v h 9730)
+  (heap-add!/v h -4429)
+  (heap-add!/v h -846)
+  (heap-add!/v h 4775)
+  (heap-add!/v h 3609)
+  (heap-add!/v h -3881)
+  (heap-add!/v h 6167)
+  (heap-add!/v h 6767)
+  (heap-remove-min!/v h)
+  (heap-add!/v h 2842)
+  (heap-add!/v h -4103)
+  (heap-add!/v h 154)
+  (heap-add!/v h 3748)
+  (heap-add!/v h -536)
+  (heap-add!/v h -5565)
+  (heap-add!/v h 4970)
+  (heap-add!/v h 4775)
+  (heap-add!/v h 4818)
+  (heap-add!/v h 5124)
+  (heap-add!/v h -8657)
+  (heap-add!/v h -6842)
+  (heap-remove-min!/v h)
+  (heap-add!/v h 2480)
+  (heap-add!/v h 8878)
+  (heap-add!/v h -1806)
+  (heap-remove-min!/v h)
+  (heap-add!/v h -8205)
+  (heap-remove!/v h 9730)
+  (heap-add!/v h -3164)
+  (heap-add!/v h 1589)
+  (heap-add!/v h 8444)
+  (heap-add!/v h -7839)
+  (heap-add!/v h -3810)
+  (heap-remove!/v h 4970)
+  ; -1612 out of position
+  (void))
+
+;; simpler test case from PR14651
+(let ([heap (make-heap <=)])
+  (heap-add!/v heap 43)
+  (heap-add!/v heap 1)
+  (heap-add!/v heap 37)
+  (heap-add!/v heap 81)
+  (heap-add!/v heap 94)
+  (heap-add!/v heap 4)
+  (heap-remove!/v heap 94))
+
+(define (random-test)
+  (define heap (make-heap <=))
+  (let loop ([ops '()]
+             [values '()])
+    (cond
+      [(not (valid-heap? heap))
+       (eprintf "crash! ~a ops\n" (length ops))
+       (pretty-write `(let ([heap (make-heap <=)]) ,@(reverse ops))
+                     (current-error-port))]
+      [(= (length ops) 50)
+       (void)]
+      [else
+       (define (do-an-add)
+         (define n (random 10))
+         (heap-add! heap n)
+         (loop (cons `(heap-add!/v heap ,n) ops)
+               (cons n values)))
+       (case (random 3)
+         [(0) (do-an-add)]
+         [(1)
+          (cond
+            [(null? values)
+             (do-an-add)]
+            [else
+             (define to-remove (list-ref values (random (length values))))
+             (heap-remove! heap to-remove)
+             (loop (cons `(heap-remove!/v heap ,to-remove) ops)
+                   (remove to-remove values))])]
+         [(2)
+          (cond
+            [(null? values)
+             (do-an-add)]
+            [else
+             (heap-remove-min! heap)
+             (define smallest (apply min values))
+             (loop (cons `(heap-remove-min!/v heap) ops)
+                   (remove smallest values))])])])))
+
+(for ([x (in-range 10000)])
+  (random-test))