]>
Commit | Line | Data |
---|---|---|
223e47cc LB |
1 | // Copyright 2013 The Rust Project Developers. See the COPYRIGHT |
2 | // file at the top-level directory of this distribution and at | |
3 | // http://rust-lang.org/COPYRIGHT. | |
4 | // | |
5 | // Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or | |
6 | // http://www.apache.org/licenses/LICENSE-2.0> or the MIT license | |
7 | // <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your | |
8 | // option. This file may not be copied, modified, or distributed | |
9 | // except according to those terms. | |
10 | ||
11 | //! An ordered map and set implemented as self-balancing binary search | |
12 | //! trees. The only requirement for the types is that the key implements | |
13 | //! `TotalOrd`. | |
14 | ||
15 | use core::prelude::*; | |
16 | ||
17 | // This is implemented as an AA tree, which is a simplified variation of | |
18 | // a red-black tree where where red (horizontal) nodes can only be added | |
19 | // as a right child. The time complexity is the same, and re-balancing | |
20 | // operations are more frequent but also cheaper. | |
21 | ||
22 | // Future improvements: | |
23 | ||
24 | // range search - O(log n) retrieval of an iterator from some key | |
25 | ||
26 | // (possibly) implement the overloads Python does for sets: | |
27 | // * intersection: & | |
28 | // * difference: - | |
29 | // * symmetric difference: ^ | |
30 | // * union: | | |
31 | // These would be convenient since the methods work like `each` | |
32 | ||
33 | pub struct TreeMap<K, V> { | |
34 | priv root: Option<~TreeNode<K, V>>, | |
35 | priv length: uint | |
36 | } | |
37 | ||
38 | impl<K: Eq + TotalOrd, V: Eq> Eq for TreeMap<K, V> { | |
39 | fn eq(&self, other: &TreeMap<K, V>) -> bool { | |
40 | if self.len() != other.len() { | |
41 | false | |
42 | } else { | |
43 | let mut x = self.iter(); | |
44 | let mut y = other.iter(); | |
45 | for self.len().times { | |
46 | if map_next(&mut x).unwrap() != | |
47 | map_next(&mut y).unwrap() { | |
48 | return false | |
49 | } | |
50 | } | |
51 | true | |
52 | } | |
53 | } | |
54 | fn ne(&self, other: &TreeMap<K, V>) -> bool { !self.eq(other) } | |
55 | } | |
56 | ||
57 | // Lexicographical comparison | |
58 | fn lt<K: Ord + TotalOrd, V>(a: &TreeMap<K, V>, | |
59 | b: &TreeMap<K, V>) -> bool { | |
60 | let mut x = a.iter(); | |
61 | let mut y = b.iter(); | |
62 | ||
63 | let (a_len, b_len) = (a.len(), b.len()); | |
64 | for uint::min(a_len, b_len).times { | |
65 | let (key_a,_) = map_next(&mut x).unwrap(); | |
66 | let (key_b,_) = map_next(&mut y).unwrap(); | |
67 | if *key_a < *key_b { return true; } | |
68 | if *key_a > *key_b { return false; } | |
69 | }; | |
70 | ||
71 | a_len < b_len | |
72 | } | |
73 | ||
74 | impl<K: Ord + TotalOrd, V> Ord for TreeMap<K, V> { | |
75 | #[inline(always)] | |
76 | fn lt(&self, other: &TreeMap<K, V>) -> bool { lt(self, other) } | |
77 | #[inline(always)] | |
78 | fn le(&self, other: &TreeMap<K, V>) -> bool { !lt(other, self) } | |
79 | #[inline(always)] | |
80 | fn ge(&self, other: &TreeMap<K, V>) -> bool { !lt(self, other) } | |
81 | #[inline(always)] | |
82 | fn gt(&self, other: &TreeMap<K, V>) -> bool { lt(other, self) } | |
83 | } | |
84 | ||
85 | impl<'self, K: TotalOrd, V> BaseIter<(&'self K, &'self V)> for TreeMap<K, V> { | |
86 | /// Visit all key-value pairs in order | |
87 | fn each(&self, f: &fn(&(&'self K, &'self V)) -> bool) { | |
88 | each(&self.root, f) | |
89 | } | |
90 | fn size_hint(&self) -> Option<uint> { Some(self.len()) } | |
91 | } | |
92 | ||
93 | impl<'self, K: TotalOrd, V> | |
94 | ReverseIter<(&'self K, &'self V)> | |
95 | for TreeMap<K, V> | |
96 | { | |
97 | /// Visit all key-value pairs in reverse order | |
98 | fn each_reverse(&self, f: &fn(&(&'self K, &'self V)) -> bool) { | |
99 | each_reverse(&self.root, f); | |
100 | } | |
101 | } | |
102 | ||
103 | impl<K: TotalOrd, V> Container for TreeMap<K, V> { | |
104 | /// Return the number of elements in the map | |
105 | fn len(&const self) -> uint { self.length } | |
106 | ||
107 | /// Return true if the map contains no elements | |
108 | fn is_empty(&const self) -> bool { self.root.is_none() } | |
109 | } | |
110 | ||
111 | impl<K: TotalOrd, V> Mutable for TreeMap<K, V> { | |
112 | /// Clear the map, removing all key-value pairs. | |
113 | fn clear(&mut self) { | |
114 | self.root = None; | |
115 | self.length = 0 | |
116 | } | |
117 | } | |
118 | ||
119 | impl<K: TotalOrd, V> Map<K, V> for TreeMap<K, V> { | |
120 | /// Return true if the map contains a value for the specified key | |
121 | fn contains_key(&self, key: &K) -> bool { | |
122 | self.find(key).is_some() | |
123 | } | |
124 | ||
125 | /// Visit all keys in order | |
126 | fn each_key(&self, f: &fn(&K) -> bool) { self.each(|&(k, _)| f(k)) } | |
127 | ||
128 | /// Visit all values in order | |
129 | fn each_value(&self, f: &fn(&V) -> bool) { | |
130 | self.each(|&(_, v)| f(v)) | |
131 | } | |
132 | ||
133 | /// Iterate over the map and mutate the contained values | |
134 | fn mutate_values(&mut self, f: &fn(&'self K, &'self mut V) -> bool) { | |
135 | mutate_values(&mut self.root, f); | |
136 | } | |
137 | ||
138 | /// Return a reference to the value corresponding to the key | |
139 | fn find(&self, key: &K) -> Option<&'self V> { | |
140 | let mut current: &'self Option<~TreeNode<K, V>> = &self.root; | |
141 | loop { | |
142 | match *current { | |
143 | Some(ref r) => { | |
144 | match key.cmp(&r.key) { | |
145 | Less => current = &r.left, | |
146 | Greater => current = &r.right, | |
147 | Equal => return Some(&r.value) | |
148 | } | |
149 | } | |
150 | None => return None | |
151 | } | |
152 | } | |
153 | } | |
154 | ||
155 | /// Return a mutable reference to the value corresponding to the key | |
156 | #[inline(always)] | |
157 | fn find_mut(&mut self, key: &K) -> Option<&'self mut V> { | |
158 | find_mut(&mut self.root, key) | |
159 | } | |
160 | ||
161 | /// Insert a key-value pair into the map. An existing value for a | |
162 | /// key is replaced by the new value. Return true if the key did | |
163 | /// not already exist in the map. | |
164 | fn insert(&mut self, key: K, value: V) -> bool { | |
165 | let ret = insert(&mut self.root, key, value); | |
166 | if ret { self.length += 1 } | |
167 | ret | |
168 | } | |
169 | ||
170 | /// Remove a key-value pair from the map. Return true if the key | |
171 | /// was present in the map, otherwise false. | |
172 | fn remove(&mut self, key: &K) -> bool { | |
173 | let ret = remove(&mut self.root, key); | |
174 | if ret { self.length -= 1 } | |
175 | ret | |
176 | } | |
177 | } | |
178 | ||
179 | pub impl<K: TotalOrd, V> TreeMap<K, V> { | |
180 | /// Create an empty TreeMap | |
181 | fn new() -> TreeMap<K, V> { TreeMap{root: None, length: 0} } | |
182 | ||
183 | /// Visit all keys in reverse order | |
184 | fn each_key_reverse(&self, f: &fn(&K) -> bool) { | |
185 | self.each_reverse(|&(k, _)| f(k)) | |
186 | } | |
187 | ||
188 | /// Visit all values in reverse order | |
189 | fn each_value_reverse(&self, f: &fn(&V) -> bool) { | |
190 | self.each_reverse(|&(_, v)| f(v)) | |
191 | } | |
192 | ||
193 | /// Get a lazy iterator over the key-value pairs in the map. | |
194 | /// Requires that it be frozen (immutable). | |
195 | fn iter(&self) -> TreeMapIterator<'self, K, V> { | |
196 | TreeMapIterator{stack: ~[], node: &self.root} | |
197 | } | |
198 | } | |
199 | ||
200 | /// Lazy forward iterator over a map | |
201 | pub struct TreeMapIterator<'self, K, V> { | |
202 | priv stack: ~[&'self ~TreeNode<K, V>], | |
203 | priv node: &'self Option<~TreeNode<K, V>> | |
204 | } | |
205 | ||
206 | /// Advance the iterator to the next node (in order) and return a | |
207 | /// tuple with a reference to the key and value. If there are no | |
208 | /// more nodes, return `None`. | |
209 | pub fn map_next<'r, K, V>(iter: &mut TreeMapIterator<'r, K, V>) | |
210 | -> Option<(&'r K, &'r V)> { | |
211 | while !iter.stack.is_empty() || iter.node.is_some() { | |
212 | match *iter.node { | |
213 | Some(ref x) => { | |
214 | iter.stack.push(x); | |
215 | iter.node = &x.left; | |
216 | } | |
217 | None => { | |
218 | let res = iter.stack.pop(); | |
219 | iter.node = &res.right; | |
220 | return Some((&res.key, &res.value)); | |
221 | } | |
222 | } | |
223 | } | |
224 | None | |
225 | } | |
226 | ||
227 | /// Advance the iterator through the map | |
228 | pub fn map_advance<'r, K, V>(iter: &mut TreeMapIterator<'r, K, V>, | |
229 | f: &fn((&'r K, &'r V)) -> bool) { | |
230 | loop { | |
231 | match map_next(iter) { | |
232 | Some(x) => { | |
233 | if !f(x) { return } | |
234 | } | |
235 | None => return | |
236 | } | |
237 | } | |
238 | } | |
239 | ||
240 | pub struct TreeSet<T> { | |
241 | priv map: TreeMap<T, ()> | |
242 | } | |
243 | ||
244 | impl<T: TotalOrd> BaseIter<T> for TreeSet<T> { | |
245 | /// Visit all values in order | |
246 | #[inline(always)] | |
247 | fn each(&self, f: &fn(&T) -> bool) { self.map.each_key(f) } | |
248 | #[inline(always)] | |
249 | fn size_hint(&self) -> Option<uint> { Some(self.len()) } | |
250 | } | |
251 | ||
252 | impl<T: TotalOrd> ReverseIter<T> for TreeSet<T> { | |
253 | /// Visit all values in reverse order | |
254 | #[inline(always)] | |
255 | fn each_reverse(&self, f: &fn(&T) -> bool) { | |
256 | self.map.each_key_reverse(f) | |
257 | } | |
258 | } | |
259 | ||
260 | impl<T: Eq + TotalOrd> Eq for TreeSet<T> { | |
261 | #[inline(always)] | |
262 | fn eq(&self, other: &TreeSet<T>) -> bool { self.map == other.map } | |
263 | #[inline(always)] | |
264 | fn ne(&self, other: &TreeSet<T>) -> bool { self.map != other.map } | |
265 | } | |
266 | ||
267 | impl<T: Ord + TotalOrd> Ord for TreeSet<T> { | |
268 | #[inline(always)] | |
269 | fn lt(&self, other: &TreeSet<T>) -> bool { self.map < other.map } | |
270 | #[inline(always)] | |
271 | fn le(&self, other: &TreeSet<T>) -> bool { self.map <= other.map } | |
272 | #[inline(always)] | |
273 | fn ge(&self, other: &TreeSet<T>) -> bool { self.map >= other.map } | |
274 | #[inline(always)] | |
275 | fn gt(&self, other: &TreeSet<T>) -> bool { self.map > other.map } | |
276 | } | |
277 | ||
278 | impl<T: TotalOrd> Container for TreeSet<T> { | |
279 | /// Return the number of elements in the set | |
280 | #[inline(always)] | |
281 | fn len(&const self) -> uint { self.map.len() } | |
282 | ||
283 | /// Return true if the set contains no elements | |
284 | #[inline(always)] | |
285 | fn is_empty(&const self) -> bool { self.map.is_empty() } | |
286 | } | |
287 | ||
288 | impl<T: TotalOrd> Mutable for TreeSet<T> { | |
289 | /// Clear the set, removing all values. | |
290 | #[inline(always)] | |
291 | fn clear(&mut self) { self.map.clear() } | |
292 | } | |
293 | ||
294 | impl<T: TotalOrd> Set<T> for TreeSet<T> { | |
295 | /// Return true if the set contains a value | |
296 | #[inline(always)] | |
297 | fn contains(&self, value: &T) -> bool { | |
298 | self.map.contains_key(value) | |
299 | } | |
300 | ||
301 | /// Add a value to the set. Return true if the value was not already | |
302 | /// present in the set. | |
303 | #[inline(always)] | |
304 | fn insert(&mut self, value: T) -> bool { self.map.insert(value, ()) } | |
305 | ||
306 | /// Remove a value from the set. Return true if the value was | |
307 | /// present in the set. | |
308 | #[inline(always)] | |
309 | fn remove(&mut self, value: &T) -> bool { self.map.remove(value) } | |
310 | ||
311 | /// Return true if the set has no elements in common with `other`. | |
312 | /// This is equivalent to checking for an empty intersection. | |
313 | fn is_disjoint(&self, other: &TreeSet<T>) -> bool { | |
314 | let mut x = self.iter(); | |
315 | let mut y = other.iter(); | |
316 | let mut a = set_next(&mut x); | |
317 | let mut b = set_next(&mut y); | |
318 | while a.is_some() && b.is_some() { | |
319 | let a1 = a.unwrap(); | |
320 | let b1 = b.unwrap(); | |
321 | match a1.cmp(b1) { | |
322 | Less => a = set_next(&mut x), | |
323 | Greater => b = set_next(&mut y), | |
324 | Equal => return false | |
325 | } | |
326 | } | |
327 | true | |
328 | } | |
329 | ||
330 | /// Return true if the set is a subset of another | |
331 | #[inline(always)] | |
332 | fn is_subset(&self, other: &TreeSet<T>) -> bool { | |
333 | other.is_superset(self) | |
334 | } | |
335 | ||
336 | /// Return true if the set is a superset of another | |
337 | fn is_superset(&self, other: &TreeSet<T>) -> bool { | |
338 | let mut x = self.iter(); | |
339 | let mut y = other.iter(); | |
340 | let mut a = set_next(&mut x); | |
341 | let mut b = set_next(&mut y); | |
342 | while b.is_some() { | |
343 | if a.is_none() { | |
344 | return false | |
345 | } | |
346 | ||
347 | let a1 = a.unwrap(); | |
348 | let b1 = b.unwrap(); | |
349 | ||
350 | match a1.cmp(b1) { | |
351 | Less => (), | |
352 | Greater => return false, | |
353 | Equal => b = set_next(&mut y), | |
354 | } | |
355 | ||
356 | a = set_next(&mut x); | |
357 | } | |
358 | true | |
359 | } | |
360 | ||
361 | /// Visit the values (in-order) representing the difference | |
362 | fn difference(&self, other: &TreeSet<T>, f: &fn(&T) -> bool) { | |
363 | let mut x = self.iter(); | |
364 | let mut y = other.iter(); | |
365 | ||
366 | let mut a = set_next(&mut x); | |
367 | let mut b = set_next(&mut y); | |
368 | ||
369 | while a.is_some() { | |
370 | if b.is_none() { | |
371 | return do a.while_some() |a1| { | |
372 | if f(a1) { set_next(&mut x) } else { None } | |
373 | } | |
374 | } | |
375 | ||
376 | let a1 = a.unwrap(); | |
377 | let b1 = b.unwrap(); | |
378 | ||
379 | let cmp = a1.cmp(b1); | |
380 | ||
381 | if cmp == Less { | |
382 | if !f(a1) { return } | |
383 | a = set_next(&mut x); | |
384 | } else { | |
385 | if cmp == Equal { a = set_next(&mut x) } | |
386 | b = set_next(&mut y); | |
387 | } | |
388 | } | |
389 | } | |
390 | ||
391 | /// Visit the values (in-order) representing the symmetric difference | |
392 | fn symmetric_difference(&self, other: &TreeSet<T>, | |
393 | f: &fn(&T) -> bool) { | |
394 | let mut x = self.iter(); | |
395 | let mut y = other.iter(); | |
396 | ||
397 | let mut a = set_next(&mut x); | |
398 | let mut b = set_next(&mut y); | |
399 | ||
400 | while a.is_some() { | |
401 | if b.is_none() { | |
402 | return do a.while_some() |a1| { | |
403 | if f(a1) { set_next(&mut x) } else { None } | |
404 | } | |
405 | } | |
406 | ||
407 | let a1 = a.unwrap(); | |
408 | let b1 = b.unwrap(); | |
409 | ||
410 | let cmp = a1.cmp(b1); | |
411 | ||
412 | if cmp == Less { | |
413 | if !f(a1) { return } | |
414 | a = set_next(&mut x); | |
415 | } else { | |
416 | if cmp == Greater { | |
417 | if !f(b1) { return } | |
418 | } else { | |
419 | a = set_next(&mut x); | |
420 | } | |
421 | b = set_next(&mut y); | |
422 | } | |
423 | } | |
424 | do b.while_some |b1| { | |
425 | if f(b1) { set_next(&mut y) } else { None } | |
426 | } | |
427 | } | |
428 | ||
429 | /// Visit the values (in-order) representing the intersection | |
430 | fn intersection(&self, other: &TreeSet<T>, f: &fn(&T) -> bool) { | |
431 | let mut x = self.iter(); | |
432 | let mut y = other.iter(); | |
433 | ||
434 | let mut a = set_next(&mut x); | |
435 | let mut b = set_next(&mut y); | |
436 | ||
437 | while a.is_some() && b.is_some() { | |
438 | let a1 = a.unwrap(); | |
439 | let b1 = b.unwrap(); | |
440 | ||
441 | let cmp = a1.cmp(b1); | |
442 | ||
443 | if cmp == Less { | |
444 | a = set_next(&mut x); | |
445 | } else { | |
446 | if cmp == Equal { | |
447 | if !f(a1) { return } | |
448 | } | |
449 | b = set_next(&mut y); | |
450 | } | |
451 | } | |
452 | } | |
453 | ||
454 | /// Visit the values (in-order) representing the union | |
455 | fn union(&self, other: &TreeSet<T>, f: &fn(&T) -> bool) { | |
456 | let mut x = self.iter(); | |
457 | let mut y = other.iter(); | |
458 | ||
459 | let mut a = set_next(&mut x); | |
460 | let mut b = set_next(&mut y); | |
461 | ||
462 | while a.is_some() { | |
463 | if b.is_none() { | |
464 | return do a.while_some() |a1| { | |
465 | if f(a1) { set_next(&mut x) } else { None } | |
466 | } | |
467 | } | |
468 | ||
469 | let a1 = a.unwrap(); | |
470 | let b1 = b.unwrap(); | |
471 | ||
472 | let cmp = a1.cmp(b1); | |
473 | ||
474 | if cmp == Greater { | |
475 | if !f(b1) { return } | |
476 | b = set_next(&mut y); | |
477 | } else { | |
478 | if !f(a1) { return } | |
479 | if cmp == Equal { | |
480 | b = set_next(&mut y); | |
481 | } | |
482 | a = set_next(&mut x); | |
483 | } | |
484 | } | |
485 | do b.while_some |b1| { | |
486 | if f(b1) { set_next(&mut y) } else { None } | |
487 | } | |
488 | } | |
489 | } | |
490 | ||
491 | pub impl <T: TotalOrd> TreeSet<T> { | |
492 | /// Create an empty TreeSet | |
493 | #[inline(always)] | |
494 | fn new() -> TreeSet<T> { TreeSet{map: TreeMap::new()} } | |
495 | ||
496 | /// Get a lazy iterator over the values in the set. | |
497 | /// Requires that it be frozen (immutable). | |
498 | #[inline(always)] | |
499 | fn iter(&self) -> TreeSetIterator<'self, T> { | |
500 | TreeSetIterator{iter: self.map.iter()} | |
501 | } | |
502 | } | |
503 | ||
504 | /// Lazy forward iterator over a set | |
505 | pub struct TreeSetIterator<'self, T> { | |
506 | priv iter: TreeMapIterator<'self, T, ()> | |
507 | } | |
508 | ||
509 | /// Advance the iterator to the next node (in order). If this iterator is | |
510 | /// finished, does nothing. | |
511 | #[inline(always)] | |
512 | pub fn set_next<'r, T>(iter: &mut TreeSetIterator<'r, T>) -> Option<&'r T> { | |
513 | do map_next(&mut iter.iter).map |&(value, _)| { value } | |
514 | } | |
515 | ||
516 | /// Advance the iterator through the set | |
517 | #[inline(always)] | |
518 | pub fn set_advance<'r, T>(iter: &mut TreeSetIterator<'r, T>, | |
519 | f: &fn(&'r T) -> bool) { | |
520 | do map_advance(&mut iter.iter) |(k, _)| { f(k) } | |
521 | } | |
522 | ||
523 | // Nodes keep track of their level in the tree, starting at 1 in the | |
524 | // leaves and with a red child sharing the level of the parent. | |
525 | struct TreeNode<K, V> { | |
526 | key: K, | |
527 | value: V, | |
528 | left: Option<~TreeNode<K, V>>, | |
529 | right: Option<~TreeNode<K, V>>, | |
530 | level: uint | |
531 | } | |
532 | ||
533 | pub impl<K: TotalOrd, V> TreeNode<K, V> { | |
534 | #[inline(always)] | |
535 | fn new(key: K, value: V) -> TreeNode<K, V> { | |
536 | TreeNode{key: key, value: value, left: None, right: None, level: 1} | |
537 | } | |
538 | } | |
539 | ||
540 | fn each<'r, K: TotalOrd, V>(node: &'r Option<~TreeNode<K, V>>, | |
541 | f: &fn(&(&'r K, &'r V)) -> bool) { | |
542 | for node.each |x| { | |
543 | each(&x.left, f); | |
544 | if f(&(&x.key, &x.value)) { each(&x.right, f) } | |
545 | } | |
546 | } | |
547 | ||
548 | fn each_reverse<'r, K: TotalOrd, V>(node: &'r Option<~TreeNode<K, V>>, | |
549 | f: &fn(&(&'r K, &'r V)) -> bool) { | |
550 | for node.each |x| { | |
551 | each_reverse(&x.right, f); | |
552 | if f(&(&x.key, &x.value)) { each_reverse(&x.left, f) } | |
553 | } | |
554 | } | |
555 | ||
556 | fn mutate_values<'r, K: TotalOrd, V>(node: &'r mut Option<~TreeNode<K, V>>, | |
557 | f: &fn(&'r K, &'r mut V) -> bool) | |
558 | -> bool { | |
559 | match *node { | |
560 | Some(~TreeNode{key: ref key, value: ref mut value, left: ref mut left, | |
561 | right: ref mut right, _}) => { | |
562 | if !mutate_values(left, f) { return false } | |
563 | if !f(key, value) { return false } | |
564 | if !mutate_values(right, f) { return false } | |
565 | } | |
566 | None => return false | |
567 | } | |
568 | true | |
569 | } | |
570 | ||
571 | // Remove left horizontal link by rotating right | |
572 | fn skew<K: TotalOrd, V>(node: &mut ~TreeNode<K, V>) { | |
573 | if node.left.map_default(false, |x| x.level == node.level) { | |
574 | let mut save = node.left.swap_unwrap(); | |
575 | node.left <-> save.right; // save.right now None | |
576 | *node <-> save; | |
577 | node.right = Some(save); | |
578 | } | |
579 | } | |
580 | ||
581 | // Remove dual horizontal link by rotating left and increasing level of | |
582 | // the parent | |
583 | fn split<K: TotalOrd, V>(node: &mut ~TreeNode<K, V>) { | |
584 | if node.right.map_default(false, | |
585 | |x| x.right.map_default(false, |y| y.level == node.level)) { | |
586 | let mut save = node.right.swap_unwrap(); | |
587 | node.right <-> save.left; // save.left now None | |
588 | save.level += 1; | |
589 | *node <-> save; | |
590 | node.left = Some(save); | |
591 | } | |
592 | } | |
593 | ||
594 | fn find_mut<'r, K: TotalOrd, V>(node: &'r mut Option<~TreeNode<K, V>>, | |
595 | key: &K) | |
596 | -> Option<&'r mut V> { | |
597 | match *node { | |
598 | Some(ref mut x) => { | |
599 | match key.cmp(&x.key) { | |
600 | Less => find_mut(&mut x.left, key), | |
601 | Greater => find_mut(&mut x.right, key), | |
602 | Equal => Some(&mut x.value), | |
603 | } | |
604 | } | |
605 | None => None | |
606 | } | |
607 | } | |
608 | ||
609 | fn insert<K: TotalOrd, V>(node: &mut Option<~TreeNode<K, V>>, key: K, value: V) -> bool { | |
610 | match *node { | |
611 | Some(ref mut save) => { | |
612 | match key.cmp(&save.key) { | |
613 | Less => { | |
614 | let inserted = insert(&mut save.left, key, value); | |
615 | skew(save); | |
616 | split(save); | |
617 | inserted | |
618 | } | |
619 | Greater => { | |
620 | let inserted = insert(&mut save.right, key, value); | |
621 | skew(save); | |
622 | split(save); | |
623 | inserted | |
624 | } | |
625 | Equal => { | |
626 | save.key = key; | |
627 | save.value = value; | |
628 | false | |
629 | } | |
630 | } | |
631 | } | |
632 | None => { | |
633 | *node = Some(~TreeNode::new(key, value)); | |
634 | true | |
635 | } | |
636 | } | |
637 | } | |
638 | ||
639 | fn remove<K: TotalOrd, V>(node: &mut Option<~TreeNode<K, V>>, | |
640 | key: &K) -> bool { | |
641 | fn heir_swap<K: TotalOrd, V>(node: &mut ~TreeNode<K, V>, | |
642 | child: &mut Option<~TreeNode<K, V>>) { | |
643 | // *could* be done without recursion, but it won't borrow check | |
644 | for child.each_mut |x| { | |
645 | if x.right.is_some() { | |
646 | heir_swap(node, &mut x.right); | |
647 | } else { | |
648 | node.key <-> x.key; | |
649 | node.value <-> x.value; | |
650 | } | |
651 | } | |
652 | } | |
653 | ||
654 | match *node { | |
655 | None => { | |
656 | return false // bottom of tree | |
657 | } | |
658 | Some(ref mut save) => { | |
659 | let (removed, this) = match key.cmp(&save.key) { | |
660 | Less => (remove(&mut save.left, key), false), | |
661 | Greater => (remove(&mut save.right, key), false), | |
662 | Equal => { | |
663 | if save.left.is_some() { | |
664 | if save.right.is_some() { | |
665 | let mut left = save.left.swap_unwrap(); | |
666 | if left.right.is_some() { | |
667 | heir_swap(save, &mut left.right); | |
668 | } else { | |
669 | save.key <-> left.key; | |
670 | save.value <-> left.value; | |
671 | } | |
672 | save.left = Some(left); | |
673 | remove(&mut save.left, key); | |
674 | } else { | |
675 | *save = save.left.swap_unwrap(); | |
676 | } | |
677 | (true, false) | |
678 | } else if save.right.is_some() { | |
679 | *save = save.right.swap_unwrap(); | |
680 | (true, false) | |
681 | } else { | |
682 | (true, true) | |
683 | } | |
684 | } | |
685 | }; | |
686 | ||
687 | if !this { | |
688 | let left_level = save.left.map_default(0, |x| x.level); | |
689 | let right_level = save.right.map_default(0, |x| x.level); | |
690 | ||
691 | // re-balance, if necessary | |
692 | if left_level < save.level - 1 || right_level < save.level - 1 { | |
693 | save.level -= 1; | |
694 | ||
695 | if right_level > save.level { | |
696 | for save.right.each_mut |x| { x.level = save.level } | |
697 | } | |
698 | ||
699 | skew(save); | |
700 | ||
701 | for save.right.each_mut |right| { | |
702 | skew(right); | |
703 | for right.right.each_mut |x| { skew(x) } | |
704 | } | |
705 | ||
706 | split(save); | |
707 | for save.right.each_mut |x| { split(x) } | |
708 | } | |
709 | ||
710 | return removed; | |
711 | } | |
712 | } | |
713 | } | |
714 | ||
715 | *node = None; | |
716 | true | |
717 | } | |
718 | ||
719 | #[cfg(test)] | |
720 | mod test_treemap { | |
721 | use core::prelude::*; | |
722 | use super::*; | |
723 | use core::rand::RngUtil; | |
724 | use core::rand; | |
725 | ||
726 | #[test] | |
727 | fn find_empty() { | |
728 | let m = TreeMap::new::<int, int>(); assert!(m.find(&5) == None); | |
729 | } | |
730 | ||
731 | #[test] | |
732 | fn find_not_found() { | |
733 | let mut m = TreeMap::new(); | |
734 | assert!(m.insert(1, 2)); | |
735 | assert!(m.insert(5, 3)); | |
736 | assert!(m.insert(9, 3)); | |
737 | assert!(m.find(&2) == None); | |
738 | } | |
739 | ||
740 | #[test] | |
741 | fn test_find_mut() { | |
742 | let mut m = TreeMap::new(); | |
743 | assert!(m.insert(1, 12)); | |
744 | assert!(m.insert(2, 8)); | |
745 | assert!(m.insert(5, 14)); | |
746 | let new = 100; | |
747 | match m.find_mut(&5) { | |
748 | None => fail!(), Some(x) => *x = new | |
749 | } | |
750 | assert_eq!(m.find(&5), Some(&new)); | |
751 | } | |
752 | ||
753 | #[test] | |
754 | fn insert_replace() { | |
755 | let mut m = TreeMap::new(); | |
756 | assert!(m.insert(5, 2)); | |
757 | assert!(m.insert(2, 9)); | |
758 | assert!(!m.insert(2, 11)); | |
759 | assert!(m.find(&2).unwrap() == &11); | |
760 | } | |
761 | ||
762 | #[test] | |
763 | fn test_clear() { | |
764 | let mut m = TreeMap::new(); | |
765 | m.clear(); | |
766 | assert!(m.insert(5, 11)); | |
767 | assert!(m.insert(12, -3)); | |
768 | assert!(m.insert(19, 2)); | |
769 | m.clear(); | |
770 | assert!(m.find(&5).is_none()); | |
771 | assert!(m.find(&12).is_none()); | |
772 | assert!(m.find(&19).is_none()); | |
773 | assert!(m.is_empty()); | |
774 | } | |
775 | ||
776 | #[test] | |
777 | fn u8_map() { | |
778 | let mut m = TreeMap::new(); | |
779 | ||
780 | let k1 = str::to_bytes(~"foo"); | |
781 | let k2 = str::to_bytes(~"bar"); | |
782 | let v1 = str::to_bytes(~"baz"); | |
783 | let v2 = str::to_bytes(~"foobar"); | |
784 | ||
785 | m.insert(copy k1, copy v1); | |
786 | m.insert(copy k2, copy v2); | |
787 | ||
788 | assert!(m.find(&k2) == Some(&v2)); | |
789 | assert!(m.find(&k1) == Some(&v1)); | |
790 | } | |
791 | ||
792 | fn check_equal<K: Eq + TotalOrd, V: Eq>(ctrl: &[(K, V)], | |
793 | map: &TreeMap<K, V>) { | |
794 | assert!(ctrl.is_empty() == map.is_empty()); | |
795 | for ctrl.each |x| { | |
796 | let &(k, v) = x; | |
797 | assert!(map.find(&k).unwrap() == &v) | |
798 | } | |
799 | for map.each |&(map_k, map_v)| { | |
800 | let mut found = false; | |
801 | for ctrl.each |x| { | |
802 | let &(ctrl_k, ctrl_v) = x; | |
803 | if *map_k == ctrl_k { | |
804 | assert!(*map_v == ctrl_v); | |
805 | found = true; | |
806 | break; | |
807 | } | |
808 | } | |
809 | assert!(found); | |
810 | } | |
811 | } | |
812 | ||
813 | fn check_left<K: TotalOrd, V>(node: &Option<~TreeNode<K, V>>, | |
814 | parent: &~TreeNode<K, V>) { | |
815 | match *node { | |
816 | Some(ref r) => { | |
817 | assert!(r.key.cmp(&parent.key) == Less); | |
818 | assert!(r.level == parent.level - 1); // left is black | |
819 | check_left(&r.left, r); | |
820 | check_right(&r.right, r, false); | |
821 | } | |
822 | None => assert!(parent.level == 1) // parent is leaf | |
823 | } | |
824 | } | |
825 | ||
826 | fn check_right<K: TotalOrd, V>(node: &Option<~TreeNode<K, V>>, | |
827 | parent: &~TreeNode<K, V>, | |
828 | parent_red: bool) { | |
829 | match *node { | |
830 | Some(ref r) => { | |
831 | assert!(r.key.cmp(&parent.key) == Greater); | |
832 | let red = r.level == parent.level; | |
833 | if parent_red { assert!(!red) } // no dual horizontal links | |
834 | // Right red or black | |
835 | assert!(red || r.level == parent.level - 1); | |
836 | check_left(&r.left, r); | |
837 | check_right(&r.right, r, red); | |
838 | } | |
839 | None => assert!(parent.level == 1) // parent is leaf | |
840 | } | |
841 | } | |
842 | ||
843 | fn check_structure<K: TotalOrd, V>(map: &TreeMap<K, V>) { | |
844 | match map.root { | |
845 | Some(ref r) => { | |
846 | check_left(&r.left, r); | |
847 | check_right(&r.right, r, false); | |
848 | } | |
849 | None => () | |
850 | } | |
851 | } | |
852 | ||
853 | #[test] | |
854 | fn test_rand_int() { | |
855 | let mut map = TreeMap::new::<int, int>(); | |
856 | let mut ctrl = ~[]; | |
857 | ||
858 | check_equal(ctrl, &map); | |
859 | assert!(map.find(&5).is_none()); | |
860 | ||
861 | let rng = rand::seeded_rng(&[42]); | |
862 | ||
863 | for 3.times { | |
864 | for 90.times { | |
865 | let k = rng.gen_int(); | |
866 | let v = rng.gen_int(); | |
867 | if !ctrl.contains(&(k, v)) { | |
868 | assert!(map.insert(k, v)); | |
869 | ctrl.push((k, v)); | |
870 | check_structure(&map); | |
871 | check_equal(ctrl, &map); | |
872 | } | |
873 | } | |
874 | ||
875 | for 30.times { | |
876 | let r = rng.gen_uint_range(0, ctrl.len()); | |
877 | let (key, _) = vec::remove(&mut ctrl, r); | |
878 | assert!(map.remove(&key)); | |
879 | check_structure(&map); | |
880 | check_equal(ctrl, &map); | |
881 | } | |
882 | } | |
883 | } | |
884 | ||
885 | #[test] | |
886 | fn test_len() { | |
887 | let mut m = TreeMap::new(); | |
888 | assert!(m.insert(3, 6)); | |
889 | assert!(m.len() == 1); | |
890 | assert!(m.insert(0, 0)); | |
891 | assert!(m.len() == 2); | |
892 | assert!(m.insert(4, 8)); | |
893 | assert!(m.len() == 3); | |
894 | assert!(m.remove(&3)); | |
895 | assert!(m.len() == 2); | |
896 | assert!(!m.remove(&5)); | |
897 | assert!(m.len() == 2); | |
898 | assert!(m.insert(2, 4)); | |
899 | assert!(m.len() == 3); | |
900 | assert!(m.insert(1, 2)); | |
901 | assert!(m.len() == 4); | |
902 | } | |
903 | ||
904 | #[test] | |
905 | fn test_each() { | |
906 | let mut m = TreeMap::new(); | |
907 | ||
908 | assert!(m.insert(3, 6)); | |
909 | assert!(m.insert(0, 0)); | |
910 | assert!(m.insert(4, 8)); | |
911 | assert!(m.insert(2, 4)); | |
912 | assert!(m.insert(1, 2)); | |
913 | ||
914 | let mut n = 0; | |
915 | for m.each |&(k, v)| { | |
916 | assert!(*k == n); | |
917 | assert!(*v == n * 2); | |
918 | n += 1; | |
919 | } | |
920 | } | |
921 | ||
922 | #[test] | |
923 | fn test_each_reverse() { | |
924 | let mut m = TreeMap::new(); | |
925 | ||
926 | assert!(m.insert(3, 6)); | |
927 | assert!(m.insert(0, 0)); | |
928 | assert!(m.insert(4, 8)); | |
929 | assert!(m.insert(2, 4)); | |
930 | assert!(m.insert(1, 2)); | |
931 | ||
932 | let mut n = 4; | |
933 | for m.each_reverse |&(k, v)| { | |
934 | assert!(*k == n); | |
935 | assert!(*v == n * 2); | |
936 | n -= 1; | |
937 | } | |
938 | } | |
939 | ||
940 | #[test] | |
941 | fn test_eq() { | |
942 | let mut a = TreeMap::new(); | |
943 | let mut b = TreeMap::new(); | |
944 | ||
945 | assert!(a == b); | |
946 | assert!(a.insert(0, 5)); | |
947 | assert!(a != b); | |
948 | assert!(b.insert(0, 4)); | |
949 | assert!(a != b); | |
950 | assert!(a.insert(5, 19)); | |
951 | assert!(a != b); | |
952 | assert!(!b.insert(0, 5)); | |
953 | assert!(a != b); | |
954 | assert!(b.insert(5, 19)); | |
955 | assert!(a == b); | |
956 | } | |
957 | ||
958 | #[test] | |
959 | fn test_lt() { | |
960 | let mut a = TreeMap::new(); | |
961 | let mut b = TreeMap::new(); | |
962 | ||
963 | assert!(!(a < b) && !(b < a)); | |
964 | assert!(b.insert(0, 5)); | |
965 | assert!(a < b); | |
966 | assert!(a.insert(0, 7)); | |
967 | assert!(!(a < b) && !(b < a)); | |
968 | assert!(b.insert(-2, 0)); | |
969 | assert!(b < a); | |
970 | assert!(a.insert(-5, 2)); | |
971 | assert!(a < b); | |
972 | assert!(a.insert(6, 2)); | |
973 | assert!(a < b && !(b < a)); | |
974 | } | |
975 | ||
976 | #[test] | |
977 | fn test_ord() { | |
978 | let mut a = TreeMap::new(); | |
979 | let mut b = TreeMap::new(); | |
980 | ||
981 | assert!(a <= b && a >= b); | |
982 | assert!(a.insert(1, 1)); | |
983 | assert!(a > b && a >= b); | |
984 | assert!(b < a && b <= a); | |
985 | assert!(b.insert(2, 2)); | |
986 | assert!(b > a && b >= a); | |
987 | assert!(a < b && a <= b); | |
988 | } | |
989 | ||
990 | #[test] | |
991 | fn test_lazy_iterator() { | |
992 | let mut m = TreeMap::new(); | |
993 | let (x1, y1) = (2, 5); | |
994 | let (x2, y2) = (9, 12); | |
995 | let (x3, y3) = (20, -3); | |
996 | let (x4, y4) = (29, 5); | |
997 | let (x5, y5) = (103, 3); | |
998 | ||
999 | assert!(m.insert(x1, y1)); | |
1000 | assert!(m.insert(x2, y2)); | |
1001 | assert!(m.insert(x3, y3)); | |
1002 | assert!(m.insert(x4, y4)); | |
1003 | assert!(m.insert(x5, y5)); | |
1004 | ||
1005 | let m = m; | |
1006 | let mut a = m.iter(); | |
1007 | ||
1008 | assert!(map_next(&mut a).unwrap() == (&x1, &y1)); | |
1009 | assert!(map_next(&mut a).unwrap() == (&x2, &y2)); | |
1010 | assert!(map_next(&mut a).unwrap() == (&x3, &y3)); | |
1011 | assert!(map_next(&mut a).unwrap() == (&x4, &y4)); | |
1012 | assert!(map_next(&mut a).unwrap() == (&x5, &y5)); | |
1013 | ||
1014 | assert!(map_next(&mut a).is_none()); | |
1015 | ||
1016 | let mut b = m.iter(); | |
1017 | ||
1018 | let expected = [(&x1, &y1), (&x2, &y2), (&x3, &y3), (&x4, &y4), | |
1019 | (&x5, &y5)]; | |
1020 | let mut i = 0; | |
1021 | ||
1022 | for map_advance(&mut b) |x| { | |
1023 | assert!(expected[i] == x); | |
1024 | i += 1; | |
1025 | ||
1026 | if i == 2 { | |
1027 | break | |
1028 | } | |
1029 | } | |
1030 | ||
1031 | for map_advance(&mut b) |x| { | |
1032 | assert!(expected[i] == x); | |
1033 | i += 1; | |
1034 | } | |
1035 | } | |
1036 | } | |
1037 | ||
1038 | #[cfg(test)] | |
1039 | mod test_set { | |
1040 | use super::*; | |
1041 | ||
1042 | #[test] | |
1043 | fn test_clear() { | |
1044 | let mut s = TreeSet::new(); | |
1045 | s.clear(); | |
1046 | assert!(s.insert(5)); | |
1047 | assert!(s.insert(12)); | |
1048 | assert!(s.insert(19)); | |
1049 | s.clear(); | |
1050 | assert!(!s.contains(&5)); | |
1051 | assert!(!s.contains(&12)); | |
1052 | assert!(!s.contains(&19)); | |
1053 | assert!(s.is_empty()); | |
1054 | } | |
1055 | ||
1056 | #[test] | |
1057 | fn test_disjoint() { | |
1058 | let mut xs = TreeSet::new(); | |
1059 | let mut ys = TreeSet::new(); | |
1060 | assert!(xs.is_disjoint(&ys)); | |
1061 | assert!(ys.is_disjoint(&xs)); | |
1062 | assert!(xs.insert(5)); | |
1063 | assert!(ys.insert(11)); | |
1064 | assert!(xs.is_disjoint(&ys)); | |
1065 | assert!(ys.is_disjoint(&xs)); | |
1066 | assert!(xs.insert(7)); | |
1067 | assert!(xs.insert(19)); | |
1068 | assert!(xs.insert(4)); | |
1069 | assert!(ys.insert(2)); | |
1070 | assert!(ys.insert(-11)); | |
1071 | assert!(xs.is_disjoint(&ys)); | |
1072 | assert!(ys.is_disjoint(&xs)); | |
1073 | assert!(ys.insert(7)); | |
1074 | assert!(!xs.is_disjoint(&ys)); | |
1075 | assert!(!ys.is_disjoint(&xs)); | |
1076 | } | |
1077 | ||
1078 | #[test] | |
1079 | fn test_subset_and_superset() { | |
1080 | let mut a = TreeSet::new(); | |
1081 | assert!(a.insert(0)); | |
1082 | assert!(a.insert(5)); | |
1083 | assert!(a.insert(11)); | |
1084 | assert!(a.insert(7)); | |
1085 | ||
1086 | let mut b = TreeSet::new(); | |
1087 | assert!(b.insert(0)); | |
1088 | assert!(b.insert(7)); | |
1089 | assert!(b.insert(19)); | |
1090 | assert!(b.insert(250)); | |
1091 | assert!(b.insert(11)); | |
1092 | assert!(b.insert(200)); | |
1093 | ||
1094 | assert!(!a.is_subset(&b)); | |
1095 | assert!(!a.is_superset(&b)); | |
1096 | assert!(!b.is_subset(&a)); | |
1097 | assert!(!b.is_superset(&a)); | |
1098 | ||
1099 | assert!(b.insert(5)); | |
1100 | ||
1101 | assert!(a.is_subset(&b)); | |
1102 | assert!(!a.is_superset(&b)); | |
1103 | assert!(!b.is_subset(&a)); | |
1104 | assert!(b.is_superset(&a)); | |
1105 | } | |
1106 | ||
1107 | #[test] | |
1108 | fn test_each() { | |
1109 | let mut m = TreeSet::new(); | |
1110 | ||
1111 | assert!(m.insert(3)); | |
1112 | assert!(m.insert(0)); | |
1113 | assert!(m.insert(4)); | |
1114 | assert!(m.insert(2)); | |
1115 | assert!(m.insert(1)); | |
1116 | ||
1117 | let mut n = 0; | |
1118 | for m.each |x| { | |
1119 | assert!(*x == n); | |
1120 | n += 1 | |
1121 | } | |
1122 | } | |
1123 | ||
1124 | #[test] | |
1125 | fn test_each_reverse() { | |
1126 | let mut m = TreeSet::new(); | |
1127 | ||
1128 | assert!(m.insert(3)); | |
1129 | assert!(m.insert(0)); | |
1130 | assert!(m.insert(4)); | |
1131 | assert!(m.insert(2)); | |
1132 | assert!(m.insert(1)); | |
1133 | ||
1134 | let mut n = 4; | |
1135 | for m.each_reverse |x| { | |
1136 | assert!(*x == n); | |
1137 | n -= 1 | |
1138 | } | |
1139 | } | |
1140 | ||
1141 | fn check(a: &[int], b: &[int], expected: &[int], | |
1142 | f: &fn(&TreeSet<int>, &TreeSet<int>, f: &fn(&int) -> bool)) { | |
1143 | let mut set_a = TreeSet::new(); | |
1144 | let mut set_b = TreeSet::new(); | |
1145 | ||
1146 | for a.each |x| { assert!(set_a.insert(*x)) } | |
1147 | for b.each |y| { assert!(set_b.insert(*y)) } | |
1148 | ||
1149 | let mut i = 0; | |
1150 | for f(&set_a, &set_b) |x| { | |
1151 | assert!(*x == expected[i]); | |
1152 | i += 1; | |
1153 | } | |
1154 | assert!(i == expected.len()); | |
1155 | } | |
1156 | ||
1157 | #[test] | |
1158 | fn test_intersection() { | |
1159 | fn check_intersection(a: &[int], b: &[int], expected: &[int]) { | |
1160 | check(a, b, expected, |x, y, z| x.intersection(y, z)) | |
1161 | } | |
1162 | ||
1163 | check_intersection([], [], []); | |
1164 | check_intersection([1, 2, 3], [], []); | |
1165 | check_intersection([], [1, 2, 3], []); | |
1166 | check_intersection([2], [1, 2, 3], [2]); | |
1167 | check_intersection([1, 2, 3], [2], [2]); | |
1168 | check_intersection([11, 1, 3, 77, 103, 5, -5], | |
1169 | [2, 11, 77, -9, -42, 5, 3], | |
1170 | [3, 5, 11, 77]); | |
1171 | } | |
1172 | ||
1173 | #[test] | |
1174 | fn test_difference() { | |
1175 | fn check_difference(a: &[int], b: &[int], expected: &[int]) { | |
1176 | check(a, b, expected, |x, y, z| x.difference(y, z)) | |
1177 | } | |
1178 | ||
1179 | check_difference([], [], []); | |
1180 | check_difference([1, 12], [], [1, 12]); | |
1181 | check_difference([], [1, 2, 3, 9], []); | |
1182 | check_difference([1, 3, 5, 9, 11], | |
1183 | [3, 9], | |
1184 | [1, 5, 11]); | |
1185 | check_difference([-5, 11, 22, 33, 40, 42], | |
1186 | [-12, -5, 14, 23, 34, 38, 39, 50], | |
1187 | [11, 22, 33, 40, 42]); | |
1188 | } | |
1189 | ||
1190 | #[test] | |
1191 | fn test_symmetric_difference() { | |
1192 | fn check_symmetric_difference(a: &[int], b: &[int], | |
1193 | expected: &[int]) { | |
1194 | check(a, b, expected, |x, y, z| x.symmetric_difference(y, z)) | |
1195 | } | |
1196 | ||
1197 | check_symmetric_difference([], [], []); | |
1198 | check_symmetric_difference([1, 2, 3], [2], [1, 3]); | |
1199 | check_symmetric_difference([2], [1, 2, 3], [1, 3]); | |
1200 | check_symmetric_difference([1, 3, 5, 9, 11], | |
1201 | [-2, 3, 9, 14, 22], | |
1202 | [-2, 1, 5, 11, 14, 22]); | |
1203 | } | |
1204 | ||
1205 | #[test] | |
1206 | fn test_union() { | |
1207 | fn check_union(a: &[int], b: &[int], | |
1208 | expected: &[int]) { | |
1209 | check(a, b, expected, |x, y, z| x.union(y, z)) | |
1210 | } | |
1211 | ||
1212 | check_union([], [], []); | |
1213 | check_union([1, 2, 3], [2], [1, 2, 3]); | |
1214 | check_union([2], [1, 2, 3], [1, 2, 3]); | |
1215 | check_union([1, 3, 5, 9, 11, 16, 19, 24], | |
1216 | [-2, 1, 5, 9, 13, 19], | |
1217 | [-2, 1, 3, 5, 9, 11, 13, 16, 19, 24]); | |
1218 | } | |
1219 | } |