;;;; This file is part of P36-lib
;;;; Copyright (C) 2016-2021 Alexa Jones-Gonzales <alexa@partition36.com>
;;;;
;;;; This program is free software: you can redistribute it and/or modify
;;;; it under the terms of the GNU General Public License as published by
;;;; the Free Software Foundation, either version 3 of the License, or
;;;; (at your option) any later version.
;;;;
;;;; This program is distributed in the hope that it will be useful,
;;;; but WITHOUT ANY WARRANTY; without even the implied warranty of
;;;; MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
;;;; GNU General Public License for more details.
;;;;
;;;; You should have received a copy of the GNU General Public License
;;;; along with this program. If not, see <http://www.gnu.org/licenses/>.
(in-package :p36-lib/data-structures)
(deftype t/rb-node-color ()
'(member :red :black))
(defclass rb-node/sentinel (base-node)
((parent
:initform nil
:accessor node-parent)
(left
:initform nil
:accessor node-left)
(right
:initform nil
:accessor node-right)
(color
:initform :black
:type t/rb-node-color
:reader node-color)))
(defclass rb-node (base-node)
((parent
:initform (make-instance 'rb-node/sentinel)
:type (or rb-node rb-node/sentinel)
:accessor node-parent)
(left
:initform (make-instance 'rb-node/sentinel)
:type (or rb-node rb-node/sentinel)
:accessor node-left)
(right
:initform (make-instance 'rb-node/sentinel)
:type (or rb-node rb-node/sentinel)
:accessor node-right)
(color
:initarg :color
:initform :red
:type t/rb-node-color
:accessor node-color)
(data
:initarg :data
:initform nil
:accessor node-data)
(parent-tree
:initarg :parent-tree
:initform nil
:type (or null rb-tree)
:accessor node-parent-tree)))
(defclass rb-tree (base-tree)
((root
:initform nil
:type (or null rb-node rb-node/sentinel)
:accessor tree-root)))
(defparameter *rb-node-sentinel* (make-instance 'rb-node/sentinel))
(defmethod print-object ((obj rb-tree) out)
(print-unreadable-object (obj out :type t)
(format out "Count: ~:d" (tree-count obj))))
(defmethod print-object ((obj rb-node) out)
(print-unreadable-object (obj out :type t)
(format out "Color: ~a, Data: ~a"
(node-color obj) (node-data obj))))
(defmethod initialize-instance :after ((tree rb-tree) &key &allow-other-keys)
(unless (slot-value tree 'test-fn)
(error "Tree has no test function"))
(unless (slot-value tree 'key-fn)
(error "Tree has no key function")))
(defmethod initialize-instance :after ((node rb-node/sentinel) &key &allow-other-keys)
(setf (node-parent node) node))
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;;;
;;; Basic Node Stuff
;;;
(defmethod node-grandparent ((node rb-node))
(declare (optimize speed))
(node-parent (node-parent node)))
(defmethod node-sibling ((node rb-node))
(declare (optimize speed))
(let ((p (node-parent node)))
(when (typep p 'rb-node/sentinel)
(return-from node-sibling *rb-node-sentinel*))
(if (eq node (node-left p))
(node-right p)
(node-left p))))
(defmethod node-uncle ((node rb-node))
(declare (optimize speed))
(node-sibling (node-parent node)))
(defmethod node-rotate-left ((node rb-node))
(declare (optimize speed))
(let ((new (node-right node))
(parent (node-parent node)))
(when (eq new *rb-node-sentinel*)
(error "Cannot rotate left, right node is a sentinel"))
(setf (node-right node) (node-left new))
(setf (node-left new) node)
(setf (node-parent node) new)
(unless (eq (node-right node) *rb-node-sentinel*)
(setf (node-parent (node-right node)) node))
(unless (typep parent 'rb-node/sentinel)
(if (eq node (node-left parent))
(setf (node-left parent) new)
(setf (node-right parent) new)))
(setf (node-parent new) parent)))
(defmethod node-rotate-right ((node rb-node))
(declare (optimize speed))
(let ((new (node-left node))
(parent (node-parent node)))
(when (eq new *rb-node-sentinel*)
(error "Cannot rotate left, right node is a sentinel"))
(setf (node-left node) (node-right new))
(setf (node-right new) node)
(setf (node-parent node) new)
(unless (eq (node-left node) *rb-node-sentinel*)
(setf (node-parent (node-left node)) node))
(unless (typep parent 'rb-node/sentinel)
(if (eq node (node-left parent))
(setf (node-left parent) new)
(setf (node-right parent) new)))
(setf (node-parent new) parent)))
(defmacro rb-node-test (node1 node2)
`(funcall (the function (slot-value (slot-value ,node1 'parent-tree) 'test-fn))
(funcall (the function (slot-value (slot-value ,node1 'parent-tree) 'key-fn)) (node-data ,node1))
(funcall (the function (slot-value (slot-value ,node2 'parent-tree) 'key-fn)) (node-data ,node2))))
(defmacro rb-node-test* (data node fn)
`(funcall (the function ,fn)
,data
(funcall (the function (slot-value (slot-value ,node 'parent-tree) 'key-fn)) (node-data ,node))))
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;;;
;;; Node Insertion
;;;
(defmethod node-insert ((root (eql nil)) (node rb-node))
(%node-insert root node)
(%node-repair node)
(setf root node)
(setf (node-parent-tree node) (node-parent-tree root))
(loop until (typep (node-parent root) 'rb-node/sentinel) do
(setf root (node-parent root)))
root)
(defmethod node-insert ((root rb-node) (node rb-node))
(%node-insert root node)
(%node-repair node)
(setf root node)
(setf (node-parent-tree node) (node-parent-tree root))
(loop until (typep (node-parent root) 'rb-node/sentinel) do
(setf root (node-parent root)))
root)
(defmethod %node-insert ((root (eql nil)) (node rb-node))
(setf (node-parent node) (make-instance 'rb-node/sentinel))
(setf (node-left node) *rb-node-sentinel*)
(setf (node-right node) *rb-node-sentinel*)
(setf (node-color node) :red)
(incf (the fixnum (slot-value (node-parent-tree node) 'count)))
nil)
(defmethod %node-insert ((root rb-node/sentinel) (node rb-node))
(setf (node-parent node) root)
(setf (node-left node) *rb-node-sentinel*)
(setf (node-right node) *rb-node-sentinel*)
(setf (node-color node) :red)
(incf (the fixnum (slot-value (node-parent-tree node) 'count)))
nil)
(defmethod %node-insert ((root rb-node) (node rb-node))
(declare (optimize (speed 3))) ;; For tail-call optimization
(cond
((rb-node-test node root)
(cond
((not (eq (node-left root) *rb-node-sentinel*))
(return-from %node-insert (%node-insert (node-left root) node)))
(t
(setf (node-left root) node))))
;; node >= root
(t
(cond
((not (eq (node-right root) *rb-node-sentinel*))
(return-from %node-insert (%node-insert (node-right root) node)))
(t
(setf (node-right root) node)))))
(setf (node-parent node) root)
(setf (node-left node) *rb-node-sentinel*)
(setf (node-right node) *rb-node-sentinel*)
(setf (node-color node) :red)
(incf (the fixnum (slot-value (node-parent-tree root) 'count)))
nil)
(defmethod %node-repair ((node rb-node))
(declare (optimize (speed 3))) ;; For tail-call optimization
(cond
((typep (node-parent node) 'rb-node/sentinel)
(setf (node-color node) :black))
((eql (node-color (node-parent node)) :black)
nil)
((and (not (eq (node-uncle node) *rb-node-sentinel*))
(eql (node-color (node-uncle node)) :red))
(setf (node-color (node-parent node)) :black)
(setf (node-color (node-uncle node)) :black)
(setf (node-color (node-grandparent node)) :red)
(%node-repair (node-grandparent node)))
(t
(let ((parent (node-parent node))
(grandma (node-grandparent node)))
(cond
((and (eq node (node-right parent))
(eq parent (node-left grandma)))
(node-rotate-left parent)
(setf node (node-left node)))
((and (eq node (node-left parent))
(eq parent (node-right grandma)))
(node-rotate-right parent)
(setf node (node-right node))))
(setf parent (node-parent node))
(setf grandma (node-grandparent node))
(if (eq node (node-left parent))
(node-rotate-right grandma)
(node-rotate-left grandma))
(setf (node-color parent) :black)
(setf (node-color grandma) :red)))))
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;;;
;;; Node Deletion
;;;
(defmethod node-replace ((node rb-node) (child rb-node))
(setf (node-parent child) (node-parent node))
(if (eq node (node-left (node-parent node)))
(setf (node-left (node-parent node)) child)
(setf (node-right (node-parent node)) child)))
;;
;; Deletion gets kinda hairy, so split up the cases into various
;; macros.
;;
(declaim (type function node-delete/case-1))
(defun node-delete/case-6 (node)
(declare (optimize (speed 3)))
(let ((sib (node-sibling node)))
(setf (node-color sib) (node-color (node-parent node)))
(setf (node-color (node-parent node)) :black)
(cond
((eq node (node-left (node-parent node)))
(setf (node-color (node-right sib)) :black)
(node-rotate-left (node-parent node)))
(t
(setf (node-color (node-left sib)) :black)
(node-rotate-right (node-parent node))))))
(defun node-delete/case-5 (node)
(declare (optimize (speed 3)))
(let ((sib (node-sibling node)))
(when (eql (node-color sib) :black)
(cond
((and (eq node (node-left (node-parent node)))
(eql (node-color (node-right sib)) :black)
(eql (node-color (node-left sib)) :red))
(setf (node-color sib) :red)
(setf (node-color (node-left sib)) :black)
(node-rotate-right sib))
((and (eq node (node-right (node-parent node)))
(eql (node-color (node-left sib)) :black)
(eql (node-color (node-right sib)) :red))
(setf (node-color sib) :red)
(setf (node-color (node-right sib)) :black)
(node-rotate-left sib)))))
(node-delete/case-6 node))
(defun node-delete/case-4 (node)
(declare (optimize (speed 3)))
(let ((sib (node-sibling node)))
(cond
((and (eql (node-color (node-parent node)) :red)
(eql (node-color sib) :black)
(eql (node-color (node-left sib)) :black)
(eql (node-color (node-right sib)) :black))
(setf (node-color sib) :red)
(setf (node-color (node-parent node)) :black))
(t (node-delete/case-5 node)))))
(defun node-delete/case-3 (node)
(declare (optimize (speed 3)))
(let ((sib (node-sibling node)))
(cond
((and (eql (node-color (node-parent node)) :black)
(eql (node-color sib) :black)
(eql (node-color (node-left sib)) :black)
(eql (node-color (node-right sib)) :black))
(setf (node-color sib) :red)
(node-delete/case-1 (node-parent node)))
(t (node-delete/case-4 node)))))
(defun node-delete/case-2 (node)
(declare (optimize (speed 3)))
(let ((sib (node-sibling node)))
(when (eql (node-color sib) :red)
(setf (node-color (node-parent node)) :red)
(setf (node-color sib) :black)
(if (eq node (node-left (node-parent node)))
(node-rotate-left (node-parent node))
(node-rotate-right (node-parent node))))
(node-delete/case-3 node)))
(defun node-delete/case-1 (node)
(declare (optimize (speed 3)))
(unless (typep (node-parent node) 'rb-node/sentinel)
(node-delete/case-2 node)))
(defmethod node-delete ((node rb-node))
(let ((child (if (eq (node-right node) *rb-node-sentinel*)
(node-left node)
(node-right node))))
(when (not child)
(error "Expected to find a child"))
(node-replace node child)
(when (eql (node-color node) :black)
(if (eql (node-color child) :red)
(setf (node-color child) :black)
(node-delete/case-1 child)))
(decf (slot-value (node-parent-tree node) 'count))))
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;;;
;;; Tree Methods
;;;
(defmethod tree-find ((tree rb-tree) data &key (test #'eq))
(unless (tree-root tree)
;; No root, nothing to find
(return-from tree-find))
(let ((node (tree-root tree))
(lt-test (tree-test-fn tree)))
(loop until (typep node 'rb-node/sentinel) do
(cond
;; Found the node
((rb-node-test* data node test)
(return-from tree-find node))
;; data < node
((rb-node-test* data node lt-test)
(setf node (node-left node)))
;; Data > node
(t (setf node (node-right node))))))
nil)
(defmethod tree-insert ((tree rb-tree) (node rb-node))
(setf (tree-root tree) (node-insert (tree-root tree) node)))
(defmethod tree-insert ((tree rb-tree) data)
(tree-insert tree (tree-make-node tree data)))
(defmethod tree-delete ((tree rb-tree) data)
;; Do a search for it first
(let ((node (tree-find tree data)))
(unless (typep node 'rb-node/sentinel)
(node-delete node)))
nil)
(defmethod tree-make-node ((tree rb-tree) &optional data)
(make-instance 'rb-node :parent-tree tree :data data))
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;;;
;;; Printing functions, mostly for debugging
;;;
(defmethod node-print ((node rb-node) (depth integer) &optional (side ""))
(format t "~a~a~a~%" (make-string (* 3 depth) :initial-element #\Space) side node)
(when (node-left node)
(node-print (node-left node) (1+ depth) "L: "))
(when (node-right node)
(node-print (node-right node) (1+ depth) "R: ")))
(defmethod node-print ((node rb-node/sentinel) depth &optional side)
(declare (ignore node depth side))
nil)
(defmethod tree-print ((tree rb-tree))
(if (tree-root tree)
(node-print (tree-root tree) 0)
(format t "Empty tree~%")))