#lang racket

;;; Procedure:
;;;   heap-new
;;; Parameters:
;;;   higher-priority?, a binary predicate
;;; Purpose:
;;;   Create a new heap.
;;; Produces:
;;;   heap, a heap
(define heap-new
  (lambda (higher-priority?)
    (vector 0 (make-vector 4) higher-priority?)))

;;;;;;;;;;;;;;;;;;;;;;;;
;;; Helper Procedures ;;
;;;;;;;;;;;;;;;;;;;;;;;;

;;; Procedure:
;;;   get-parent-idx
;;; Parameters:
;;;   pos, an index into a vector (non-negative integer)
;;; Purpose:
;;;   Determine the location of the parent entry to the entry
;;;     in the location pos
;;; Produces:
;;;   parent, a non-negative integer
(define get-parent-idx
  (lambda (pos)
    (if (= pos 0)
        0
        (floor (/ (- pos 1) 2)))))

;;; Procedure:
;;;   get-lc-idx
;;; Parameters:
;;;   pos, an index into a vector (non-negative integer)
;;;   size, the size of the heap's internal vector
;;; Purpose:
;;;   Determine the location of the left-child entry to the entry
;;;     in the location pos
;;; Produces:
;;;   left-child, a non-negative integer
(define get-lc-idx
  (lambda (pos size)
    (let ([candidate (+ 1 (* pos 2))])
      (if (< candidate size)
          candidate
          pos))))

;;; Procedure:
;;;   get-rc-idx
;;; Parameters:
;;;   pos, an index into a vector (non-negative integer)
;;;   size, the size of the heap's internal vector
;;; Purpose:
;;;   Determine the location of the right-child entry to the entry
;;;     in the location pos
;;; Produces:
;;;   right-child, a non-negative integer
(define get-rc-idx
  (lambda (pos size)
    (let ([candidate (+ 2 (* pos 2))])
      (if (< candidate size)
          candidate
          pos))))

;;; Procedure:
;;;   vector-swap!
;;; Parameters:
;;;   vec, a vector
;;;   i, a non-negative integer
;;;   j, a non-negative integer
;;; Purpose:
;;;   Swap the i-th and j-th elements of vec
;;; Produces:
;;;   [nothing, called for side-effect]
;;; Pre-conditions:
;;;   0 <= i < (vector-length vec)
;;;   0 <= j < (vector-length vec)
;;; Post-conditions:
;;;   vec has been mutated
(define vector-swap!
  (lambda (vec i j)
    (let ([temp (vector-ref vec i)])
      (unless (= i j)
        (vector-set! vec i (vector-ref vec j))
        (vector-set! vec j temp)))))

;;; Procedure:
;;;   vector-insert!
;;; Parameters:
;;;   vec, a vector
;;;   pos, the position where to insert the value (non-negative integer)
;;;   val, the value to be inserted
;;; Purpose:
;;;   Inserts the val at the pos-th index of vec. If pos is greater than the
;;;     length of vec, it creates a new vector twice the length of vec, copies
;;;     all the elements of vec into this vector, and inserts val.
;;; Produces:
;;;   [nothing, called for side-effect]
;;; Pre-conditions:
;;;   pos < 2 * (vector-length vec)
;;; Post-conditions:
;;;   vec has been mutated
;;;   (= (vector-ref vec pos) val) evaluates to true
(define vector-insert!
  (lambda (vec pos val)
    (let* ([length (vector-length vec)]
           [make_vec? (>= pos length)]
           [local_vec (if make_vec?
                          (make-vector (* 2 length))
                          vec)])
      (when make_vec?
        (vector-copy! local_vec 0 vec))
      (vector-set! local_vec pos val)
      local_vec)))

;;; Procedure:
;;;   heap-percolate-up!
;;; Parameters:
;;;   heap, a heap created with make-heap
;;; Purpose:
;;;   To place the last element of the heap's internal vector in the correct
;;;     position i.e. ensure that the heap's semantic invariant is preserved:
;;;       for each node in the heap,
;;;         (and (higher-priority? node left-child)
;;;              (higher-priority? node left-child))
;;;       evaluates to true.
;;; Produces:
;;;   [nothing, called for side-effect]
;;; Pre-conditions:
;;;   [no additional]
;;; Post-conditions:
;;;   heap may have been mutated
(define heap-percolate-up!
  (lambda (heap)
    (let* ([last (- (vector-ref heap 0) 1)]
           [vec (vector-ref heap 1)]
           [higher-priority? (vector-ref heap 2)])
      (let kernel ([cur last]
                   [parent (get-parent-idx last)])
        (when (and (> cur 0)
                   (higher-priority? (vector-ref vec cur) (vector-ref vec parent)))
          (vector-swap! vec parent cur)
          (kernel parent (get-parent-idx parent)))))))

