(in-package :trees)

(declaim (inline redp blackp redden blacken))
(defun redp (node) (eq (color node) :red))
(defun blackp (node) (eq (color node) :black))
(defun redden (node) (setf (color node) :red))
(defun blacken (node) (setf (color node) :black))

(defun red-black-rebalance/insert (tree direction-stack)
  (loop with stack = direction-stack
     for parent = (caar stack)
     for pp = (caadr stack)
     until (or (null parent) (null pp) (blackp parent))
     do (macrolet ((frob (ppfun opp r1 r2)
                     `(let ((y (,opp pp)))
                        (cond
                          ((and (not (null y)) (redp y))
                           (blacken parent)
                           (blacken y)
                           (redden pp)
                           (pop stack)
                           (pop stack))
                          (t
                           (let ((x pp))
                             (cond
                               ((eq (cdar stack) ',ppfun)
                                (setf y parent))
                               (t
                                (let ((x parent))
                                  (setf y (,opp x)
                                        (,ppfun pp) (,r1 x)))))
                             (redden x)
                             (blacken y)
                             (,r2 x)
                             (let ((ppp (caddr stack)))
                               (if (null ppp)
                                   (setf (root tree) y)
                                   (insert-child-for-stack-entry ppp y))
                               (loop-finish))))))))
          (if (eq (cdadr stack) 'left)
              (frob left right rotate-left rotate-right)
              (frob right left rotate-right rotate-left)))
     finally (progn
               (blacken (root tree))
               (return (values)))))

(defun red-black-rebalance/delete (tree node replacement new-stack)
  (unless (null (right node))
    (rotatef (color replacement) (color node)))
  (when (and (blackp node) new-stack)
    (loop while new-stack
       do
       (symbol-macrolet ((tos (car new-stack))
                         (tos-node (car tos))
                         (tos-dir (cdr tos))
                         (poptop (cadr new-stack))
                         (poptop-node (car poptop))
                         (poptop-dir (cdr poptop)))
         #+nil (format *trace-output* "processing ~A~%" tos)
         (let ((x (funcall tos-dir tos-node)))
           #+nil (format *trace-output* "z's ~A child is ~A~%" tos-dir x)
           (when (and x (redp x))
             (blacken x)
             (loop-finish))
           (macrolet ((frob (rf1 rf2 op1 op2)
                        `(let ((w (,op1 tos-node)))
                           #+nil(format *trace-output* "w is ~A~%" w)
                           (when (redp w)
                             (blacken w)
                             (redden tos-node)
                             (let ((r (,rf2 tos-node)))
                               #+nil (format *trace-output* "1 inserting ~A into ~A~%" r poptop)
                               (if (null (cdr new-stack))
                                   (setf (root tree) r)
                                   (insert-child-for-stack-entry poptop r)))
                             (push (cons tos-node tos-dir) new-stack)
                             (setf tos-dir ',op2
                                   poptop-node w)
                             (setf w (,op1 tos-node))
                             #+nil (format *trace-output* "w is now ~A~%" w)
                             (assert w)
                             (assert (blackp w)))
                           (cond
                             ((let ((s1 (,op1 w))
                                    (s2 (,op2 w)))
                                (and (or (null s1) (blackp s1))
                                     (or (null s2) (blackp s2))))
                              #+NIL (format *trace-output* "red'ing ~A~%" w)
                              (redden w))
                             (t
                              (when (let ((r (,op1 w)))
                                      (or (null r) (blackp r)))
                                #+nil (format *trace-output* "converting cases~%")
                                (assert (redp (,op2 w)))
                                (blacken (,op2 w))
                                (redden w)
                                (let ((y (,rf1 w)))
                                  (insert-child-for-stack-entry (cons tos-node ',op1) y)
                                  (assert tos-node)
                                  (setf w (,op1 tos-node))
                                  #+nil (format *trace-output* "w has become ~A~%" w)
                                  (assert (eq y w))))
                              #+NIL (format *trace-output* "handling case 3?~%")
                              (setf (color w) (color tos-node))
                              (blacken tos-node)
                              (blacken (,op1 w))
                              (let ((r (,rf2 tos-node)))
                                #+NIL (format *trace-output* "2 inserting ~A into ~A~%" r poptop)
                                (if (null (cdr new-stack))
                                    (setf (root tree) r)
                                    (insert-child-for-stack-entry poptop r))
                                (loop-finish)))))))
             (case tos-dir
               (left (frob rotate-right rotate-left right left))
               (right (frob rotate-left rotate-right left right)))
             (pop new-stack))))))
  (let ((root (root tree)))
    (when root
      (blacken root))))

(unless (assoc :red-black *binary-tree-info*)
  (push (list :red-black
              #'make-red-black-node
              #'red-black-rebalance/insert
              #'red-black-rebalance/delete)
        *binary-tree-info*))