(in-package :binary-trees)

(defmethod make-binary-tree-node ((tree red-black-tree) item)
  (let ((null-node (sentinel-node tree)))
    (make-instance 'red-black-tree-node
                   :left null-node
                   :right null-node
                   :parent null-node
                   :datum item)))

(defmethod make-binary-tree ((type (eql :red-black)) &key compfun eqfun keyfun)
  ;; colored according to CLRS
  (let* ((sentinel-node (make-sentinel-node 'red-black-tree-node
                                            '(:color :black))))
    (make-instance 'red-black-tree
                   :compfun compfun
                   :eqfun eqfun
                   :keyfun keyfun
                   :sentinel-node sentinel-node
                   :root-node sentinel-node)))

(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))

(defmethod insert-at-node ((tree red-black-tree) item parent direction-stack)
  ;; do insertion
  (let ((new-node (insert-node tree parent (first direction-stack) item)))
    ;; rebalance
    (do* ((parent parent (parent new-node))
          (pp (parent parent) (parent parent)))
         ((or (null-node-p parent tree)
              (null-node-p pp tree)
              (blackp parent))
          (when (null-node-p (parent new-node) tree)
            (setf (root-node tree) new-node))
          (blacken (root-node tree))
          (values))
      (macrolet ((frob (ppfun rotfun1 rotfun2)
                   `(if (and (not (null-node-p (,ppfun pp) tree))
                         (redp (,ppfun pp)))
                     (progn
                       (blacken parent)
                       (blacken (,ppfun pp))
                       (redden pp)
                       (setf new-node pp))
                     (progn
                       (when (eq new-node (,ppfun parent))
                         (setf new-node parent)
                         (,rotfun1 tree new-node)
                         (setf parent (parent new-node)))
                       (blacken parent)
                       (unless (null-node-p (setf pp (parent parent)) tree)
                         (redden pp)
                         (,rotfun2 tree pp))))))
        (if (eq (left pp) parent)
            (frob right rotate-left rotate-right)
            (frob left rotate-right rotate-left))))))

(defun red-black-delete-fixup (tree node)
  (with-slots (root-node) tree
    (loop while (and (not (eq root-node node))
                     (blackp node))
      do
      (macrolet ((frob (rf1 rf2 op1 op2)
                   `(let ((w (,op1 (parent node))))
                     (when (redp w)
                       (blacken w)
                       (redden (parent node))
                       (,rf2 tree (parent node))
                       (setf w (,op1 (parent node))))
                     (if (and (blackp (,op2 w)) (blackp (,op1 w)))
                         (progn
                           (redden w)
                           (setf node (parent node)))
                         (progn
                           (when (blackp (,op1 w))
                             (blacken (,op2 w))
                             (redden w)
                             (,rf1 tree w)
                             (setf w (,op1 (parent node))))
                           (setf (color w) (color (parent node)))
                           (blacken (parent node))
                           (blacken (,op1 w))
                           (,rf2 tree (parent node))
                           (setf node root-node))))))
        (cond
          ((eq node (left (parent node)))
           (frob rotate-right rotate-left right left))
          (t
           (frob rotate-left rotate-right left right)))))
    (blacken node)))
                   
(defmethod tree-delete-nonempty ((tree red-black-tree) deleted child low-subtree)
  (when (blackp deleted)
    (red-black-delete-fixup tree child))
  t)