;; Persistent Red-Black Tree Implementation
;; Author: Timothy Beyer
;; License: MIT License
;; Description: A persistent Red-Black Tree, in the style of Chris Okasaki's
;; "Purely Functional Data Structures" version
;;
;; Inspired by solution on "Programming Praxis" Blog (but not based on)
;; url: http://programmingpraxis.com/2009/10/02/red-black-trees/
;;
;; Note: Lacks most operations right now, runs in Chicken Scheme 4.x

(module red-black-tree
    (key value color left right make-node search insert traverse-inorder
         traverse-preorder traverse-postorder leaf)

(import chicken scheme matchable data-structures)
(require-library matchable utf8-srfi-13)

(define-syntax define*
  (syntax-rules ()
    ((_ name body ...) (define name (match-lambda* body ...)))))

;; _ :: [string] -> string
(define ++ string-append)

;; _ :: node a b -> maybe a
(define (key n) (vector-ref n 0))

;; _ :: node a b -> maybe b
(define (value n) (vector-ref n 1))

;; _ :: node a b -> maybe symbol
(define (color n) (vector-ref n 2))

;; _ :: node a b -> maybe node a b
(define (left n) (vector-ref n 3))
(define (right n) (vector-ref n 4))

;; _ :: a -> b -> symbol -> node a b -> node a b -> node a b
(define (make-node k v c lc rc) (vector k v c lc rc))

;; _ :: node a b
(define (leaf) (make-node 'null 'null 'B 'null 'null))
(define root leaf)

;; _ :: a -> b -> symbol -> node a b
(define (make-leaves key val color) (make-node key val color (leaf) (leaf)))

;; _ :: node a b -> boolean
(define (leaf? n) (equal? n (leaf)))

;; _ :: node a b -> a
(define* search
  ((T K) (search T K <))
  (((? leaf? T) _ _) #f)
  ((#(K* V* _ L* R*) K Cmp)
   (cond ((Cmp K K*) (search L* K Cmp))
         ((Cmp K* K) (search R* K Cmp))
         (else V*))))

;; _ :: node a b -> node a b
(define (recolor-parent n)
  (make-node (key n) (value n) 'B (left n) (right n)))

;; _ :: node a b -> a -> b -> function-symbol -> node a
(define* insert
  ((T K V) (insert T K V <))
  (((? leaf? T) K V Cmp) (recolor-parent (make-leaves K V 'R)))
  ((#(K* V* C* L* R*) K V Cmp)
   (recolor-parent
    (cond ((Cmp K K*) (balance (make-node K* V* C* (insert L* K V Cmp) R*)))
          ((Cmp K* K) (balance (make-node K* V* C* L* (insert R* K V Cmp))))
          (else (make-node K V C* L* R*))))))

;; _ node a b -> node a b
(define* balance
  ;; 1) red left child has red left grandchild
  (#(K V 'B #(K* V* 'R #(K** V** 'R L** R**) R*) R)
   (make-node K* V* 'R (make-node K** V** 'B L** R**) (make-node K V 'B R* R)))
  ;; 2) red left child has red right grandchild
  (#(K V 'B #(K* V* 'R L* #(K** V** 'R L** R**)) R)
   (make-node K** V** 'R (make-node K* V* 'B L* L**) (make-node K V 'B R** R)))
  ;; 3) red right child has red left grandchild
  (#(K V 'B L #(K* V* 'R #(K** V** 'R L** R**) R*))
   (make-node K** V** 'R (make-node K V 'B L L**) (make-node K* V* 'B R** R*)))
  ;; 4) red right child has red right grandchild
  (#(K V 'B L #(K* V* 'R L* #(K** V** 'R L** R**)))
   (make-node K* V* 'R (make-node K V 'B L L*) (make-node K** V** 'B L** R**)))
  ((T) T))

;; _ :: node a b -> maybe IO ()
(define (traverse-inorder t)
  (if (leaf? t) ""
    (let ((key* (->string (key t))) (val (->string (value t)))
          (color* (->string (color t))))
      (traverse-inorder (left t))
      (display (++ "(" key* "," val "," color* ") "))
      (traverse-inorder (right t)))))
(define (traverse-preorder t)
  (if (leaf? t) ""
    (let ((key* (->string (key t))) (val (->string (value t)))
          (color* (->string (color t))))
      (display (++ "(" key* "," val "," color* ") "))
      (traverse-preorder (left t))
      (traverse-preorder (right t)))))
(define (traverse-postorder t)
  (if (leaf? t) ""
    (let ((key* (->string (key t))) (val (->string (value t)))
          (color* (->string (color t))))
      (traverse-preorder (left t))
      (traverse-preorder (right t))
      (display (++ "(" key* "," val "," color* ") ")))))

;; example:
;;(define t (root))
;;(define t (insert t 2 "b" <))
;;(display (++ (traverse-inorder t) "\n"))
;;(define t (insert t 5 "e" <))
;;(display (++ (traverse-inorder t) "\n"))
;;(define t (insert t 3 "c" <))
;;(display (++ (traverse-inorder t) "\n"))
;;(define t (insert t 4 "d" <))
;;(display (++ (traverse-inorder t) "\n"))
;;(define t (insert t 1 "a" <))
;;(display (++ (traverse-inorder t) "\n"))

)

;;; Local Variables:
;;; scheme-program-name: csc-library
;;; End:

Line 13:4: require: bad module path in: (key value color left right make-node search insert traverse-inorder traverse-preorder traverse-postorder leaf)