; avl trees

(define (tree k v l r) (vector k v l r (+ (max (ht l) (ht r)) 1)))
(define (key t) (vector-ref t 0))
(define (val t) (vector-ref t 1))
(define (lkid t) (vector-ref t 2))
(define (rkid t) (vector-ref t 3))
(define (ht t) (vector-ref t 4))
(define (bal t) (- (ht (lkid t)) (ht (rkid t))))
(define nil (vector 'nil 'nil 'nil 'nil 0))
(define (nil? t) (eq? t nil))

(define (rot-left t)
  (if (nil? t) t
    (tree (key (rkid t))
          (val (rkid t))
          (tree (key t) (val t) (lkid t) (lkid (rkid t)))
          (rkid (rkid t)))))

(define (rot-right t)
  (if (nil? t) t
    (tree (key (lkid t))
          (val (lkid t))
          (lkid (lkid t))
          (tree (key t) (val t) (rkid (lkid t)) (rkid t)))))

(define (balance t)
  (let ((b (bal t)))
    (cond ((and (= b 2) (= (bal (lkid t)) -1))
            (rot-right (tree (key t) (val t) (rot-left (lkid t)) (rkid t))))
          ((= b 2) (rot-right t))
          ((and (= b -2) (= (bal (rkid t)) 1))
            (rot-left (tree (key t) (val t) (lkid t) (rot-right (rkid t)))))
          ((= b -2) (rot-left t))
          (else t))))

(define (lookup lt? t k)
  (cond ((nil? t) #f)
        ((lt? k (key t)) (lookup lt? (lkid t) k))
        ((lt? (key t) k) (lookup lt? (rkid t) k))
        (else (cons k (val t)))))

(define (insert lt? t k v)
  (cond ((nil? t) (tree k v nil nil))
        ((lt? k (key t))
          (balance (tree (key t) (val t) (insert lt? (lkid t) k v) (rkid t))))
        ((lt? (key t) k)
          (balance (tree (key t) (val t) (lkid t) (insert lt? (rkid t) k v))))
        (else (tree k v (lkid t) (rkid t)))))

(define (delete lt? t k)
  (define (delete-successor t)
    (cond ((nil? (lkid t)) (values (rkid t) (key t) (val t)))
          (else (call-with-values
                  (lambda () (delete-successor (lkid t)))
                  (lambda (l k v)
                    (values (balance (tree (key t) (val t) l (rkid t))) k v))))))
  (let delete ((t t))
    (cond ((nil? t) nil)
          ((lt? k (key t))
            (balance (tree (key t) (val t) (delete (lkid t)) (rkid t))))
          ((lt? (key t) k)
            (balance (tree (key t) (val t) (lkid t) (delete (rkid t)))))
          ((nil? (lkid t)) (rkid t))
          ((nil? (rkid t)) (lkid t))
          (else (call-with-values
                  (lambda () (delete-successor (rkid t)))
                  (lambda (r k v) (balance (tree k v (lkid t) r))))))))

(define (enlist t)
  (cond ((nil? t) (list))
        ((and (nil? (lkid t)) (nil? (rkid t)))
          (list (cons (key t) (val t))))
        (else (append (enlist (lkid t))
                      (list (cons (key t) (val t)))
                      (enlist (rkid t))))))

(define t (insert < nil 4 4))
(set! t (insert < t 1 1))
(set! t (insert < t 3 3))
(set! t (insert < t 5 5))
(set! t (insert < t 2 2))
(display (enlist t)) (newline)
(display (lookup < t 3)) (newline)
(display (lookup < t 9)) (newline)
(set! t (delete < t 2))
(set! t (delete < t 5))
(set! t (delete < t 4))
(set! t (delete < t 1))
(set! t (delete < t 3))
(display (enlist t)) (newline)

((1 . 1) (2 . 2) (3 . 3) (4 . 4) (5 . 5))
(3 . 3)
#f
()