Persistent sorted set implemented as a red-black tree in Common Lisp

sset.lisp

;; todo: how do you actually manage dependencies?
(ql:quickload "generic-cl")
(ql:quickload "trivia")

;; Persistent red-black tree implemnetation per Okasaki's paper:
;; https://www.cs.tufts.edu/~nr/cs257/archive/chris-okasaki/redblack99.pdf

(defpackage sset
  (:use generic-cl))

(in-package sset)

(defstruct rbtree color left element right)

(defstruct empty-rbtree)

(defun empty-set () (make-empty-rbtree))

(defun set-member (rbtree item)
  (trivia:match rbtree
    ((empty-rbtree) nil)
    ((rbtree (color 'red)
                (left l)
                (element e)
                (right r))
     (cond
       ((< item e) (set-member l item))
       ((equalp item e) t)
       ((> item e) (set-member r item))))))

(defun balance (rbtree)
  (trivia:match rbtree
    ;; match all black subtrees with red child with another red child, so
    ;; we can reconstruct them as a red subtree with two black
    ;; children (Okasaki fig. 1)
    ((or (rbtree :color 'black
                 :left (rbtree :color 'red
                             :left (rbtree :color 'red
                                         :left a
                                         :element x
                                         :right b)
                             :element y
                             :right c)
                 :element z
                 :right d)
         (rbtree :color 'black
                 :left (rbtree :color 'red
                               :left a
                               :element x
                               :right (rbtree :color 'red
                                              :left b
                                              :element y
                                              :right c))
                 :element z
                 :right d)
         (rbtree :color 'black
                 :left a
                 :element x
                 :right (rbtree :color 'red
                                :left (rbtree :color 'red
                                              :left b
                                              :element y
                                              :right c)
                                :element z
                                :right d))
         (rbtree :color 'black
                 :left a
                 :element x
                 :right (rbtree :color 'red
                                :left b
                                :element y
                                :right (rbtree :color 'red
                                               :left c
                                               :element z
                                               :right d))))
     (make-rbtree :color 'red
                  :left (make-rbtree :color 'black
                                     :left a
                                     :element x
                                     :right b)
                  :element y
                  :right (make-rbtree :color 'black
                                      :left c
                                      :element z
                                      :right d)))
    (rbtree rbtree)))

(defun add (rbtree item)
  (labels ((ins (tree)
             (trivia:match tree
               ((empty-rbtree) (make-rbtree :color 'red
                                            :left (make-empty-rbtree)
                                            :element item
                                            :right (make-empty-rbtree)))
               ((rbtree :color color :left left :element element :right right)
                (cond
                  ((< item element) (balance (make-rbtree :color color
                                                          :left (ins left)
                                                          :element element
                                                          :right right)))
                  ;; return a new tree even though it's identical,
                  ;; because we will be changing the color to black
                  ((= item element) (make-rbtree :color color
                                                 :left left
                                                 :element element
                                                 :right right))
                  ((> item element) (balance (make-rbtree :color color
                                                          :left left
                                                          :element element
                                                          :right (ins right)))))))))
    (let ((new-tree (ins rbtree)))
      (setf (slot-value new-tree 'color) 'black)
      new-tree)))

(defun elements (rbtree)
  (trivia:match rbtree
    ((empty-rbtree) '())
    ((rbtree :color _ :left left :element elt :right right)
     (concatenate (elements left) (list elt) (elements right)))))

(defun list->rbtree (list)
  (reduce #'add list :initial-value (make-empty-rbtree)))

(defun range (rbtree &key from to)
  ;; todo: iterator, not list?
  (trivia:match rbtree
    ((empty-rbtree) '())
    ((rbtree :left l :element x :right r)
     (concatenate (when (< from x)
                    (range l :from from :to to))
                  (when (and (<= from x) (< x to))
                    (list x))
                  (when (< x to)
                    (range r :from from :to to))))))

;; (range (list->rbtree '(1 2 3 4 5 6 7 8 9 10)) :from 4 :to 7)

;; (elements (list->rbtree '("bear" "apple" "pear" "orange")))

;; trees should be sorted:
;; (time
;;  (= (elements (list->rbtree (loop for x from 10000 downto 0 collect x)))
;;     (elements (list->rbtree (loop for x from 0 to 10000 collect x)))))