1
0
Fork 0
tinmop/src/rb-tree.lisp

519 lines
22 KiB
Common Lisp

;; tinmop: an humble gemini and pleroma client
;; Copyright (C) 2020 cage
;; 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 :rb-tree)
(alexandria:define-constant +rb-red+ :red :test #'eq)
(alexandria:define-constant +rb-black+ :black :test #'eq)
(alexandria:define-constant +rb-neg-black+ :neg-black :test #'eq)
(alexandria:define-constant +rb-black-black+ :black-black :test #'eq)
(alexandria:define-constant +rb-color+ :color :test #'eq)
(defun incf-black (color)
(cond
((eq color +rb-black+)
+rb-black-black+)
((eq color +rb-red+)
+rb-black+)
((eq color +rb-neg-black+)
+rb-red+)))
(defun decf-black (color)
(cond
((eq color +rb-black-black+)
+rb-black+)
((eq color +rb-black+)
+rb-red+)
((eq color +rb-red+)
+rb-neg-black+)))
(defclass rb-node (node)
((color
:initarg :color
:initform +rb-black+
:accessor color)))
(defgeneric balance (object))
(defgeneric bubble (object))
(defgeneric balancedp (object))
(defgeneric left-balance (object))
(defgeneric right-balance (object))
(defgeneric remove-max-node (object key key-datum compare equal))
(defgeneric %remove-node (object needle key key-datum compare equal))
(defgeneric remove-node (object needle &key key key-datum compare equal))
(defmethod node->string ((object rb-node))
(if (null (data object))
(format nil "leaf color ~a" (color object))
(format nil "~a (~a)~% [~a] [~a]"
(data object)
(color object)
(node->string (left object))
(node->string (right object)))))
(defun make-rb-node (color data left right parent)
(make-instance 'rb-node :color color :left left :right right :data data :parent parent))
(defun make-rb-leaf (color parent)
(make-instance 'rb-node :color color :parent parent :left nil :right nil))
(defun make-root-rb-node (datum color)
(let* ((tree (make-rb-node color datum nil nil nil))
(l-leaf (make-rb-leaf +rb-black+ tree))
(r-leaf (make-rb-leaf +rb-black+ tree)))
(setf (left tree) l-leaf
(right tree) r-leaf)
tree))
(defmethod insert ((object rb-node) datum &key
(key #'identity) (key-datum #'identity)
(compare #'<) (equal #'=))
(macrolet ((make-leaf-node (datum left right parent)
`(make-rb-node +rb-red+ ,datum ,left ,right ,parent))
(make-leaf (new-node)
`(make-rb-leaf +rb-black+ ,new-node))
(make-node (data left right parent)
`(make-rb-node (color node) ,data ,left ,right ,parent)))
(with-insert-local-function (make-node make-node make-leaf-node make-leaf
left-balance right-balance)
(let ((balanced (%insert object datum key key-datum compare equal)))
(setf (color balanced) +rb-black+)
balanced))))
(defmacro with-match-tree ((color left data right) tree &body body)
`(and (or (not ,color)
(eq ,color (color ,tree)))
(let ((,data (data ,tree)))
(declare (ignorable ,data))
,(if (consp left)
`(with-match-tree ,left (left ,tree)
,(if (consp right)
`(with-match-tree ,right (right ,tree)
,@body)
`(let ((,right (right ,tree)))
(declare (ignorable ,right))
,@body)))
`(let ((,left (left ,tree)))
(declare (ignorable ,left))
,(if (consp right)
`(with-match-tree ,right (right ,tree)
,@body)
`(let ((,right (right ,tree)))
(declare (ignorable ,right))
,@body)))))))
(defmethod balance ((object rb-node))
(macrolet ((setf-parent (new-node)
`(setf (parent (left ,new-node)) ,new-node
(parent (right ,new-node)) ,new-node
(parent (left (left ,new-node))) (left ,new-node)
(parent (right (left ,new-node))) (left ,new-node)
(parent (right (right ,new-node))) (right ,new-node)
(parent (left (right ,new-node))) (right ,new-node))))
(with-match-tree (+rb-black+ (+rb-red+ (+rb-red+ a x b) y c) z d) object
(return-from balance
(let ((new-node (make-rb-node +rb-red+ y
(make-rb-node +rb-black+ x a b nil)
(make-rb-node +rb-black+ z c d nil)
(parent object))))
(setf-parent new-node)
new-node)))
(with-match-tree (+rb-black+ (+rb-red+ a x (+rb-red+ b y c)) z d) object
(return-from balance
(let ((new-node (make-rb-node +rb-red+ y
(make-rb-node +rb-black+ x a b nil)
(make-rb-node +rb-black+ z c d nil)
(parent object))))
(setf-parent new-node)
new-node)))
(with-match-tree (+rb-black+ a x (+rb-red+ (+rb-red+ b y c) z d)) object
(return-from balance
(let ((new-node (make-rb-node +rb-red+ y
(make-rb-node +rb-black+ x a b nil)
(make-rb-node +rb-black+ z c d nil)
(parent object))))
(setf-parent new-node)
new-node)))
(with-match-tree (+rb-black+ a x (+rb-red+ b y (+rb-red+ c z d))) object
(return-from balance
(let ((new-node (make-rb-node +rb-red+ y
(make-rb-node +rb-black+ x a b nil)
(make-rb-node +rb-black+ z c d nil)
(parent object))))
(setf-parent new-node)
new-node)))
;; for deletion
(with-match-tree (+rb-black-black+ (+rb-red+ a x (+rb-red+ b y c)) z d) object
(return-from balance
(let ((new-node (make-rb-node +rb-black+
y
(make-rb-node +rb-black+ x a b nil)
(make-rb-node +rb-black+ z c d nil)
(parent object))))
(setf-parent new-node)
new-node)))
(with-match-tree (+rb-black-black+ (+rb-red+ (+rb-red+ a x b) y c) z d) object
(return-from balance
(let ((new-node (make-rb-node +rb-black+
y
(make-rb-node +rb-black+ x a b nil)
(make-rb-node +rb-black+ z c d nil)
(parent object))))
(setf-parent new-node)
new-node)))
(with-match-tree (+rb-black-black+ a x (+rb-red+ (+rb-red+ b y c) z d)) object
(return-from balance
(let ((new-node (make-rb-node +rb-black+
y
(make-rb-node +rb-black+ x a b nil)
(make-rb-node +rb-black+ z c d nil)
(parent object))))
(setf-parent new-node)
new-node)))
(with-match-tree (+rb-black-black+ a x (+rb-red+ b y (+rb-red+ c z d))) object
(return-from balance
(let ((new-node (make-rb-node +rb-black+
y
(make-rb-node +rb-black+ x a b nil)
(make-rb-node +rb-black+ z c d nil)
(parent object))))
(setf-parent new-node)
new-node)))
(with-match-tree (+rb-black-black+ (+rb-neg-black+ (+rb-black+ a w b)
x
(+rb-black+ c y d))
z
e) object
(return-from balance
(let ((new-node (make-rb-node +rb-black+
y
(balance (make-rb-node +rb-black+
x
(make-rb-node +rb-red+ w a b nil)
c
nil))
(make-rb-node +rb-black+ z d e nil)
(parent object))))
(setf-parent new-node)
(balance new-node))))
(with-match-tree (+rb-black-black+ a
x
(+rb-neg-black+ (+rb-black+ b y c)
z
(+rb-black+ j n i))) object
(return-from balance
(let ((new-node (make-rb-node +rb-black+
y
(make-rb-node +rb-black+
x
a
b
nil)
(balance (make-rb-node +rb-black+
z
c
(make-rb-node +rb-red+ n j i nil)
nil))
(parent object))))
(setf-parent new-node)
(balance new-node))))
object))
(defmethod bubble ((object rb-node))
(if (or (eq (color (left object)) +rb-black-black+)
(eq (color (right object)) +rb-black-black+))
;; TODO reconstruct-parent!
(balance (make-rb-node (incf-black (color object))
(data object)
(make-rb-node (decf-black (color (left object)))
(data (left object))
(left (left object))
(right (left object))
nil)
(make-rb-node (decf-black (color (right object)))
(data (right object))
(left (right object))
(right (right object))
nil)
nil))
(balance object)))
(defmethod left-balance ((object rb-node))
(macrolet ((setf-parent (new-node)
`(setf (parent (left ,new-node)) ,new-node
(parent (right ,new-node)) ,new-node
(parent (left (left ,new-node))) (left ,new-node)
(parent (right (left ,new-node))) (left ,new-node)
(parent (right (right ,new-node))) (right ,new-node)
(parent (left (right ,new-node))) (right ,new-node))))
(with-match-tree (+rb-black+ (+rb-red+ (+rb-red+ a x b) y c) z d) object
(return-from left-balance
(let ((new-node (make-rb-node +rb-red+ y
(make-rb-node +rb-black+ x a b nil)
(make-rb-node +rb-black+ z c d nil)
(parent object))))
(setf-parent new-node)
new-node)))
(with-match-tree (+rb-black+ (+rb-red+ a x (+rb-red+ b y c)) z d) object
(return-from left-balance
(let ((new-node (make-rb-node +rb-red+ y
(make-rb-node +rb-black+ x a b nil)
(make-rb-node +rb-black+ z c d nil)
(parent object))))
(setf-parent new-node)
new-node)))
object))
(defmethod right-balance ((object rb-node))
(macrolet ((setf-parent (new-node)
`(setf (parent (left ,new-node)) ,new-node
(parent (right ,new-node)) ,new-node
(parent (left (left ,new-node))) (left ,new-node)
(parent (right (left ,new-node))) (left ,new-node)
(parent (right (right ,new-node))) (right ,new-node)
(parent (left (right ,new-node))) (right ,new-node))))
(with-match-tree (+rb-black+ a x (+rb-red+ (+rb-red+ b y c) z d)) object
(return-from right-balance
(let ((new-node (make-rb-node +rb-red+ y
(make-rb-node +rb-black+ x a b nil)
(make-rb-node +rb-black+ z c d nil)
(parent object))))
(setf-parent new-node)
new-node)))
(with-match-tree (+rb-black+ a x (+rb-red+ b y (+rb-red+ c z d))) object
(return-from right-balance
(let ((new-node (make-rb-node +rb-red+ y
(make-rb-node +rb-black+ x a b nil)
(make-rb-node +rb-black+ z c d nil)
(parent object))))
(setf-parent new-node)
new-node)))
object))
(defmethod map ((object rb-node) function)
(with-accessors ((color color) (data data) (left left) (right right)) object
(if (leafp object)
(make-rb-leaf color nil)
(make-rb-node color (funcall function data)
(map left function)
(map right function) nil))))
(defmethod map-node ((object rb-node) function)
(with-accessors ((color color) (data data) (left left) (right right)) object
(if (leafp object)
(funcall function object (make-rb-leaf color object))
(funcall function object (make-rb-node color data
(map-node left function)
(map-node right function) nil)))))
(defmethod reconstruct-parent ((object rb-node) &optional (parent (parent object)))
(with-accessors ((color color) (data data) (left left) (right right)) object
(if (leafp object)
(make-rb-leaf color parent)
(make-rb-node color data
(reconstruct-parent left object)
(reconstruct-parent right object) parent))))
(defmethod remove-max-node ((object node) key key-datum compare equal)
(if (leafp (right object))
(%remove-node object (data object) key key-datum compare equal)
(bubble (make-rb-node (color object)
(data object)
(left object)
(remove-max-node (right object)
key
key-datum
compare
equal)
nil))))
(defmethod %remove-node ((object node) needle key key-datum compare equal)
(cond
((leafp object)
object)
((funcall equal (%key key (data object)) (%key key-datum needle)) ;; equal, delete
(cond
((and (all-children-leaf-p object) ; no children, red color
(eq (color object) +rb-red+))
(make-rb-leaf +rb-black+ nil))
((and (all-children-leaf-p object) ; no children, black color
(eq (color object) +rb-black+))
(make-rb-leaf +rb-black-black+ nil))
((and (leafp (left object)) ;; one child on right
(not (leafp (right object))))
(make-rb-node +rb-black+
(data (right object))
(left (right object))
(right (right object))
nil))
((and (leafp (right object)) ; one child on left
(not (leafp (left object))))
(make-rb-node +rb-black+
(data (left object))
(left (left object))
(right (left object))
nil))
(t ; two children
(let* ((max-node (find-max-node (left object)))
(max-data (data max-node)))
(bubble (make-rb-node (color object)
max-data
(remove-max-node (left object)
key
key
compare
equal)
(right object)
nil)))))) ;; end removing block
((funcall compare (%key key-datum needle) (%key key (data object)))
(bubble (make-rb-node (color object)
(data object)
(%remove-node (left object) needle key key-datum compare equal)
(right object)
nil)))
(t ; go right, needle is greater then this node
(bubble (make-rb-node (color object)
(data object)
(left object)
(%remove-node (right object) needle key key-datum compare equal)
nil)))))
(defmethod remove-node ((object rb-node) needle &key
(key #'identity) (key-datum #'identity)
(compare #'<) (equal #'=))
(let ((new (%remove-node object needle key key-datum compare equal)))
(setf (color new) +rb-black+)
new))
(defmethod node->dot ((object rb-node))
(labels ((color->hex (node)
(cond
((eq (color node) +rb-red+)
"#ff0000")
((eq (color node) +rb-black-black+)
"#ff00ff")
((eq (color node) +rb-neg-black+)
"#a1a1a1")
(t
"#ffffff")))
(nodes ()
(append
(list
`(:node ((:id ,(format nil "~a" (data object)))
(:label ,(format nil "~ap~a" (data object)
(data (parent object))))
(:style "filled")
(:fillcolor ,(color->hex object)))))
(if (not (leafp (left object)))
(node->dot (left object))
(list
`(:node ((:id ,(format nil "nil-l~a" (data object)))
(:label "nil")
(:style "filled")
(:fillcolor ,(color->hex (left object)))))))
(if (not (leafp (right object)))
(node->dot (right object))
(list
`(:node ((:id ,(format nil "nil-r~a" (data object)))
(:label "nil")
(:style "filled")
(:fillcolor ,(color->hex (right object)))))))))
(edges ()
(append
(if (data (left object))
(list `(:edge
((:from ,(format nil "~a" (data object)))
(:to ,(format nil "~a" (data (left object)))))))
(list `(:edge
((:from ,(format nil "~a" (data object)))
(:to ,(format nil "nil-l~a" (data object)))))))
(if (data (right object))
(list `(:edge
((:from ,(format nil "~a" (data object)))
(:to ,(format nil "~a" (data (right object)))))))
(list `(:edge
((:from ,(format nil "~a" (data object)))
(:to ,(format nil "nil-r~a" (data object))))))))))
(append (nodes) (edges))))
(defmethod to-sexp ((object rb-node))
(let ((*print-circle* t))
(list +data+ (to-sexp (data object))
+rb-color+ (color object)
+left+ (to-sexp (left object))
+right+ (to-sexp (right object))
+parent+ (to-sexp (data (parent object))))))
(defmethod from-sexp ((object rb-node) sexp)
(declare (ignorable object))
(labels ((%from-sexp (sexp)
(if (null sexp)
(make-rb-leaf +rb-black+ nil)
(make-rb-node (getf sexp +rb-color+)
(getf sexp +data+)
(from-sexp object (getf sexp +left+))
(from-sexp object (getf sexp +right+)) nil))))
(let ((new-tree (%from-sexp sexp)))
(reconstruct-parent new-tree))))
(defmethod balancedp ((object rb-node))
(let ((black-paths-count '()))
(labels ((all-red-child-black-p (node)
(map-node node #'(lambda (n b)
(declare (ignore b))
(if (or (eq (color n) +rb-black+)
(and (eq (color n) +rb-red+)
(eq (color (left n)) +rb-black+)
(eq (color (right n)) +rb-black+)))
t
(return-from all-red-child-black-p nil)))))
(all-path-black-same-length (node &optional (ct 0))
(when (eq (color node) +rb-black+)
(incf ct))
(if (leafp node)
(progn
(push ct black-paths-count)
(decf ct))
(progn
(all-path-black-same-length (left node) ct)
(all-path-black-same-length (right node) ct)))))
(all-path-black-same-length object)
(and (every #'(lambda (a) (= a (elt black-paths-count 0))) black-paths-count)
(all-red-child-black-p object)))))