;;;; java version @ http://pasted.co/12767a26
(defun compare (a b)
  "Return -1, 0, +1 as a < b, a = b, a > b, respectively."
  (cond ((< a b) -1)
        ((= a b)  0)
        (t       +1)))

(defconstant +red+ t)
(defconstant +black+ nil)

(defstruct node
  key
  val
  (lft nil)
  (rgt nil)
  (rob +red+)  ; color of parent link
  (num 1))     ; # of nodes in subtree rooted here

;;; t if +red+, nil if +black+ or tree is nil
(defun red? (tree)
  (and tree (node-rob tree)))

(defun black? (tree)
  (and tree (not (node-rob tree))))

(defun rb-size (tree)
  (if tree
      (node-num tree)
      0))

;;; update size _AND RETURN tree_
(defun update-size (tree)
  (setf (node-num tree) (+ (rb-size (node-lft tree))
                           (rb-size (node-rgt tree))
                           1))
  tree)

(defun rotate-right (tree)
  (assert (and tree (red? (node-lft tree))))

  (let ((x (node-lft tree)))
    (setf (node-lft tree) (node-rgt x))
    (setf (node-rgt x) tree)
    (setf (node-rob x) (node-rob (node-rgt x)))
    (setf (node-rob (node-rgt x)) +red+)
    (setf (node-num x) (node-num tree))
    (update-size tree)
    x))

(defun rotate-left (tree)
  (assert (and tree (red? (node-rgt tree))))

  (let ((x (node-rgt tree)))
    (setf (node-rgt tree) (node-lft x))
    (setf (node-lft x) tree)
    (setf (node-rob x) (node-rob (node-lft x)))
    (setf (node-rob (node-lft x)) +red+)
    (setf (node-num x) (node-num tree))
    (update-size tree)
    x))

(defun flip-colors (tree)
  ;;; debugging output -- i only ever see T T T when assert fails (only 3rd one ever fails)
  (format t "t = ~a, l = ~a, r = ~a~%" (node-rob tree) (node-rob (node-lft tree)) (node-rob (node-rgt tree)))

  (assert tree)
  (assert (and (node-lft tree) (node-rgt tree)))
  (assert (or (and (black? tree) (red? (node-lft tree)) (red? (node-rgt tree)))
              (and (red? tree) (black? (node-lft tree)) (black? (node-rgt tree)))))

  (setf (node-rob tree) (not (node-rob tree)))
  (setf (node-rob (node-lft tree)) (not (node-rob (node-lft tree))))
  (setf (node-rob (node-rgt tree)) (not (node-rob (node-rgt tree)))))

(defun rb-put (tree key val)
  (labels ((put (tree)
             (if (null tree)
                 (make-node :key key :val val)
                 (let ((cmp (compare key (node-key tree))))
                   (cond ((< cmp 0) (setf (node-lft tree) (put (node-lft tree))))
                         ((> cmp 0) (setf (node-rgt tree) (put (node-rgt tree))))
                         (t (setf (node-val tree) val)))
                   (when (and (red? (node-rgt tree)) (black? (node-lft tree)))
                     (setf tree (rotate-left tree)))
                   (when (and (red? (node-lft tree)) (red? (node-lft (node-lft tree))))
                     (setf tree (rotate-right tree)))
                   (when (and (red? (node-lft tree)) (red? (node-rgt tree)))
                     (flip-colors tree))
                   (update-size tree)))))
    (let ((tree (put tree)))
      (setf (node-rob tree) +black+)
      tree)))




;;; nope! last ASSERT in FLIP-COLORS will fail. Try (TEST-PUT 10) a few times.
(defun test-put (n)
  (let ((ht (make-hash-table :size n))
        (rb nil))
    ;; insert n distinct keys
    (loop for i below n do
         (let ((k (loop for k = (random 1000000000) while (gethash k ht)
                     finally (return k)))
               (v (random 1.0)))
           (setf (gethash k ht) v)
           (setf rb (rb-put rb k v))))
    (= n (rb-size rb))))

;;; HOWEVER:
;;; comment out the ASSERTs in FLIP-COLORS and everything seems to work:
;;; tree is built and all the keys are found, but the FORMAT statement in
;;; FLIP-COLORS is showing many t = T, l = T, r = T, which would be caught
;;; by the asserts
(defun rb-get (tree key)
  (labels ((iter (tree)
             (let ((cmp (compare key (node-key tree))))
               (cond ((< cmp 0) (iter (node-lft tree)))
                     ((> cmp 0) (iter (node-rgt tree)))
                     (t (node-val tree))))))
    (when tree
      (iter tree))))

(defun test-get (n)
  (let ((ht (make-hash-table :size n))
        (rb nil))
    (loop for i below n do
         (let ((k (loop for k = (random 1000000000) while (gethash k ht)
                     finally (return k)))
               (v (random 1000000000)))
           (setf (gethash k ht) v)
           (setf rb (rb-put rb k v))))
    (= n (loop for k being the hash-key using (hash-value v) of ht
            count (= v (rb-get rb k))))))