;;; -*-lisp-*-
;;;
;;; $Id$
;;;
;;; Heap data structure; useful for priority queues and suchlike
;;;
;;; (c) 2006 Straylight/Edgeware
;;;

;;;----- Licensing notice ---------------------------------------------------
;;;
;;; This program is free software; you can redistribute it and/or modify
;;; it under the terms of the GNU General Public License as published by
;;; the Free Software Foundation; either version 2 of the License, or
;;; (at your option) any later version.
;;;
;;; This program is distributed in the hope that it will be useful,
;;; but WITHOUT ANY WARRANTY; without even the implied warranty of
;;; MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
;;; GNU General Public License for more details.
;;;
;;; You should have received a copy of the GNU General Public License
;;; along with this program; if not, write to the Free Software Foundation,
;;; Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.

(defpackage #:heap
  (:use #:common-lisp)
  (:export #:make-heap #:heap-count #:heap-empty-p
	   #:heap-insert #:heap-head #:heap-remove
	   #:heap-sort))
(in-package #:heap)

;;;--------------------------------------------------------------------------
;;; Useful indexing functions.

(declaim (inline parent left-child right-child))
(deftype index () '(and unsigned-byte fixnum))
(defun parent (i)
  (declare (type index i))
  (the index (floor (- i 1) 2)))
(defun left-child (i)
  (declare (type index i))
  (the index (+ (* 2 i) 1)))
(defun right-child (i)
  (declare (type index i))
  (the index (+ (* 2 i) 2)))

;;;--------------------------------------------------------------------------
;;; Low-level heap operations.

(defun upheap (v key cmp n x)
  "Insert the element X in the highest place possible in the heap."
  (declare (type vector v)
	   (type function key cmp)
	   (type index n))
  (let ((i n) (xk (funcall key x)))
    (loop (when (zerop i) (return))
	  (let* ((j (parent i))
		 (y (aref v j)))
	    (when (funcall cmp (funcall key y) xk) (return))
	    (setf (aref v i) y
		  i j)))
    (setf (aref v i) x)))

(defun downheap (v key cmp n x)
  "Insert the element X in the lowest place possible in the heap."
  (declare (type vector v)
	   (type function key cmp)
	   (type index n))
  (let ((i 0) (xk (funcall key x)))
    (loop (let ((l (left-child i))
		(r (right-child i)))
	    (when (>= l n) (return))
	    (multiple-value-bind
		(j y yk)
		(let* ((y (aref v l))
		       (yk (funcall key y)))
		  (if (= r n)
		      (values l y yk)
		      (let* ((z (aref v r))
			     (zk (funcall key z)))
			(if (funcall cmp yk zk)
			    (values l y yk)
			    (values r z zk)))))
	      (when (funcall cmp xk yk)
		(return))
	      (setf (aref v i) y
		    i j))))
    (setf (aref v i) x)))

(defun check (v key cmp n)
  "Verify the heap invariant on the heap."
  (declare (type vector v)
	   (type function key cmp)
	   (type index n))
  (dotimes (i n)
    (let* ((item (aref v i))
	   (item-key (funcall key item))
	   (l (left-child i))
	   (r (right-child i)))
      (when (< l n)
	(let ((left-item (aref v l)))
	  (assert (funcall cmp item-key (funcall key left-item))))
	(when (< r n)
	  (let ((right-item (aref v r)))
	    (assert (funcall cmp item-key (funcall key right-item)))))))))

;;;--------------------------------------------------------------------------
;;; High-level heap things

(defstruct (heap (:predicate heapp) (:constructor %make-heap))
  "Data structure for a heap."
  (v (make-array 16) :type vector)
  (n 0 :type index)
  (key #'identity :type function)
  (compare #'<= :type function))

(defun make-heap
    (&key (compare #'<=) (key #'identity)
          (type 't) (init-size 16) (contents nil contentsp))
  "Return a new heap.

   COMPARE is a partial-order predicate: (COMPARE X Y) should return true if
   X <= Y in some order.

   The TYPE is the element type of the heap.

   INIT-SIZE is the initial allocation for the heap; the heap will grow
   automatically if necessary, so this isn't a big deal.  This is only a
   hint; make-heap may ignore it completely.

   KEY is a function to extract the key from an element.  The default is to
   use the item unmolested.

   CONTENTS is the initial contents of the heap.  If omitted, the heap is
   initially empty."
  (let ((n (if contentsp (length contents) 0)))
    (loop while (< init-size n)
	  do (setf init-size (ash init-size 1)))
    (let ((v (make-array init-size :element-type type)))
      (when contentsp
	(reduce (lambda (i item)
		  (upheap v key compare i item)
		  (1+ i))
		contents
		:initial-value 0))
      (%make-heap :compare compare :key key :n n :v v))))

(defun heap-count (heap)
  "Return the number of elements in HEAP."
  (declare (type heap heap))
  (heap-n heap))

(defun heap-empty-p (heap)
  "True if HEAP is empty."
  (declare (type heap heap))
  (zerop (heap-count heap)))

(defun heap-insert (heap item)
  "Insert ITEM into the HEAP."
  (declare (type heap heap))
  (let* ((v (heap-v heap))
	 (n (heap-n heap))
	 (sz (array-dimension v 0)))
    (when (= n sz)
      (setf v (adjust-array v (* 2 n))
	    (heap-v heap) v))
    (upheap v (heap-key heap) (heap-compare heap) n item)
    (setf (heap-n heap) (1+ n))))

(defun heap-head (heap)
  "Peep at the head item on HEAP."
  (declare (type heap heap))
  (assert (not (heap-empty-p heap)))
  (aref (heap-v heap) 0))

(defun heap-remove (heap)
  "Remove the head item from HEAP and return it."
  (declare (type heap heap))
  (assert (not (heap-empty-p heap)))
  (let ((v (heap-v heap))
	(n (1- (heap-n heap))))
    (prog1 (aref v 0)
      (setf (heap-n heap) n)
      (downheap v (heap-key heap) (heap-compare heap) n (aref v n)))))

(defun heap-sort (items compare &key (key #'identity))
  "Return the ITEMS, least-first, as sorted by the ordering COMPARE."
  (let ((heap (make-heap :compare compare :contents items :key key)))
    (loop repeat (heap-n heap)
	  collect (heap-remove heap))))

;;;----- That's all, folks --------------------------------------------------