Red Black Tree for Common Lisp.

rbtree.lisp

;; The following implementation of rb-tree is based on http://www.cs.kent.ac.uk/people/staff/smk/redblack/.

(defun change-to-black (tree)
  (pattern-match tree
    ((:pattern (_ . rest) :variable rest :ignore _) `(:B . ,rest))
    (:otherwise nil)))

(defun rb-insert (tree obj cmp)
  (change-to-black (rb-insert% tree obj cmp)))

(defun rb-insert% (tree obj cmp)
  (pattern-match tree
    ((:pattern (color left x right) :guard (funcall cmp obj x) :variable (color left x right))
     (balance `(,color ,(rb-insert% left obj cmp) ,x ,right)))
    ((:pattern (color left x right) :guard :otherwise          :variable (color left x right))
     (balance `(,color ,left ,x ,(rb-insert% right obj cmp))))
    ((:pattern nil)
     `(:R nil ,obj nil))
    (:otherwise
     (error "The first argument is not a proper rb-tree."))))

(defun rb-remove (tree obj cmp)
  (change-to-black (rb-remove% tree obj cmp)))

(defun rb-remove% (tree obj cmp)
  (pattern-match tree
    ((:pattern (color left x right) :variable (color left x right))
     (let ((less    (funcall cmp obj x))
           (greater (funcall cmp x obj)))
       (cond ((and less greater) (balance (rb-remove-root tree cmp)))
             (less               (balance `(,color ,(rb-remove% left obj cmp) ,x ,right)))
             (greater            (balance `(,color ,left ,x ,(rb-remove% right obj cmp))))
             (:otherwise         (balance (rb-remove-root tree cmp))))))
    ((:pattern nil)
      nil)
    (:otherwise
     (error "The first argument is not a proper rb-tree."))))

