;;; This file and the accompanying README were derived from
;;; Oleg's code for Gambit available from
;;;
;;;        http://okmij.org/ftp/Scheme/lib/treap.scm
;;;

(define-record-type treap
  (really-make-treap key-compare size root)
  treap?
  (key-compare treap-key-compare)
  (size treap-size set-treap-size!)
  (root treap-root set-treap-root!))

(define (make-treap key-compare)
  (really-make-treap key-compare 0 #f))

;; a node of a tree, a record of
;;   key, anything that key-compare could be applied to
;;   value, any object associated with the key
;;   left-kid, #f if absent
;;   right-kid
;;   prio, a priority of the node (a FIXNUM random number)


(define-record-type node
  (make-node key value left-kid right-kid priority)
  node?
  (key node:key)
  (value node:value node:value-set!)
  (left-kid node:left-kid node:left-kid-set!)
  (right-kid node:right-kid node:right-kid-set!)
  (priority node:priority))

(define random
  (let ((max (- (expt 2 15) 1)))
    (lambda ()
      (random-integer max))))

(define (new-leaf key value)
  (make-node key value #f #f (random)))

(define (node:key-value node)
  (cons (node:key node)
        (node:value node)))

(define (node:unsubordination? parent kid)
  (> (node:priority parent) (node:priority kid)))

(define-syntax node:dispatch-on-key
  (syntax-rules ()
    ((node:dispatch-on-key treap node key on-less on-equal on-greater)
     (let ((result ((treap-key-compare treap) key (node:key node))))
       (cond
        ((zero? result) on-equal)
        ((positive? result) on-greater)
        (else on-less))))))

(define (n-display . args)
  (for-each display args))

(define (node:debugprint node)
  (n-display " " (node:key-value node) ", kids "
             (cons (not (not (node:left-kid node)))
                   (not (not (node:right-kid node))))
             ", prio " (node:priority node) #\newline))




;; Looking up assocaitions in a treap: just like in any search tree
;; Given  a key, return the corresponding (key . value) association
;; in the treap, or #f if the treap does not contain an association
;; with that key
;; This procedure takes as many comparisons (evaluations of the
;; key-compare procedure) as the depth of the found node
(define (locate-assoc treap key)
  (let loop ((node (treap-root treap)))
    (and node
         (node:dispatch-on-key treap node key
                               (loop (node:left-kid node))
                               (node:key-value node)
                               (loop (node:right-kid node))))))

(define (locate-extremum-node treap branch-selector)
  (let ((root (treap-root treap)))
    (if (not root) (error "empty tree")
        (let loop ((node root) (parent #f))
          (if node (loop (branch-selector node) node)
              (node:key-value parent))))))

                                        ; in-order traversal of the treap
(define (for-each-inorder treap primary-branch-selector secondary-branch-selector)
  (let ((root (treap-root treap)))
    (lambda (proc)
      (if (not root) (error "empty tree")
          (let loop ((node root))
            (if node
                (begin
                  (loop (primary-branch-selector node))
                  (proc (node:key-value node))
                  (loop (secondary-branch-selector node)))))))))

(define (get-depth treap)
  (let ((root (treap-root treap)))
    (let loop ((node root) (level 0))
      (if (not node) level
          (max (loop (node:left-kid node) (+ 1 level))
               (loop (node:right-kid node) (+ 1 level)))))))

;; debug printing of all nodes of the tree in-order
;; in an ascending order of keys
(define (debugprint treap)
  (let ((root (treap-root treap)))
    (n-display  #\newline
                "The treap contains " (treap-size treap) " nodes"
                #\newline)
    (let loop ((node root) (level 0))
      (if node
          (begin
            (loop (node:left-kid node) (+ 1 level))
            (n-display "  level " level)
            (node:debugprint node)
            (loop (node:right-kid node) (+ 1 level))))
      (newline))))

;; Adding a new association to the treap (or replacing the old one
;; if existed). Return the (key . value) pair of an old (existed
;; and replaced association), or #f if a new association was really
;; added
(define (insert! treap key value)
  (let ((root (treap-root treap)))
    (letrec ((new-node (new-leaf key value))
             (old-key-value #f)
             ;; If the left branch of parent is empty, insert the
             ;; new node there, check priorities
             ;; Otherwise, descend recursively
             ;; If the parent got inverted due to a right rotation,
             ;; return the new parent of the branch; otherwise,
             ;; return #f (indicating no further checks are necessary)
             (insert-into-left-branch
              (lambda (key parent)
                (let ((old-left-kid (node:left-kid parent)))
                  ;; Found a place to insert the 'new-node': as the left
                  ;; leaf of the parent
                  (if (not old-left-kid)
                      (cond
                       ((node:unsubordination? parent new-node)
                        ;; Right rotation over the new-leaf
                        (node:right-kid-set! new-node parent)
                        new-node)	;; becomes a new parent
                       (else
                        (node:left-kid-set! parent new-node)
                        #f))
                      ;; Insert the new-leaf into a branch rooted
                      ;; on old-left-kid
                      (let ((new-left-kid
                             (node:dispatch-on-key treap old-left-kid key
                                                   (insert-into-left-branch key old-left-kid)
                                                   (update-existing-node old-left-kid)
                                                   (insert-into-right-branch key old-left-kid))))
                        (and new-left-kid
                             ;; That branch got a new root
                             (cond
                              ((node:unsubordination? parent new-left-kid)
                               ;; Right rotation over the new-left-kid
                               (node:left-kid-set! parent
                                                   (node:right-kid new-left-kid))
                               (node:right-kid-set! new-left-kid parent)
                               new-left-kid) ;; becomes a new parent
                              (else
                               (node:left-kid-set! parent new-left-kid)
                               #f))))
                      ))))

             ;; If the right branch of parent is empty, insert the
             ;; new node there, check priorities
                                        ; Otherwise, descend recursively
             ;; If the parent got inverted due to a left rotation,
             ;; return the new parent of the branch; otherwise,
             ;; return #f (indicating no further checks are necessary)
             (insert-into-right-branch
              (lambda (key parent)
                (let ((old-right-kid (node:right-kid parent)))
                  ;; Found a place to insert the 'new-node': as the right
                  ;; leaf of the parent
                  (if (not old-right-kid)
                      (cond
                       ((node:unsubordination? parent new-node)
                        ;; Left rotation over the new-leaf
                        (node:left-kid-set! new-node parent)
                        new-node)	; becomes a new parent
                       (else
                        (node:right-kid-set! parent new-node)
                        #f))
                      ;; Insert the new-leaf into a branch rooted
                      ;; on old-right-kid
                      (let ((new-right-kid
                             (node:dispatch-on-key treap old-right-kid key
                                                   (insert-into-left-branch key old-right-kid)
                                                   (update-existing-node old-right-kid)
                                                   (insert-into-right-branch key old-right-kid))))
                        (and new-right-kid
                             ;; That branch got a new root
                             (cond
                              ((node:unsubordination? parent new-right-kid)
                               ;; Left rotation over the new-right-kid
                               (node:right-kid-set! parent
                                                    (node:left-kid new-right-kid))
                               (node:left-kid-set! new-right-kid parent)
                               new-right-kid) ; becomes a new parent
                              (else
                               (node:right-kid-set! parent new-right-kid)
                               #f))))
                      ))))

             (update-existing-node
              (lambda (node)
                (set! old-key-value (node:key-value node))
                (node:value-set! node value)
                #f))
             )                          ; end of letrec

      ;; insert's body
      (cond
       ;; insert into an empty tree
       ((not root) (set-treap-root! treap new-node))

       (else
        (let ((new-root
               (node:dispatch-on-key treap root key
                                     (insert-into-left-branch key root)
                                     (update-existing-node root)
                                     (insert-into-right-branch key root))))
          (if new-root
              (set-treap-root! treap  new-root)))))
      (if (not old-key-value)
          (set-treap-size! treap (+ (treap-size treap) 1))) ; if the insertion has really occurred
      old-key-value)))


;; Deleting existing associations from the treap

(define (delete-extremum-node! treap branch-selector
                               branch-setter the-other-branch-selector)
  (let ((root (treap-root treap)))
    (cond
     ((not root) (error "empty tree"))
     ((not (branch-selector root))	; root is the extreme node
      (let ((result (node:key-value root)))
        (set-treap-root! treap (the-other-branch-selector root))
        (set-treap-size! treap (- (treap-size treap) 1))
        result))
     (else
      (let loop ((node (branch-selector root)) (parent root))
        (let ((kid (branch-selector node)))
          (if kid (loop kid node)
              (let ((result (node:key-value node)))
                (branch-setter parent (the-other-branch-selector node))
                (set-treap-size! treap (- (treap-size treap) 1))
                result))))))))

;; Given two treap branches (both of which could be empty)
;; which satisfy both the order invariant and the priority invariant
;; (all keys of all the nodes in the right branch are strictly bigger
;; than the keys of left branch nodes), join them
;; while keeping the sorted and priority orders intact
(define (join! treap left-branch right-branch)
  (cond
   ((not left-branch) right-branch)	; left-branch was empty
   ((not right-branch) left-branch)	; right-branch was empty
   ((node:unsubordination? left-branch right-branch)
    ;; the root of the right-branch should be the new root
    (node:left-kid-set! right-branch
                        (join! treap left-branch (node:left-kid right-branch)))
    right-branch)
   (else
    ;; the root of the left-branch should be the new root
    (node:right-kid-set! left-branch
                         (join! treap (node:right-kid left-branch) right-branch))
    left-branch)))


;; Find an association with a given KEY, and delete it.
;; Return the (key . value) pair of the deleted association, or
;; #f if it couldn't be found
(define (delete! treap key)
  (define (delete-node! node parent from-left?)
    (let ((old-assoc (node:key-value node))
          (new-kid (join! treap (node:left-kid node) (node:right-kid node))))
      (set-treap-size! treap (- (treap-size treap) 1))
      (if parent
          (if from-left?
              (node:left-kid-set! parent new-kid)
              (node:right-kid-set! parent new-kid))
          ;; Deleting of the root node
          (set-treap-root! treap new-kid))
      old-assoc))

  (let loop ((node (treap-root treap)) (parent #f) (from-left? #t))
    (and node
         (node:dispatch-on-key treap node key
                               (loop (node:left-kid node) node #t)
                               (delete-node! node parent from-left?)
                               (loop (node:right-kid node) node #f)))))

(define (apply-default-clause key default-clause)
  (cond
   ((null? default-clause)
    (error "key " key " was not found in the treap "))
   ((pair? (cdr default-clause))
    (error "default argument must be a single clause"))
   ((procedure? (car default-clause)) ((car default-clause)))
   (else (car default-clause))))

(define (treap-get treap key . default-clause)
  (or (locate-assoc treap key) (apply-default-clause key default-clause)))

(define (treap-delete! treap key . default-clause)
  (or (delete! treap key) (apply-default-clause key default-clause)))

(define (treap-get-min treap)
  (locate-extremum-node treap node:left-kid))

(define (treap-get-max treap)
  (locate-extremum-node treap node:right-kid))

(define (treap-delete-min! treap)
  (delete-extremum-node! treap
                         node:left-kid node:left-kid-set!
                         node:right-kid))

(define (treap-delete-max! treap)
  (delete-extremum-node! treap
                         node:right-kid node:right-kid-set!
                         node:left-kid))

(define (treap-empty? treap)
  (not (treap-root treap)))

(define (treap-depth treap)
  (get-depth treap))

(define (treap-clear! treap)
  (set-treap-root! treap #f)
  (set-treap-size! treap 0))

(define treap-put! insert!)

(define (treap-for-each-ascending treap proc)
  ((for-each-inorder treap node:left-kid node:right-kid) proc))

(define (treap-for-each-descending treap proc)
  ((for-each-inorder treap node:right-kid node:left-kid) proc))

(define treap-debugprint debugprint)