 | cs-llrb - Chez Scheme implementation of left-leaning red-black trees.
git clone https://benconnors.ca/git-repos/cs-llrb |
llrb.scm (9247B) - raw
1 (library (llrb (1)) 2 (export 3 make-tree 4 tree-insert! 5 tree-delete! 6 tree-search 7 tree-get 8 tree-has? 9 ) 10 (import (chezscheme)) 11 12 (define (make-node key val) 13 (vector #t key val '() '())) 14 15 (define (node-color node) 16 (vector-ref node 0)) 17 (define (node-key node) 18 (vector-ref node 1)) 19 (define (node-val node) 20 (vector-ref node 2)) 21 (define (node-left node) 22 (vector-ref node 3)) 23 (define (node-right node) 24 (vector-ref node 4)) 25 26 (define (node-set-color! node color) 27 (vector-set! node 0 color)) 28 (define (node-set-key! node key) 29 (vector-set! node 1 key)) 30 (define (node-set-val! node val) 31 (vector-set! node 2 val)) 32 (define (node-set-left! node left) 33 (vector-set! node 3 left)) 34 (define (node-set-right! node right) 35 (vector-set! node 4 right)) 36 37 (define (node-flip-colors! node) 38 (define (flip-color! node) 39 (if (not (null? node)) 40 (node-set-color! node (not (node-color node))))) 41 (flip-color! node) 42 (flip-color! (node-left node)) 43 (flip-color! (node-right node))) 44 45 (define (node-rotate-left! node) 46 (define x (node-right node)) 47 (node-set-right! node (node-left x)) 48 (node-set-left! x node) 49 (node-set-color! x (node-color node)) 50 (node-set-color! node #t) 51 x) 52 53 (define (node-rotate-right! node) 54 (define x (node-left node)) 55 (node-set-left! node (node-right x)) 56 (node-set-right! x node) 57 (node-set-color! x (node-color node)) 58 (node-set-color! node #t) 59 x) 60 61 (define (node-move-red-left! node) 62 (node-flip-colors! node) 63 (let ((right (node-right node))) 64 (if (and (not (null? right)) (node-red? (node-left right))) 65 (begin 66 (node-set-right! node (node-rotate-right! right)) 67 (set! node (node-rotate-left! node)) 68 (node-flip-colors! node)))) 69 node) 70 71 (define (node-move-red-right! node) 72 (node-flip-colors! node) 73 (let ((left (node-left node))) 74 (if (and (not (null? left)) (node-red? (node-left left))) 75 (begin 76 (set! node (node-rotate-right! node)) 77 (node-flip-colors! node)))) 78 node) 79 80 (define (make-tree cmp) 81 (vector cmp '())) 82 83 (define (tree-cmp t) 84 (vector-ref t 0)) 85 (define (tree-root t) 86 (vector-ref t 1)) 87 88 (define (tree-set-root! t root) 89 (vector-set! t 1 root)) 90 91 ; Search the tree for a given key. Raises 'not-found if the key is not present (use 92 ; tree-get to avoid this) 93 (define (tree-search t key) 94 (define cmp (tree-cmp t)) 95 (define (search-inner node) 96 (if (null? node) (raise 'not-found) 97 (let ((res (cmp key (node-key node)))) 98 (if (= 0 res) (node-val node) 99 (search-inner 100 (if (> 0 res) (node-left node) 101 (node-right node))))))) 102 (search-inner (tree-root t))) 103 104 ; Check if the tree has a key 105 (define (tree-has? t key) 106 (guard (ex 107 ((eq? ex 'not-found) #f)) 108 (begin 109 (tree-search t key) 110 #t))) 111 112 ; Get the value of a key from the tree, returning not-found if the key isn't present 113 (define (tree-get t key not-found) 114 (guard (ex 115 ((eq? ex 'not-found) not-found)) 116 (begin 117 (tree-search t key)))) 118 119 (define (node-fixup! node) 120 (if (and (node-red? (node-right node)) (not (node-red? (node-left node)))) 121 (set! node (node-rotate-left! node))) 122 (if (and (node-red? (node-left node)) (node-red? (node-left (node-left node)))) 123 (set! node (node-rotate-right! node))) 124 (if (and (node-red? (node-left node)) (node-red? (node-right node))) 125 (node-flip-colors! node))1 126 node) 127 128 ; Insert an element into the tree 129 (define (tree-insert! t key val) 130 (define cmp (tree-cmp t)) 131 (define (inner-insert node) 132 (cond ((null? node) (make-node key val)) 133 (else 134 (let ((left (node-left node)) 135 (right (node-right node)) 136 (res (cmp key (node-key node)))) 137 (cond ((= res 0) (node-set-val! node val)) 138 ((< res 0) (node-set-left! node (inner-insert (node-left node)))) 139 (else (node-set-right! node (inner-insert (node-right node))))) 140 (node-fixup! node))))) 141 (tree-set-root! t (inner-insert (tree-root t))) 142 (node-set-color! (tree-root t) #f) 143 t) 144 145 ; Check if a node is not null and red 146 (define (node-red? node) 147 (and (not (null? node)) (node-color node))) 148 149 ; Delete an element from the tree, returning #t if the element was found 150 (define (tree-delete! t key) 151 (define cmp (tree-cmp t)) 152 (define (delete-min node) 153 (let ((left (node-left node))) 154 (cond ((null? left) '()) 155 (else 156 (if (and (not (node-red? left)) (not (node-red? (node-left left)))) 157 (set! node (node-move-red-left! node))) 158 (node-set-left! node (delete-min (node-left node))) 159 (node-fixup! node))))) 160 (define (node-min node) 161 (if (null? (node-left node)) node 162 (node-min (node-left node)))) 163 (define (inner-delete node) 164 (cond ((< (cmp key (node-key node)) 0) 165 (let ((left (node-left node))) 166 (if (not (null? left)) 167 (begin 168 (if (and (not (node-red? left)) (not (node-red? (node-left left)))) 169 (set! node (node-move-red-left! node))) 170 (let ((res (inner-delete left))) 171 (node-set-left! node res) 172 res)) 173 'not-found))) 174 (else 175 (if (node-red? (node-left node)) 176 (set! node (node-rotate-right! node))) 177 (if (and (= (cmp key (node-key node)) 0) (null? (node-right node))) 178 (begin 179 (set! node '()) 180 '()) 181 (let ((right (node-right node))) 182 (if (not (null? right)) 183 (begin 184 (if (and (not (node-red? right)) (not (node-red? (node-left right)))) 185 (set! node (node-move-red-right! node))) 186 (if (= (cmp key (node-key node)) 0) 187 (begin 188 (let ((the-min (node-min (node-right node)))) 189 (node-set-val! node (node-val the-min)) 190 (node-set-key! node (node-key the-min)) 191 (delete-min (node-right node)))) 192 (let ((res (inner-delete (node-right node)))) 193 (node-set-right! (node-right node) res) 194 res)) 195 'not-found)))))) 196 (if (or (null? node) (eq? node 'not-found)) 197 node 198 (node-fixup! node))) 199 (let ((res (inner-delete (tree-root t)))) 200 (tree-set-root! t res) 201 (not (eq? res 'not-found)))) 202 203 204 ; Determine if any element of `l` satisfies `pred?` 205 (define (any? pred? l) 206 (cond ((null? l) #f) 207 ((pred? (car l)) #t) 208 (else (any? pred? (cdr l))))) 209 210 ; This version makes a tree and uses a message-passing style to run functions on it 211 (define (make-tree-disp cmp) 212 (let ((t (make-tree cmp))) 213 (lambda (f . args) 214 (define (match? l) (any? (lambda (i) (eq? f i)) l)) 215 (cond ((match? '(search s)) (apply tree-search (cons t args))) 216 ((match? '(delete d)) (apply tree-delete! (cons t args))) 217 ((match? '(insert i)) (apply tree-insert! (cons t args))) 218 ((match? '(has h in)) (apply tree-has? (cons t args))) 219 ((match? '(get g)) (apply tree-get (cons t args))) 220 (else (error #f "No such function!")))))) 221 222 )