(defun rb-remove-root (tree cmp)
  (pattern-match tree
    ((:pattern (:B a _ Nil) :variable a :ignore _)
     (mark a))
    ((:pattern (:B Nil _ a) :variable a :ignore _)
     (mark a))
    ((:pattern (:R Nil _ Nil) :ignore _)
     nil)
    ((:pattern (color a _ b) :variable (color a b) :ignore _)
     (let ((min (rb-minimum b cmp)))
       `(,color ,a ,min ,(rb-remove% b min cmp))))))

(defun markedp (x)
  (pattern-match x
    ((:pattern :X)                 t)
    ((:pattern (:X . _) :ignore _) t)
    (:otherwise                    nil)))

(defun mark (x)
  (pattern-match x
    ((:pattern Nil)                        :X)
    ((:pattern (:B . rest) :variable rest) `(:X . ,rest))
    ((:pattern (:R . rest) :variable rest) `(:B . ,rest))))

(defun unmark (x)
  (pattern-match x
    ((:pattern :X)                         `Nil)
    ((:pattern (:X . rest) :variable rest) `(:B . ,rest))
    ((:pattern (:B . rest) :variable rest) `(:R . ,rest))))

(defun balance (tree)
  (pattern-match tree
    ((:pattern (:B (:R a x b) y (:R c z d)) :variable (a x b y c z d))
     `(:R (:B ,a ,x ,b) ,y (:B ,c ,z ,d)))
    ((:pattern (:B (:R (:R a x b) y c) z d) :variable (a x b y c z d))
     `(:R (:B ,a ,x ,b) ,y (:B ,c ,z ,d)))
    ((:pattern (:B (:R a x (:R b y c)) z d) :variable (a x b y c z d))
     `(:R (:B ,a ,x ,b) ,y (:B ,c ,z ,d)))
    ((:pattern (:B a x (:R (:R b y c) z d)) :variable (a x b y c z d))
     `(:R (:B ,a ,x ,b) ,y (:B ,c ,z ,d)))
    ((:pattern (:B a x (:R b y (:R c z d))) :variable (a x b y c z d))
     `(:R (:B ,a ,x ,b) ,y (:B ,c ,z ,d)))
    ((:pattern (:B a x (:R b y c)) :guard (markedp a) :variable (a x b y c))
     (balance `(:B ,(balance `(:R ,a ,x ,b)) ,y ,c)))
    ((:pattern (:B (:R a x b) y c) :guard (markedp c) :variable (a x b y c))
     (balance `(:B ,a ,x ,(balance `(:R ,b ,y ,c)))))
    ((:pattern (:B a x (:B . b)) :guard (markedp a) :variable (a x b))
     (mark (balance `(:B ,(unmark a) ,x (:R . ,b)))))
    ((:pattern (:R a x (:B . b)) :guard (markedp a) :variable (a x b))
     (balance `(:B ,(unmark a) ,x (:R . ,b))))
    ((:pattern (:B (:B . a) x b) :guard (markedp b) :variable (a x b))
     (mark (balance `(:B (:R . ,a) ,x ,(unmark b)))))
    ((:pattern (:R (:B . a) x b) :guard (markedp b) :variable (a x b))
     (balance `(:B (:R . ,a) ,x ,(unmark b))))
    (:otherwise
     tree)))

(defun rb-minimum (tree cmp)
  (pattern-match tree
    ((:pattern (_ Nil x _) :variable x :ignore _)
     (values x t))
    ((:pattern (_ left _ _) :variable left :ignore _)
     (rb-minimum left cmp))
    ((:pattern nil)
     (values nil nil))
    (:otherwise
     (error "The first argument is not a proper rb-tree."))))

(defun rb-find (tree obj cmp)
  (pattern-match tree
    ((:pattern (_ left x right) :variable (left x right) :ignore _)
     (let ((less    (funcall cmp obj x))
           (greater (funcall cmp x obj)))
       (cond ((and less greater) (values x t))
             (less               (rb-find left  obj cmp))
             (greater            (rb-find right obj cmp))
             (:otherwise         (values x t)))))
    ((:pattern nil)
     (values nil nil))
    (:otherwise
     (error "The first argument is not a proper rb-tree."))))

;; The implementation of rb-tree ends here.

(defmacro with-gensyms (gensym-variables &body body)
  `(let ,(loop for gensym-variable in gensym-variables
               collect `(,gensym-variable (gensym (symbol-name ',gensym-variable))))
     ,@body))

(defun matcher (target pattern variables values ignores)
  (labels ((matcher% (target pattern)
             (cond ((consp pattern)
                    (with-gensyms (cartarget cdrtarget)
                                  `(and (consp ,target)
                                        (let ((,cartarget (car ,target))
                                              (,cdrtarget (cdr ,target)))
                                          (and ,(matcher% cartarget (car pattern))
                                               ,(matcher% cdrtarget (cdr pattern)))))))
                   ((symbolp pattern)
                    (cond ((member pattern ignores)
                           t)
                          ((member pattern variables)
                           (setf variables (remove pattern variables))
                           (setf values    (cons   pattern values))
                           `(progn (setf ,pattern ,target) t))
                          ((member pattern values)
                           `(equal ,target ,pattern))
                          (:otherwise
                           `(eq ,target ',pattern))))
                   (:otherwise
                    `(equal ,target ,pattern)))))
    (matcher% target pattern)))

(defmacro if-matches (target-expr (&key (variable nil) (value nil) (guard t) (pattern nil pattern-bound-p) (ignore nil)) then else)
  (let ((variable (if (listp variable) variable (list variable)))
        (value    (if (listp value)    value    (list value)))
        (ignore   (if (listp ignore)   ignore   (list ignore))))
    (with-gensyms (target)
      `(let ,variable
         (let ((,target ,target-expr))
           (if (and ,(if pattern-bound-p 
                         (matcher target pattern variable value ignore)
                       t)
                    ,guard)
               ,then
             ,else))))))

(defmacro pattern-match (target-expr &body match-clauses)
  (with-gensyms (target block-name)
    `(block ,block-name
       (let ((,target ,target-expr))
         ,@(loop for match-clause in match-clauses
                 collect (if (eq (first match-clause) :otherwise)
                             `(return-from ,block-name (progn ,@(rest match-clause)))
                             `(if-matches ,target ,(first match-clause)
                                      (return-from ,block-name (progn ,@(rest match-clause)))
                                      nil)))))))