;;; Procedure:
;;;   heap-percolate-down!
;;; Parameters:
;;;   heap, a heap created with make-heap
;;; Purpose:
;;;   To place the fist element of the heap's internal vector in the correct
;;;     position i.e. ensure that the heap's semantic invariant is preserved:
;;;       for each node in the heap,
;;;         (and (higher-priority? node left-child)
;;;              (higher-priority? node left-child))
;;;       evaluates to true.
;;; Produces:
;;;   [nothing, called for side-effect]
;;; Pre-conditions:
;;;   [no additional]
;;; Post-conditions:
;;;   heap has been mutated
(define heap-percolate-down!
  (lambda (heap)
    (let* ([size (vector-ref heap 0)]
           [vec (vector-ref heap 1)]
           [higher-priority? (vector-ref heap 2)])
      (let kernel ([cur 0]
                   [lc-idx (get-lc-idx 0 size)]
                   [rc-idx (get-rc-idx 0 size)])
        ; Recurse while the current node is out of place
        (when (and (< cur size)
                   (or (higher-priority? (vector-ref vec lc-idx)
                                         (vector-ref vec cur))
                       (higher-priority? (vector-ref vec rc-idx)
                                         (vector-ref vec cur))))

          (cond
            [(higher-priority? (vector-ref vec lc-idx)
                               (vector-ref vec rc-idx))
             (vector-swap! vec cur lc-idx)
             (kernel lc-idx (get-lc-idx lc-idx size) (get-rc-idx lc-idx size))]
            [else
             (vector-swap! vec cur rc-idx)
             (kernel rc-idx (get-lc-idx rc-idx size) (get-rc-idx rc-idx size))]))))))


;;;;;;;;;;;;;;;;;;;;;;;;
;;; Problem Solutions ;;
;;;;;;;;;;;;;;;;;;;;;;;;

;;; Procedure:
;;;   first
;;; Parameters:
;;;   heap, a heap
;;; Purpose:
;;;   Returns the first element of the priority heap
;;; Produces:
;;;   [nothing, called for side-effect]
(define (first heap)
  (vector-ref (vector-ref heap 1) 0))

;;; Procedure:
;;;   get
;;; Parameters:
;;;   heap, a heap
;;; Purpose:
;;;   Returns and removes the highest-priority element in the heap
;;; Produces:
;;;   [nothing, called for side-effect]
(define (get! heap)
  (let ([size (vector-ref heap 0)]
        [vec (vector-ref heap 1)]
        [higher-priority? (vector-ref heap 2)])
    (when (> size 0)
      (let ([root (vector-ref vec 0)])
        ; 'Discard' root by moving it to the end of the vector
        (vector-swap! vec 0 (- size 1))
        (vector-set! heap 0 (- size 1))
        (heap-percolate-down! heap)
        root))))

;;; Procedure:
;;;   add!
;;; Parameters:
;;;   heap, a heap created with make-heap
;;;   value, the value to be added to the heap
;;; Purpose:
;;;   Add a new element to the heap
;;; Produces:
;;;   [nothing, called for side-effect]
;;; Pre-conditions:
;;;   value can be compared with every element of the heap using
;;;     the higher-priority? predicate passed to make-heap
;;; Post-conditions:
;;;   heap has been mutated
;;;   heap contains value, and is in the correct position within the heap i.e.
;;;     the parent of the node containing value has a higher priority than value
;;;     and the children of that node have a lower priority
(define add!
  (lambda (heap value)
    (let ([size (vector-ref heap 0)]
          [vec (vector-ref heap 1)]
          [higher-priority? (vector-ref heap 2)])
      (vector-set! heap 1 (vector-insert! vec size value))
      (vector-set! heap 0 (+ 1 size))
      (heap-percolate-up! heap))))


;;;;;;;;;;;;;;
;;; Testing ;;
;;;;;;;;;;;;;;
(define h (heap-new >))
h
(add! h 3)
h
(add! h 1)
h
(add! h 5)
h
(add! h 8)
h
(add! h 12)
h
(add! h 9)
h
(get! h)
h
(get! h)
h
(add! h 100)
h
(get! h)
h
(get! h)
h
(get! h)
h
(get! h)
h
(get! h)
h