red_black_tree.clj

(ns forestry.rb
  (:require [clojure.core.match   :refer [match]]
            [forestry.node :as n])
  (:refer-clojure :exclude [find key val replace]))

(defprotocol -RedBlackNode
  (color   [node])
  (blacken [node])
  (-add     [node k v])
  (balance [node])
  (red?    [node]))

(defn search [v k]
  (when v
    (let [cmp (compare k (n/key v))]
      (cond (zero? cmp) (n/val v)
            (neg?  cmp) (recur (n/left  v) k)
            :else       (recur (n/right v) k)))))

(extend-type clojure.lang.IPersistentVector
  n/AssociativeNode
  (key [v] (nth v 2))
  (val [v] (nth v 3))
  (add [n k v]
    (blacken (-add n k v)))
  n/BinaryNode
  (left  [v] (second v))
  (right [v] (last v))
  -RedBlackNode
  (color   [v] (first v))
  (blacken [v] (assoc v 0 :B))
  (red?    [v] (identical? :R (color v)))
  (-add    [[c l k v r] k' v']
    (let [cmp (compare k' k)]
      (if (zero? cmp)
        [c l k v' r]
        (balance
         (if (neg? cmp)
           [c (-add l k' v') k v r]
           [c l k v (-add r k' v')])))))
  (balance [[_ L K V R :as node]]
    (if (red? node)
      node
      (match [L K V R]
        [[:R [:R a xk xv b] yk yv c] zk zv d] [:R [:B a xk xv b] yk yv [:B c zk zv d]]
        [[:R a xk xv [:R b yk yv c]] zk zv d] [:R [:B a xk xv b] yk yv [:B c zk zv d]]
        [a xk xv [:R [:R b yk yv c] zk zv d]] [:R [:B a xk xv b] yk yv [:B c zk zv d]]
        [a xk xv [:R b yk yv [:R c zk zv d]]] [:R [:B a xk xv b] yk yv [:B c zk zv d]]
        :else node))))

(extend-type nil
  -RedBlackNode
  (-add [_ k v] [:R nil k v nil]))