;;; -*-lisp-*-
;;;
;;; $Id$
;;;
;;; Simple binary heaps, for priority queues
;;;
;;; (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 #:mdw.base)
  (:export #:heap #:heapp
	   #:heap-key< #:heap-key
	   #:heap-check
	   #:heap-count #:heap-empty-p
	   #:heap-error #:heap-underflow-error
	   #:heap-add #:heap-min #:heap-remove))
(in-package #:heap)

(defstruct (heap
	     (:predicate heapp)
	     (:constructor make-heap
			   (key<-name
			    &optional
			    (key-name 'identity)
			    (max 32)
			    &aux
			    (key< (functionify key<-name))
			    (key (functionify key-name))
			    (store (make-array max :adjustable t)))))
  "Structure representing a heap."
  (key< (uninitialized-slot) :read-only t :type (function (t t) t))
  (key #'identity :read-only t :type (function (t) t))
  (store (uninitialized-slot) :type vector)
  (size 0 :type unsigned-fixnum))

(defmacro with-heap
    ((heap &key (key< 'key<) (key 'key) (store 'store) (size 'size))
     &body body)
  "Pull aspects of the HEAP into local variables and functions so that they
   can be messed with conveniently."
  (let*/gensyms (heap)
    (with-gensyms (cmp getkey x y)
      `(let ((,getkey (heap-key ,heap))
	     (,cmp (heap-key< ,heap))
	     (,store (heap-store ,heap))
	     (,size (heap-size ,heap)))
	(declare (ignorable ,size)
	         (type (function (t t) t) ,cmp)
	         (type (function (t) t) ,getkey))
	(flet ((,key< (,x ,y)
		 (funcall ,cmp ,x ,y))
	       (,key (,x)
		 (funcall ,getkey ,x)))
	  (declare (inline ,key< ,key))
	  ,@body)))))

(declaim (inline heap-left heap-right heap-parent))
(defun heap-left (i)
  "Compute the left child index of the item with index I."
  (declare (type unsigned-fixnum i))
  (the unsigned-fixnum (+ (* 2 i) 1)))
(defun heap-right (i)
  "Compute the right child index of the item with index I."
  (declare (type unsigned-fixnum i))
  (the unsigned-fixnum (+ (* 2 i) 2)))
(defun heap-parent (i)
  "Compute the parent index of the item with index I.  Do not use this on
   index zero."
  (declare (type unsigned-fixnum i))
  (the unsigned-fixnum (ash (1- i) -1)))

(defun heap-check (heap)
  "Verifies the heap property on the heap HEAP."
  (declare (type heap heap))
  (with-heap (heap)
    (labels ((check (i key)
	       (let ((left (heap-left i))
		     (right (heap-right i)))
		 (when (< left size)
		   (let ((lkey (key (aref store left))))
		     (assert (not (key< lkey key)))
		     (check left lkey))
		   (when (< right size)
		     (let ((rkey (key (aref store right))))
		       (assert (not (key< rkey key)))
		       (check right rkey)))))))
      (when (plusp size)
	(check 0 (key (aref store 0)))))))

(defun heap-count (heap)
  "Returns a count of the number of items in the HEAP."
  (declare (type heap heap))
  (heap-size heap))

(defun heap-empty-p (heap)
  "Returns true if the HEAP is currently empty, or nil if it contains at
   least one item.  Note that heap-min and heap-remove are invalid operations
   on an empty heap."
  (declare (type heap heap))
  (zerop (heap-size heap)))

(define-condition heap-error (error)
  ((heap :initarg :heap :type heap :reader heap-error-heap))
  (:documentation
   "Parent class for error conditions affecting a heap.  The accessor
    heap-error-heap will extract the offending heap."))

(define-condition heap-underflow-error (heap-error)
  ()
  (:report "Heap underflow.")
  (:documentation
   "Reports a heap underflow: i.e., an attempt to do something which requires
    a nonempty heap to an empty one."))

(defun upheap (heap i item)
  "Fixes up the HEAP after a (virtual) attempt to add the new ITEM at index
   I."
  (declare (type heap heap))
  (with-heap (heap)
    (let ((key (key item)))
      (loop (cond ((zerop i)
		   (return))
		  (t
		   (let* ((parent (heap-parent i))
			  (pitem (aref store parent)))
		     (when (key< (key pitem) key)
		       (return))
		     (setf (aref store i) pitem
			   i parent)))))
      (setf (aref store i) item))))

(defun downheap (heap item)
  "Fixes up the HEAP after a removal of the head element.  The ITEM is the
   element at the end of the heap.  Does something very bad if the heap is
   empty." 
  (declare (type heap heap))
  (with-heap (heap)
    (let* ((i 0)
	   (key (key item)))
      (loop (let ((left (heap-left i))
		  (right (heap-right i)))
	      (cond ((>= left size)
		     (return))
		    (t
		     (multiple-value-bind
			 (child citem ckey)
			 (let* ((litem (aref store left))
				(lkey (key litem)))
			   (if (>= right size)
			       (values left litem lkey)
			       (let* ((ritem (aref store right))
				      (rkey (key ritem)))
			       (if (key< lkey rkey)
				   (values left litem lkey)
				   (values right ritem rkey)))))
		       (when (key< key ckey)
			 (return))
		       (setf (aref store i) citem
			     i child))))))
      (setf (aref store i) item))))

(defun heap-add (heap item)
  "Add the given ITEM to the HEAP."
  (declare (type heap heap))
  (let ((i (heap-size heap))
	(store (heap-store heap)))
    (when (>= i (array-dimension store 0))
      (setf (heap-store heap) (adjust-array store (ash i 1))))
    (setf (heap-size heap) (1+ i))
    (upheap heap i item)))

(defun heap-min (heap)
  "Returns the smallest item in the HEAP."
  (declare (type heap heap))
  (when (zerop (heap-size heap))
    (error 'heap-underflow-error :heap heap))
  (aref (heap-store heap) 0))

(defun heap-remove (heap)
  "Removes the smallest item in the HEAP, returning something which isn't
   useful." 
  (declare (type heap heap))
  (let ((i (1- (heap-size heap)))
	(store (heap-store heap)))
    (cond ((minusp i)
	   (error 'heap-underflow-error :heap heap))
	  ((zerop i)
	   (setf (heap-size heap) 0
		 (aref store 0) nil))
	  (t
	   (let ((item (aref store i)))
	     (setf (aref store i) nil
		   (heap-size heap) i)
	     (downheap heap item))))))

(defun test ()
  (let ((heap (make-heap #'<))
	(list (cons nil nil)))
    (flet ((add ()
	     (let ((n (random 4096)))
	       (heap-add heap n)
	       (loop for c on list
		     when (or (null (cdr c))
			      (< n (cadr c)))
		     do (setf (cdr c) (cons n (cdr c))) (return))))
	   (remove ()
	     ;;(show (values (cdr list) heap))
	     (cond ((cdr list)
		    (let ((n (heap-min heap)))
		      (assert (= n (pop (cdr list))))
		      (heap-remove heap)))
		   (t
		    (handler-case
			(progn (heap-remove heap) (error "Bummer"))
		      (heap-underflow-error () nil))))))
      (dotimes (i 1024)
	(add))
      (dotimes (i 65536)
	(case (random 2)
	  (0 (add))
	  (1 (remove))))
      (loop for i in (cdr list)
	    for j = (heap-min heap)
	    do (assert (= i j)) (heap-remove heap))
      (assert (heap-empty-p heap)))))

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