; avl trees

(define (tree k v l r)
  (vector k v l r (+ (max (ht l) (ht r)) 1)
                  (+ (size l) (size 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 (size t) (vector-ref t 5))
(define (bal t) (- (ht (lkid t)) (ht (rkid t))))
(define nil (vector 'nil 'nil 'nil 'nil 0 0))
(vector-set! nil 2 nil)
(vector-set! nil 3 nil)
(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 ((< (abs b) 2) t)
          ((positive? b)
            (if (< -1 (bal (lkid t))) (rot-right t)
              (rot-right (tree (key t) (val t)
                (rot-left (lkid t)) (rkid t)))))
          ((negative? b)
            (if (< (bal (rkid t)) 1) (rot-left t)
              (rot-left (tree (key t) (val t)
                (lkid t) (rot-right (rkid 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-successor t)
  (if (nil? (lkid t)) (values (rkid t) (key t) (val t))
    (call-with-values
      (lambda () (delete-successor (lkid t)))
      (lambda (l k v)
        (values (balance (tree (key t) (val t) l (rkid t))) k v)))))

(define (delete lt? t k)
  (cond ((nil? t) nil)
        ((lt? k (key t))
          (balance (tree (key t) (val t)
            (delete lt? (lkid t) k) (rkid t))))
        ((lt? (key t) k)
          (balance (tree (key t) (val t)
            (lkid t) (delete lt? (rkid t) k))))
        ((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 (nth t n)
  (if (negative? n) (error 'nth "must be non-negative")
    (let ((s (size (lkid t))))
      (cond ((< n s) (nth (lkid t) n))
            ((< s n) (nth (rkid t) (- n s 1)))
            ((nil? t) #f)
            (else (cons (key t) (val t)))))))

(define (rank lt? t k)
  (let loop ((t t) (s (size (lkid t))))
    (cond ((nil? t) #f)
          ((lt? k (key t))
            (loop (lkid t) (size (lkid (lkid t)))))
          ((lt? (key t) k)
            (loop (rkid t) (+ s (size (lkid (rkid t))) 1)))
          (else s))))

(define (avl-map proc t) ; (proc key value)
  (if (nil? t) nil
    (tree (key t) (proc (key t) (val t))
          (avl-map proc (lkid t))
          (avl-map proc (rkid t)))))

(define (avl-fold proc base t) ; (proc key value base)
  (if (nil? t) base
    (avl-fold proc
              (proc (key t) (val t)
                    (avl-fold proc base (lkid t)))
              (rkid t))))

(define (avl-for-each proc t) ; (proc key value)
  (unless (nil? t)
    (avl-for-each proc (lkid t))
    (proc (key t) (val t))
    (avl-for-each proc (rkid t))))

(define (list->avl lt? xs)
  (let loop ((xs xs) (t nil))
    (if (null? xs) t
      (loop (cdr xs) (insert lt? t (caar xs) (cdar xs))))))

(avl-for-each (lambda (k v) (display k) (display " ") (display v) (newline))
  (avl-map (lambda (k v) (+ v v))
    (list->avl <
      (map (lambda (x) (cons x x))
        (list 0 1 2 3 4 5 6 7 8 9)))))

(display
  (avl-fold (lambda (k v base) (+ v base)) 0
    (list->avl <
      (map (lambda (x) (cons x x))
        (list 0 1 2 3 4 5 6 7 8 9)))))

0 0
1 2
2 4
3 6
4 8
5 10
6 12
7 14
8 16
9 18
45