1use std::cmp::Ordering;
2use std::fmt::Debug;
3
4mod kv;
5
6#[cfg(test)]
7mod tests;
8
9pub trait Node: Sized {
10 type Key: Ord + Debug;
11 type Ptr: NodePtr<Self>;
12
13 fn new(node: &Self) -> Self::Ptr;
14 fn free(&mut self);
15 fn left(&self) -> &Self::Ptr;
16 fn left_mut(&mut self) -> &mut Self::Ptr;
17 fn right(&self) -> &Self::Ptr;
18 fn right_mut(&mut self) -> &mut Self::Ptr;
19 fn key(&self) -> &Self::Key;
20 fn update(&mut self, node: &Self);
21
22 fn is_black(&self) -> bool;
23 fn is_red(&self) -> bool {
24 !self.is_black()
25 }
26
27 fn set_black(&mut self);
28 fn set_red(&mut self);
29}
30
31pub trait NodePtr<N: Node<Ptr = Self>>: Copy {
32 const NIL: Self;
33 fn is_nil(&self) -> bool;
34 fn node<'a>(&self) -> &'a N;
35 fn node_mut<'a>(&self) -> &'a mut N;
36
37 fn is_black(&self) -> bool {
38 self.is_nil() || self.node().is_black()
39 }
40
41 fn is_red(&self) -> bool {
42 !self.is_black()
43 }
44}
45
46struct Context<N: Node> {
47 parent: Option<(*mut Self, bool)>,
48 current: *mut N::Ptr
49}
50
51impl<N: Node> Clone for Context<N> {
52 fn clone(&self) -> Self {
53 Self {
54 parent: self.parent,
55 current: self.current
56 }
57 }
58}
59
60impl<N: Node> Copy for Context<N> {
61}
62
63impl<N: Node> Context<N> {
64 fn ptr(&self) -> &mut N::Ptr {
65 unsafe { &mut *self.current }
66 }
67
68 fn parent_ptr(&self) -> &mut N::Ptr {
69 self.parent_ctx().ptr()
70 }
71
72 fn sibling_ptr(&self) -> &mut N::Ptr {
73 if self.is_left_child() {
74 self.parent_ptr().node_mut().right_mut()
75 } else {
76 self.parent_ptr().node_mut().left_mut()
77 }
78 }
79
80 fn is_root(&self) -> bool {
81 self.parent.is_none()
82 }
83
84 fn is_left_child(&self) -> bool {
85 self.parent.unwrap().1
86 }
87
88 fn parent_ctx(&self) -> &Self {
89 unsafe { &*self.parent.unwrap().0 }
90 }
91
92 fn parent_ctx_mut(&self) -> &mut Self {
93 unsafe { &mut *self.parent.unwrap().0 }
94 }
95
96 fn left_ctx(&mut self) -> Self {
97 Context::<N> {
98 parent: Some((self as *mut Self, true)),
99 current: self.ptr().node_mut().left_mut()
100 }
101 }
102
103 fn right_ctx(&mut self) -> Self {
104 Context {
105 parent: Some((self as *mut Self, false)),
106 current: self.ptr().node_mut().right_mut()
107 }
108 }
109
110 fn has_left(&self) -> bool {
111 !self.ptr().node().left().is_nil()
112 }
113
114 fn has_left_and_right(&self) -> bool {
115 let node = self.ptr().node();
116 !node.left().is_nil() && !node.right().is_nil()
117 }
118}
119
120pub struct RBTree<N: Node> {
121 size: usize,
122 root: N::Ptr
123}
124
125impl<N: Node> RBTree<N> {
126 pub fn new() -> RBTree<N> {
127 RBTree {
128 size: 0,
129 root: N::Ptr::NIL
130 }
131 }
132
133 pub fn size(&self) -> usize {
134 self.size
135 }
136
137 fn root_context(&mut self) -> Context<N> {
138 Context {
139 parent: None,
140 current: &mut self.root as *mut N::Ptr
141 }
142 }
143
144 pub fn search(&self, key: &N::Key) -> Option<&N> {
145 let mut ptr = &self.root;
146 loop {
147 if ptr.is_nil() {
148 return None;
149 }
150 let node = ptr.node();
151 match node.key().cmp(key) {
152 Ordering::Equal => { return Some(node) }
153 Ordering::Less => { ptr = node.right() }
154 Ordering::Greater => { ptr = node.left() }
155 }
156 }
157 }
158
159 pub fn insert(&mut self, node: &N) -> bool {
160 let inserted = Self::do_insert(self.root_context(), node);
161 if inserted {
162 self.size += 1;
163 }
164 return inserted;
165 }
166
167 fn do_insert(mut ctx: Context<N>, node: &N) -> bool {
168 let current_ptr = ctx.ptr();
169 if current_ptr.is_nil() {
170 *current_ptr = N::new(node);
171 if ctx.is_root() {
172 current_ptr.node_mut().set_black();
173 }
174 return true;
175 }
176 let current_node = current_ptr.node_mut();
177 let next_ctx = match current_node.key().cmp(node.key()) {
178 Ordering::Equal => {
179 current_node.update(node);
180 return false;
181 }
182 Ordering::Less => { ctx.right_ctx() }
183 Ordering::Greater => { ctx.left_ctx() }
184 };
185 let inserted = Self::do_insert(next_ctx, node);
186 if inserted && ctx.ptr().node().is_red() {
187 if ctx.is_root() {
188 ctx.ptr().node_mut().set_black();
189 } else if next_ctx.ptr().node().is_red() {
190 Self::insert_repair(ctx, next_ctx.is_left_child())
191 }
192 }
193 inserted
194 }
195
196 fn insert_repair(ctx: Context<N>, inserted_at_left: bool) {
197 if !ctx.sibling_ptr().is_nil() {
198 let sibling_node = ctx.sibling_ptr().node_mut();
199 if sibling_node.is_red() {
200 ctx.ptr().node_mut().set_black();
201 sibling_node.set_black();
202 ctx.parent_ptr().node_mut().set_red();
203 return;
204 }
205 }
206 if ctx.is_left_child() {
207 if !inserted_at_left {
208 Self::rotate_left(ctx.ptr());
209 }
210 ctx.ptr().node_mut().set_black();
211 ctx.parent_ptr().node_mut().set_red();
212 Self::rotate_right(ctx.parent_ptr());
213 } else {
214 if inserted_at_left {
215 Self::rotate_right(ctx.ptr());
216 }
217 ctx.ptr().node_mut().set_black();
218 ctx.parent_ptr().node_mut().set_red();
219 Self::rotate_left(ctx.parent_ptr());
220 }
221 }
222
223 pub fn delete(&mut self, key: &N::Key) -> bool {
224 let mut deleted_node: N::Ptr = N::Ptr::NIL;
225 Self::do_delete(self.root_context(), key, &mut deleted_node);
226 let deleted = !deleted_node.is_nil();
227 if deleted {
228 deleted_node.node_mut().free();
229 self.size -= 1;
230 }
231 return deleted;
232 }
233
234 fn do_delete(mut ctx: Context<N>, key: &N::Key, deleted_node: &mut N::Ptr) -> bool {
235 let current_ptr = ctx.ptr();
236 if current_ptr.is_nil() {
237 return false;
238 }
239 let current_node = current_ptr.node_mut();
240 let next_ctx = match current_node.key().cmp(key) {
241 Ordering::Equal => {
242 if ctx.has_left_and_right() {
243 (*deleted_node).clone_from(ctx.ptr());
244 let need_repair = Self::delete_left_most(ctx.right_ctx(), deleted_node);
245 return need_repair && Self::delete_repair(ctx);
246 } else {
247 return Self::delete_node(ctx, deleted_node);
248 };
249 }
250 Ordering::Less => { ctx.right_ctx() }
251 Ordering::Greater => { ctx.left_ctx() }
252 };
253 Self::do_delete(next_ctx, key, deleted_node) && Self::delete_repair(ctx)
254 }
255
256 fn delete_left_most<'a>(mut ctx: Context<N>, deleted_node: &mut N::Ptr) -> bool {
257 if ctx.has_left() {
258 Self::delete_left_most(ctx.left_ctx(), deleted_node) && Self::delete_repair(ctx)
259 } else {
260 (*deleted_node).node_mut().update(ctx.ptr().node());
261 Self::delete_node(ctx, deleted_node)
262 }
263 }
264
265 fn delete_node<'a>(ctx: Context<N>, deleted_node: &mut N::Ptr) -> bool {
266 *deleted_node = *ctx.ptr();
267 let n = deleted_node.node_mut();
268 let c = if !n.left().is_nil() { n.left() } else { n.right() };
269 ctx.ptr().clone_from(c);
270
271 if n.is_red() {
272 return false;
273 }
274
275 if c.is_red() {
276 c.node_mut().set_black();
277 return false;
278 }
279
280 return Self::delete_repair(ctx);
281 }
282
283 fn delete_repair(mut ctx: Context<N>) -> bool {
284 if ctx.is_root() {
285 return false;
286 }
287
288 let p = ctx.parent_ptr();
289 let s = ctx.sibling_ptr();
290 if s.is_red() {
291 p.node_mut().set_red();
292 s.node_mut().set_black();
293 if ctx.is_left_child() {
294 Self::rotate_left(p);
295 ctx = ctx.parent_ctx_mut().left_ctx().left_ctx();
296 } else {
297 Self::rotate_right(p);
298 ctx = ctx.parent_ctx_mut().right_ctx().right_ctx();
299 }
300 }
301
302 let p = ctx.parent_ptr();
303 let s = ctx.sibling_ptr();
304 if s.is_black() && s.node().left().is_black() && s.node().right().is_black() {
305 if p.is_black() {
306 s.node_mut().set_red();
307 return true;
308 }
309 s.node_mut().set_red();
310 p.node_mut().set_black();
311 return false;
312 }
313
314 if ctx.is_left_child() {
315 if s.node().right().is_black() {
316 s.node_mut().set_red();
317 s.node().left().node_mut().set_black();
318 Self::rotate_right(s);
319 }
320 } else {
321 if s.node().left().is_black() {
322 s.node_mut().set_red();
323 s.node().right().node_mut().set_black();
324 Self::rotate_left(s);
325 }
326 }
327
328 if p.node().is_red() {
329 s.node_mut().set_red();
330 p.node_mut().set_black();
331 }
332 if ctx.is_left_child() {
333 s.node_mut().right_mut().node_mut().set_black();
334 Self::rotate_left(p);
335 } else {
336 s.node_mut().left_mut().node_mut().set_black();
337 Self::rotate_right(p);
338 }
339 return false;
340 }
341
342 fn rotate_left(ptr: &mut N::Ptr) {
343 let me = ptr.clone();
344 let r = me.node().right().clone();
345 let rl = r.node().left().clone();
346 *ptr = r;
347 *me.node_mut().right_mut() = rl;
348 *r.node_mut().left_mut() = me;
349 }
350
351 fn rotate_right(ptr: &mut N::Ptr) {
352 let me = ptr.clone();
353 let l = me.node().left().clone();
354 let lr = l.node().right().clone();
355 *ptr = l;
356 *me.node_mut().left_mut() = lr;
357 *l.node_mut().right_mut() = me;
358 }
359}
360
361pub type KeyValue<K, V> = kv::KeyValue<K, V>;