(define-module (algorithms priority-queue)
  #:use-module (ice-9 match)
  #:export (make-priority-queue
            pq-length
            pq-push!
            pq-pop!
            pq-peek
            pq-empty?
            pq-drain!))

;; The heap implementation is based on "Chapter 6: Heapsort" from "Introduction
;; to Algorithms" by Cormen, Leiserson, Rivest, and Stein.

;; We store the heap in a vector. The first element of the vector is the current
;; heap size, so elements of the heap are indexed starting from 1.

(define (parent i)
  (floor-quotient i 2))

(define (left-child i)
  (* 2 i))

(define (right-child i)
  (1+ (* 2 i)))

(define (vector-swap! A i j)
  (let ((tmp (vector-ref A i)))
    (vector-set! A i (vector-ref A j))
    (vector-set! A j tmp)))

;; The sink function is called to maintain the heap property when the root node
;; is removed and swapped with the last node in the heap. The new root node
;; might need to "sink" down to maintain the heap ordering.
(define (sink A i priority cmp)
  (let ((heap-size (vector-ref A 0)))
    (let loop ((i i))
      (let ((l (left-child i))
            (r (right-child i)))
        (cond
         ((and (<= r heap-size)
               (cmp (priority (vector-ref A r)) (priority (vector-ref A l)))
               (cmp (priority (vector-ref A r)) (priority (vector-ref A i))))
          (vector-swap! A i r)
          (loop r))
         ((and (<= l heap-size)
               (cmp (priority (vector-ref A l))
                    (priority (vector-ref A i))))
          (vector-swap! A i l)
          (loop l)))))))

;; The swim function is used to restore the heap property when a new node is
;; added to the heap: it's initially added at the end of the heap and may need to
;; "swim" up the heap to maintain the order.
(define (swim A i priority cmp)
  (let loop ((i i))
    (when (and (> (parent i) 0) (cmp (priority (vector-ref A i))
                                     (priority (vector-ref A (parent i)))))
      (vector-swap! A i (parent i))
      (loop (parent i)))))

;; make-priority-queue constructs a new heap-based priority queue
;; cmp should be > for a max-priority queue or < for a min-priority queue
;; priority is a function that returns the priority of an element
(define* (make-priority-queue #:key (cmp >) (priority identity) (initial-capacity 16))
  (let ((A (make-vector initial-capacity)))
    (vector-set! A 0 0)
    (match-lambda*
      (('size)
       (vector-ref A 0))
      (('peek)
       (let ((heap-size (vector-ref A 0)))
         (if (< heap-size 1)
             (error "heap underflow")
             (vector-ref A 1))))
      (('pop!)
       (let ((heap-size (vector-ref A 0)))
         (if (< heap-size 1)
             (error "heap underflow")
             (let ((result (vector-ref A 1)))
               (vector-set! A 1 (vector-ref A heap-size))
               (vector-set! A 0 (1- heap-size))
               (sink A 1 priority cmp)
               result))))
      (('push! v)
       (let ((heap-size (vector-ref A 0))
             (capacity (vector-length A)))
         (when (>= heap-size (1- capacity))
           (let ((new-A (make-vector (* 2 capacity))))
             (vector-copy! new-A 0 A)
             (set! A new-A)))
         (let ((heap-size (1+ heap-size)))
           (vector-set! A 0 heap-size)
           (vector-set! A heap-size v)
           (swim A heap-size priority cmp)))))))

(define (pq-length q)
  (q 'size))

(define (pq-push! q v)
  (q 'push! v))

(define (pq-pop! q)
  (q 'pop!))

(define (pq-peek q)
  (q 'peek))

(define pq-empty? (compose zero? pq-length))

;; Pop all elements of a queue onto a list, emptying the queue in the process.
;; Return the resulting list.
(define (pq-drain! q)
  (let loop ((result '()))
    (if (pq-empty? q)
        (reverse result)
        (loop (cons (pq-pop! q) result)))))