(library (llrb (1)) (export make-tree tree-insert! tree-delete! tree-search tree-get tree-has? ) (import (chezscheme)) (define (make-node key val) (vector #t key val '() '())) (define (node-color node) (vector-ref node 0)) (define (node-key node) (vector-ref node 1)) (define (node-val node) (vector-ref node 2)) (define (node-left node) (vector-ref node 3)) (define (node-right node) (vector-ref node 4)) (define (node-set-color! node color) (vector-set! node 0 color)) (define (node-set-key! node key) (vector-set! node 1 key)) (define (node-set-val! node val) (vector-set! node 2 val)) (define (node-set-left! node left) (vector-set! node 3 left)) (define (node-set-right! node right) (vector-set! node 4 right)) (define (node-flip-colors! node) (define (flip-color! node) (if (not (null? node)) (node-set-color! node (not (node-color node))))) (flip-color! node) (flip-color! (node-left node)) (flip-color! (node-right node))) (define (node-rotate-left! node) (define x (node-right node)) (node-set-right! node (node-left x)) (node-set-left! x node) (node-set-color! x (node-color node)) (node-set-color! node #t) x) (define (node-rotate-right! node) (define x (node-left node)) (node-set-left! node (node-right x)) (node-set-right! x node) (node-set-color! x (node-color node)) (node-set-color! node #t) x) (define (node-move-red-left! node) (node-flip-colors! node) (let ((right (node-right node))) (if (and (not (null? right)) (node-red? (node-left right))) (begin (node-set-right! node (node-rotate-right! right)) (set! node (node-rotate-left! node)) (node-flip-colors! node)))) node) (define (node-move-red-right! node) (node-flip-colors! node) (let ((left (node-left node))) (if (and (not (null? left)) (node-red? (node-left left))) (begin (set! node (node-rotate-right! node)) (node-flip-colors! node)))) node) (define (make-tree cmp) (vector cmp '())) (define (tree-cmp t) (vector-ref t 0)) (define (tree-root t) (vector-ref t 1)) (define (tree-set-root! t root) (vector-set! t 1 root)) ; Search the tree for a given key. Raises 'not-found if the key is not present (use ; tree-get to avoid this) (define (tree-search t key) (define cmp (tree-cmp t)) (define (search-inner node) (if (null? node) (raise 'not-found) (let ((res (cmp key (node-key node)))) (if (= 0 res) (node-val node) (search-inner (if (> 0 res) (node-left node) (node-right node))))))) (search-inner (tree-root t))) ; Check if the tree has a key (define (tree-has? t key) (guard (ex ((eq? ex 'not-found) #f)) (begin (tree-search t key) #t))) ; Get the value of a key from the tree, returning not-found if the key isn't present (define (tree-get t key not-found) (guard (ex ((eq? ex 'not-found) not-found)) (begin (tree-search t key)))) (define (node-fixup! node) (if (and (node-red? (node-right node)) (not (node-red? (node-left node)))) (set! node (node-rotate-left! node))) (if (and (node-red? (node-left node)) (node-red? (node-left (node-left node)))) (set! node (node-rotate-right! node))) (if (and (node-red? (node-left node)) (node-red? (node-right node))) (node-flip-colors! node))1 node) ; Insert an element into the tree (define (tree-insert! t key val) (define cmp (tree-cmp t)) (define (inner-insert node) (cond ((null? node) (make-node key val)) (else (let ((left (node-left node)) (right (node-right node)) (res (cmp key (node-key node)))) (cond ((= res 0) (node-set-val! node val)) ((< res 0) (node-set-left! node (inner-insert (node-left node)))) (else (node-set-right! node (inner-insert (node-right node))))) (node-fixup! node))))) (tree-set-root! t (inner-insert (tree-root t))) (node-set-color! (tree-root t) #f) t) ; Check if a node is not null and red (define (node-red? node) (and (not (null? node)) (node-color node))) ; Delete an element from the tree, returning #t if the element was found (define (tree-delete! t key) (define cmp (tree-cmp t)) (define (delete-min node) (let ((left (node-left node))) (cond ((null? left) '()) (else (if (and (not (node-red? left)) (not (node-red? (node-left left)))) (set! node (node-move-red-left! node))) (node-set-left! node (delete-min (node-left node))) (node-fixup! node))))) (define (node-min node) (if (null? (node-left node)) node (node-min (node-left node)))) (define (inner-delete node) (cond ((< (cmp key (node-key node)) 0) (let ((left (node-left node))) (if (not (null? left)) (begin (if (and (not (node-red? left)) (not (node-red? (node-left left)))) (set! node (node-move-red-left! node))) (let ((res (inner-delete left))) (node-set-left! node res) res)) 'not-found))) (else (if (node-red? (node-left node)) (set! node (node-rotate-right! node))) (if (and (= (cmp key (node-key node)) 0) (null? (node-right node))) (begin (set! node '()) '()) (let ((right (node-right node))) (if (not (null? right)) (begin (if (and (not (node-red? right)) (not (node-red? (node-left right)))) (set! node (node-move-red-right! node))) (if (= (cmp key (node-key node)) 0) (begin (let ((the-min (node-min (node-right node)))) (node-set-val! node (node-val the-min)) (node-set-key! node (node-key the-min)) (delete-min (node-right node)))) (let ((res (inner-delete (node-right node)))) (node-set-right! (node-right node) res) res)) 'not-found)))))) (if (or (null? node) (eq? node 'not-found)) node (node-fixup! node))) (let ((res (inner-delete (tree-root t)))) (tree-set-root! t res) (not (eq? res 'not-found)))) ; Determine if any element of `l` satisfies `pred?` (define (any? pred? l) (cond ((null? l) #f) ((pred? (car l)) #t) (else (any? pred? (cdr l))))) ; This version makes a tree and uses a message-passing style to run functions on it (define (make-tree-disp cmp) (let ((t (make-tree cmp))) (lambda (f . args) (define (match? l) (any? (lambda (i) (eq? f i)) l)) (cond ((match? '(search s)) (apply tree-search (cons t args))) ((match? '(delete d)) (apply tree-delete! (cons t args))) ((match? '(insert i)) (apply tree-insert! (cons t args))) ((match? '(has h in)) (apply tree-has? (cons t args))) ((match? '(get g)) (apply tree-get (cons t args))) (else (error #f "No such function!")))))) )