;;
;; 
;; Binary heap implementation.  Based on the Ocaml heap implementation
;; by Jean-Christophe Filliatre. Comments in the code are from the
;; original implementation.
;;
;; Copyright 2009-2011 Ivan Raikov.
;;
;;
;; 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 3 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.
;;
;; A full copy of the GPL license can be found at
;; <http://www.gnu.org/licenses/>.

(module binary-heap

  (make-binary-heap)

  (import scheme chicken data-structures srfi-1)

  (require-extension datatype matchable)


;;
;;  Heaps are encoded as binary trees that have the heap property,
;;  namely the value of any node is greater or equal than the nodes in
;;  its left and right subtrees.
;;
;;  The representation invariant is the following: the number of nodes
;;  in the left subtree is equal to the number of nodes in the right
;;  subtree, or exceeds it by exactly once. In the first case, we use
;;  the constructor [Same] and in the second the constructor [Diff].
;;  Then it can be proved that [2^(h-1) <= n <= 2^h] when [n] is the
;;  number of elements and [h] the height of the tree.
;;


(define-datatype tree tree?
  (Empty )
  (Same (left tree?) (elt pair?) (right tree?)) ;; same number of elements on both sides 
  (Diff (left tree?) (elt pair?) (right tree?)) ;; left has [n+1] nodes and right has [n] 
  )

(define (tree-tag x)
  (cases tree x
	 (Empty () 'Empty)
	 (Same (l x r) 'Same)
	 (Diff (l x r) 'Diff)))

(define-record-printer (tree x out)
  (cases tree x 
	 (Empty () (display "#(Empty)" out))
	 (Same (l x r)
	       (display "#(Same " out)
	       (display (tree-tag l) out)
	       (display (conc " " x " ") out)
	       (display (tree-tag r) out)
	       (display ")" out))
	 (Diff (l x r)
	       (display "#(Diff " out)
	       (display (tree-tag l) out)
	       (display (conc " " x " ") out)
	       (display (tree-tag r) out)
	       (display ")" out))))

;;
;; This macro was borrowed from treap.scm by Oleg Kiselyov
;;
(define-syntax dispatch-on-key 
  (lambda (x r c)
    (let ((key (second x)) (node-key (third x))
	  (on-greater (fourth x)) (on-less (fifth x)))
      (let ((%let   (r 'let))
	    (%cond  (r 'cond))
	    (%else  (r 'else))
	    (%positive?  (r 'positive?))
	    (result      (r 'result)))
	`(,%let ((,result (key-compare ,key ,node-key )))
		(,%cond
		 ((,%positive? ,result) ,on-greater)
		 (,%else                ,on-less)))))))

(define (make-binary-heap key-compare)

  (let ((root (Empty)) (size 0))

    (define (insert key value root) 
      (cases tree root
	     (Empty ()       (Same (Empty) (cons key value) (Empty)))
	     ;; insertion to the left
	     (Same (l y r)   (dispatch-on-key
			      key (car y)
			      (Diff (insert (car y) (cdr y) l) (cons key value) r)
			      (Diff (insert key value l) y r)))
	     ;; insertion to the right
	     (Diff (l y r)   (dispatch-on-key
			      key (car y)
			      (Same l (cons key value) (insert (car y) (cdr y) r))
			      (Same l y (insert key value r))))))
      
    (define (maximum root)
      (cases tree root
	     (Empty () #f)
	     (Same (l x r)  x)
	     (Diff (l x r)  x)))

    ;; extracts one element on the bottom level of the tree, while
    ;; maintaining the representation invariant
    (define (extract-last root)
      (match root
	     (($ tree 'Empty) #f)
	     (($ tree 'Same ($ tree 'Empty) x ($ tree 'Empty)) 
	      (list x (Empty)))
	     (($ tree 'Same l x r)  
	      (match-let (((y r1)  (extract-last r)))
			 (list y (Diff l x r1))))
	     (($ tree 'Diff l x r)  
	      (match-let (((y l1)  (extract-last l)))
			 (list y (Same l1 x r))))))
	     

    ;; removes the topmost element of the tree and inserts a new element 
    (define (descent key value root)
      (match root
	     (($ tree 'Empty)  #f)
	     (($ tree 'Same ($ tree 'Empty) _ ($ tree 'Empty))
	      (Same (Empty) (cons key value) (Empty)))
	     (($ tree 'Diff (and l ($ tree 'Same _ z _)) _ ($ tree 'Empty))
	      (dispatch-on-key key (car z)
			       (Diff l (cons key value) (Empty))
			       (Diff (Same (Empty ) (cons key value) (Empty)) z (Empty))))
	     (($ tree tag l _ r)
	      (let ((op (case tag ((Same) Same) ((Diff) Diff)))
		    (ml (maximum l)) (mr (maximum r)))
		(if (and (positive? (key-compare key (car ml)))
			 (positive? (key-compare key (car mr))))
		    (op l (cons key value) r)
		    (dispatch-on-key (car ml) (car mr)
				     (op (descent key value l) ml r) 
				     (op l mr (descent key value r))))))))

    (define (remove root)
      (match root
	     (($ tree 'Empty)  #f)
	     (($ tree 'Empty ($ tree 'Empty) x ($ tree 'Empty))
	      (Empty))
	     (else  (match-let (((y root1) (extract-last root)))
			       (descent (car y) (cdr y) root1)))))

    (define (heap-for-each f root)
      (match root
	     (($ tree 'Empty)  (begin))
	     (($ tree tag l x r)
	      (begin (heap-for-each f l)
		     (f x)
		     (heap-for-each f r)))))

    (define (heap-fold f init root)
      (define (foldf tree ax)
	(match tree
	       (($ tree 'Empty)  (begin))
	       (($ tree tag l x r)
		(foldf l (f x (foldf r ax))))))
      (foldf root (f init)))

    (define (make-heap-dispatcher root size)
      ;; Dispatcher
      (lambda (selector)
	(case selector

	  ((empty?)  (cases tree root
			    (Empty () #t)
			    (else #f)))

	  ((size)    size)

	  ((put)
	   (lambda (key value)
	     (let ((new-root  (insert key value root)))
 	       (make-heap-dispatcher new-root (+ 1 size)))))

	  ((get-max)     (lambda () (maximum root)))
	  
	  ((delete-max)  (lambda () (let ((new-root (remove root)))
				      (make-heap-dispatcher new-root (- size 1)))))
	  
	  ((fold)        (lambda (f init) (heap-fold f init root)))
	
	  ((for-each)    (lambda (f) (heap-for-each f root)))

	  (else          (error "unknown message " selector " sent to a binary heap"))
	  )))

  (make-heap-dispatcher root size)))

)