1 // Copyright 2017-2018 By tickdream125@hotmail.com. 2 // 3 // Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or 4 // http://www.apache.org/licenses/LICENSE-2.0> or the MIT license 5 // <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your 6 // option. This file may not be copied, modified, or distributed 7 // except according to those terms. 8 9 // This file brings from https://github.com/tickbh/rbtree-rs/blob/master/src/lib.rs 10 // Thanks to the author tickbh 11 12 #![allow(dead_code)] 13 14 use core::cmp::Ord; 15 use core::cmp::Ordering; 16 use core::fmt::{self, Debug}; 17 use core::iter::{FromIterator, IntoIterator}; 18 use core::marker; 19 use core::mem; 20 use core::ops::Index; 21 use core::ptr; 22 23 use alloc::boxed::Box; 24 25 use crate::kdebug; 26 27 #[derive(Debug, Copy, Clone, PartialEq, Eq)] 28 enum Color { 29 Red, 30 Black, 31 } 32 33 /*****************RBTreeNode***************************/ 34 struct RBTreeNode<K: Ord, V> { 35 color: Color, 36 left: NodePtr<K, V>, 37 right: NodePtr<K, V>, 38 parent: NodePtr<K, V>, 39 key: K, 40 value: V, 41 } 42 43 impl<K: Ord, V> RBTreeNode<K, V> { 44 #[inline] pair(self) -> (K, V)45 fn pair(self) -> (K, V) { 46 (self.key, self.value) 47 } 48 } 49 50 impl<K, V> Debug for RBTreeNode<K, V> 51 where 52 K: Ord + Debug, 53 V: Debug, 54 { fmt(&self, f: &mut fmt::Formatter) -> fmt::Result55 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { 56 write!(f, "k:{:?} v:{:?} c:{:?}", self.key, self.value, self.color) 57 } 58 } 59 60 /*****************NodePtr***************************/ 61 #[derive(Debug)] 62 struct NodePtr<K: Ord, V>(*mut RBTreeNode<K, V>); 63 64 impl<K: Ord, V> Clone for NodePtr<K, V> { clone(&self) -> NodePtr<K, V>65 fn clone(&self) -> NodePtr<K, V> { 66 NodePtr(self.0) 67 } 68 } 69 70 impl<K: Ord, V> Copy for NodePtr<K, V> {} 71 72 impl<K: Ord, V> Ord for NodePtr<K, V> { cmp(&self, other: &NodePtr<K, V>) -> Ordering73 fn cmp(&self, other: &NodePtr<K, V>) -> Ordering { 74 unsafe { (*self.0).key.cmp(&(*other.0).key) } 75 } 76 } 77 78 impl<K: Ord, V> PartialOrd for NodePtr<K, V> { partial_cmp(&self, other: &NodePtr<K, V>) -> Option<Ordering>79 fn partial_cmp(&self, other: &NodePtr<K, V>) -> Option<Ordering> { 80 unsafe { Some((*self.0).key.cmp(&(*other.0).key)) } 81 } 82 } 83 84 impl<K: Ord, V> PartialEq for NodePtr<K, V> { eq(&self, other: &NodePtr<K, V>) -> bool85 fn eq(&self, other: &NodePtr<K, V>) -> bool { 86 self.0 == other.0 87 } 88 } 89 90 impl<K: Ord, V> Eq for NodePtr<K, V> {} 91 92 impl<K: Ord, V> NodePtr<K, V> { new(k: K, v: V) -> NodePtr<K, V>93 fn new(k: K, v: V) -> NodePtr<K, V> { 94 let node = RBTreeNode { 95 color: Color::Black, 96 left: NodePtr::null(), 97 right: NodePtr::null(), 98 parent: NodePtr::null(), 99 key: k, 100 value: v, 101 }; 102 NodePtr(Box::into_raw(Box::new(node))) 103 } 104 105 #[inline] set_color(&mut self, color: Color)106 fn set_color(&mut self, color: Color) { 107 if self.is_null() { 108 return; 109 } 110 unsafe { 111 (*self.0).color = color; 112 } 113 } 114 115 #[inline] set_red_color(&mut self)116 fn set_red_color(&mut self) { 117 self.set_color(Color::Red); 118 } 119 120 #[inline] set_black_color(&mut self)121 fn set_black_color(&mut self) { 122 self.set_color(Color::Black); 123 } 124 125 #[inline] get_color(&self) -> Color126 fn get_color(&self) -> Color { 127 if self.is_null() { 128 return Color::Black; 129 } 130 unsafe { (*self.0).color } 131 } 132 133 #[inline] is_red_color(&self) -> bool134 fn is_red_color(&self) -> bool { 135 if self.is_null() { 136 return false; 137 } 138 unsafe { (*self.0).color == Color::Red } 139 } 140 141 #[inline] is_black_color(&self) -> bool142 fn is_black_color(&self) -> bool { 143 if self.is_null() { 144 return true; 145 } 146 unsafe { (*self.0).color == Color::Black } 147 } 148 149 #[inline] is_left_child(&self) -> bool150 fn is_left_child(&self) -> bool { 151 self.parent().left() == *self 152 } 153 154 #[inline] is_right_child(&self) -> bool155 fn is_right_child(&self) -> bool { 156 self.parent().right() == *self 157 } 158 159 #[inline] min_node(self) -> NodePtr<K, V>160 fn min_node(self) -> NodePtr<K, V> { 161 let mut temp = self.clone(); 162 while !temp.left().is_null() { 163 temp = temp.left(); 164 } 165 return temp; 166 } 167 168 #[inline] max_node(self) -> NodePtr<K, V>169 fn max_node(self) -> NodePtr<K, V> { 170 let mut temp = self.clone(); 171 while !temp.right().is_null() { 172 temp = temp.right(); 173 } 174 return temp; 175 } 176 177 #[inline] next(self) -> NodePtr<K, V>178 fn next(self) -> NodePtr<K, V> { 179 if !self.right().is_null() { 180 self.right().min_node() 181 } else { 182 let mut temp = self; 183 loop { 184 if temp.parent().is_null() { 185 return NodePtr::null(); 186 } 187 if temp.is_left_child() { 188 return temp.parent(); 189 } 190 temp = temp.parent(); 191 } 192 } 193 } 194 195 #[inline] prev(self) -> NodePtr<K, V>196 fn prev(self) -> NodePtr<K, V> { 197 if !self.left().is_null() { 198 self.left().max_node() 199 } else { 200 let mut temp = self; 201 loop { 202 if temp.parent().is_null() { 203 return NodePtr::null(); 204 } 205 if temp.is_right_child() { 206 return temp.parent(); 207 } 208 temp = temp.parent(); 209 } 210 } 211 } 212 213 #[inline] set_parent(&mut self, parent: NodePtr<K, V>)214 fn set_parent(&mut self, parent: NodePtr<K, V>) { 215 if self.is_null() { 216 return; 217 } 218 unsafe { (*self.0).parent = parent } 219 } 220 221 #[inline] set_left(&mut self, left: NodePtr<K, V>)222 fn set_left(&mut self, left: NodePtr<K, V>) { 223 if self.is_null() { 224 return; 225 } 226 unsafe { (*self.0).left = left } 227 } 228 229 #[inline] set_right(&mut self, right: NodePtr<K, V>)230 fn set_right(&mut self, right: NodePtr<K, V>) { 231 if self.is_null() { 232 return; 233 } 234 unsafe { (*self.0).right = right } 235 } 236 237 #[inline] parent(&self) -> NodePtr<K, V>238 fn parent(&self) -> NodePtr<K, V> { 239 if self.is_null() { 240 return NodePtr::null(); 241 } 242 unsafe { (*self.0).parent.clone() } 243 } 244 245 #[inline] left(&self) -> NodePtr<K, V>246 fn left(&self) -> NodePtr<K, V> { 247 if self.is_null() { 248 return NodePtr::null(); 249 } 250 unsafe { (*self.0).left.clone() } 251 } 252 253 #[inline] right(&self) -> NodePtr<K, V>254 fn right(&self) -> NodePtr<K, V> { 255 if self.is_null() { 256 return NodePtr::null(); 257 } 258 unsafe { (*self.0).right.clone() } 259 } 260 261 #[inline] null() -> NodePtr<K, V>262 fn null() -> NodePtr<K, V> { 263 NodePtr(ptr::null_mut()) 264 } 265 266 #[inline] is_null(&self) -> bool267 fn is_null(&self) -> bool { 268 self.0.is_null() 269 } 270 } 271 272 impl<K: Ord + Clone, V: Clone> NodePtr<K, V> { deep_clone(&self) -> NodePtr<K, V>273 unsafe fn deep_clone(&self) -> NodePtr<K, V> { 274 let mut node = NodePtr::new((*self.0).key.clone(), (*self.0).value.clone()); 275 if !self.left().is_null() { 276 node.set_left(self.left().deep_clone()); 277 node.left().set_parent(node); 278 } 279 if !self.right().is_null() { 280 node.set_right(self.right().deep_clone()); 281 node.right().set_parent(node); 282 } 283 node 284 } 285 } 286 287 /// A red black tree implemented with Rust 288 /// It is required that the keys implement the [`Ord`] traits. 289 290 /// # Examples 291 /// ```rust 292 /// use rbtree::RBTree; 293 /// // type inference lets us omit an explicit type signature (which 294 /// // would be `RBTree<&str, &str>` in this example). 295 /// let mut book_reviews = RBTree::new(); 296 /// 297 /// // review some books. 298 /// book_reviews.insert("Adventures of Huckleberry Finn", "My favorite book."); 299 /// book_reviews.insert("Grimms' Fairy Tales", "Masterpiece."); 300 /// book_reviews.insert("Pride and Prejudice", "Very enjoyable."); 301 /// book_reviews.insert("The Adventures of Sherlock Holmes", "Eye lyked it alot."); 302 /// 303 /// // check for a specific one. 304 /// if !book_reviews.contains_key(&"Les Misérables") { 305 /// kdebug!( 306 /// "We've got {} reviews, but Les Misérables ain't one.", 307 /// book_reviews.len() 308 /// ); 309 /// } 310 /// 311 /// // oops, this review has a lot of spelling mistakes, let's delete it. 312 /// book_reviews.remove(&"The Adventures of Sherlock Holmes"); 313 /// 314 /// // look up the values associated with some keys. 315 /// let to_find = ["Pride and Prejudice", "Alice's Adventure in Wonderland"]; 316 /// for book in &to_find { 317 /// match book_reviews.get(book) { 318 /// Some(review) => kdebug!("{}: {}", book, review), 319 /// None => kdebug!("{} is unreviewed.", book), 320 /// } 321 /// } 322 /// 323 /// // iterate over everything. 324 /// for (book, review) in book_reviews.iter() { 325 /// kdebug!("{}: \"{}\"", book, review); 326 /// } 327 /// 328 /// book_reviews.print_tree(); 329 /// ``` 330 /// 331 /// // A `RBTree` with fixed list of elements can be initialized from an array: 332 /// ``` 333 /// use rbtree::RBTree; 334 /// let timber_resources: RBTree<&str, i32> = 335 /// [("Norway", 100), 336 /// ("Denmark", 50), 337 /// ("Iceland", 10)] 338 /// .iter().cloned().collect(); 339 /// // use the values stored in rbtree 340 /// ``` 341 pub struct RBTree<K: Ord, V> { 342 root: NodePtr<K, V>, 343 len: usize, 344 } 345 346 // Drop all owned pointers if the tree is dropped 347 impl<K: Ord, V> Drop for RBTree<K, V> { 348 #[inline] drop(&mut self)349 fn drop(&mut self) { 350 self.clear(); 351 } 352 } 353 354 /// If key and value are both impl Clone, we can call clone to get a copy. 355 impl<K: Ord + Clone, V: Clone> Clone for RBTree<K, V> { clone(&self) -> RBTree<K, V>356 fn clone(&self) -> RBTree<K, V> { 357 unsafe { 358 let mut new = RBTree::new(); 359 new.root = self.root.deep_clone(); 360 new.len = self.len; 361 new 362 } 363 } 364 } 365 366 impl<K, V> Debug for RBTree<K, V> 367 where 368 K: Ord + Debug, 369 V: Debug, 370 { fmt(&self, f: &mut fmt::Formatter) -> fmt::Result371 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { 372 f.debug_map().entries(self.iter()).finish() 373 } 374 } 375 376 /// This is a method to help us to get inner struct. 377 impl<K: Ord + Debug, V: Debug> RBTree<K, V> { tree_print(&self, node: NodePtr<K, V>, direction: i32)378 fn tree_print(&self, node: NodePtr<K, V>, direction: i32) { 379 if node.is_null() { 380 return; 381 } 382 if direction == 0 { 383 unsafe { 384 kdebug!("'{:?}' is root node", (*node.0)); 385 } 386 } else { 387 let direct = if direction == -1 { "left" } else { "right" }; 388 unsafe { 389 kdebug!( 390 "{:?} is {:?}'s {:?} child ", 391 (*node.0), 392 *node.parent().0, 393 direct 394 ); 395 } 396 } 397 self.tree_print(node.left(), -1); 398 self.tree_print(node.right(), 1); 399 } 400 print_tree(&self)401 pub fn print_tree(&self) { 402 if self.root.is_null() { 403 kdebug!("This is a empty tree"); 404 return; 405 } 406 kdebug!("This tree size = {:?}, begin:-------------", self.len()); 407 self.tree_print(self.root, 0); 408 kdebug!("end--------------------------"); 409 } 410 } 411 412 /// all key be same, but it has multi key, if has multi key, it perhaps no correct 413 impl<K, V> PartialEq for RBTree<K, V> 414 where 415 K: Eq + Ord, 416 V: PartialEq, 417 { eq(&self, other: &RBTree<K, V>) -> bool418 fn eq(&self, other: &RBTree<K, V>) -> bool { 419 if self.len() != other.len() { 420 return false; 421 } 422 423 self.iter() 424 .all(|(key, value)| other.get(key).map_or(false, |v| *value == *v)) 425 } 426 } 427 428 impl<K, V> Eq for RBTree<K, V> 429 where 430 K: Eq + Ord, 431 V: Eq, 432 { 433 } 434 435 impl<'a, K, V> Index<&'a K> for RBTree<K, V> 436 where 437 K: Ord, 438 { 439 type Output = V; 440 441 #[inline] index(&self, index: &K) -> &V442 fn index(&self, index: &K) -> &V { 443 self.get(index).expect("no entry found for key") 444 } 445 } 446 447 impl<K: Ord, V> FromIterator<(K, V)> for RBTree<K, V> { from_iter<T: IntoIterator<Item = (K, V)>>(iter: T) -> RBTree<K, V>448 fn from_iter<T: IntoIterator<Item = (K, V)>>(iter: T) -> RBTree<K, V> { 449 let mut tree = RBTree::new(); 450 tree.extend(iter); 451 tree 452 } 453 } 454 455 /// RBTree into iter 456 impl<K: Ord, V> Extend<(K, V)> for RBTree<K, V> { extend<T: IntoIterator<Item = (K, V)>>(&mut self, iter: T)457 fn extend<T: IntoIterator<Item = (K, V)>>(&mut self, iter: T) { 458 let iter = iter.into_iter(); 459 for (k, v) in iter { 460 self.insert(k, v); 461 } 462 } 463 } 464 465 /// provide the rbtree all keys 466 /// # Examples 467 /// ``` 468 /// use rbtree::RBTree; 469 /// let mut m = RBTree::new(); 470 /// for i in 1..6 { 471 /// m.insert(i, i); 472 /// } 473 /// let vec = vec![1, 2, 3, 4, 5]; 474 /// let key_vec: Vec<_> = m.keys().cloned().collect(); 475 /// assert_eq!(vec, key_vec); 476 /// ``` 477 pub struct Keys<'a, K: Ord + 'a, V: 'a> { 478 inner: Iter<'a, K, V>, 479 } 480 481 impl<'a, K: Ord, V> Clone for Keys<'a, K, V> { clone(&self) -> Keys<'a, K, V>482 fn clone(&self) -> Keys<'a, K, V> { 483 Keys { 484 inner: self.inner.clone(), 485 } 486 } 487 } 488 489 impl<'a, K: Ord + Debug, V> fmt::Debug for Keys<'a, K, V> { fmt(&self, f: &mut fmt::Formatter) -> fmt::Result490 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { 491 f.debug_list().entries(self.clone()).finish() 492 } 493 } 494 495 impl<'a, K: Ord, V> Iterator for Keys<'a, K, V> { 496 type Item = &'a K; 497 498 #[inline] next(&mut self) -> Option<&'a K>499 fn next(&mut self) -> Option<&'a K> { 500 self.inner.next().map(|(k, _)| k) 501 } 502 503 #[inline] size_hint(&self) -> (usize, Option<usize>)504 fn size_hint(&self) -> (usize, Option<usize>) { 505 self.inner.size_hint() 506 } 507 } 508 509 /// provide the rbtree all values order by key 510 /// # Examples 511 /// ``` 512 /// use rbtree::RBTree; 513 /// let mut m = RBTree::new(); 514 /// m.insert(2, 5); 515 /// m.insert(1, 6); 516 /// m.insert(3, 8); 517 /// m.insert(4, 9); 518 /// let vec = vec![6, 5, 8, 9]; 519 /// let key_vec: Vec<_> = m.values().cloned().collect(); 520 /// assert_eq!(vec, key_vec); 521 /// ``` 522 pub struct Values<'a, K: 'a + Ord, V: 'a> { 523 inner: Iter<'a, K, V>, 524 } 525 526 impl<'a, K: Ord, V> Clone for Values<'a, K, V> { clone(&self) -> Values<'a, K, V>527 fn clone(&self) -> Values<'a, K, V> { 528 Values { 529 inner: self.inner.clone(), 530 } 531 } 532 } 533 534 impl<'a, K: Ord + Debug, V: Debug> fmt::Debug for Values<'a, K, V> { fmt(&self, f: &mut fmt::Formatter) -> fmt::Result535 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { 536 f.debug_list().entries(self.clone()).finish() 537 } 538 } 539 540 impl<'a, K: Ord, V> Iterator for Values<'a, K, V> { 541 type Item = &'a V; 542 543 #[inline] next(&mut self) -> Option<&'a V>544 fn next(&mut self) -> Option<&'a V> { 545 self.inner.next().map(|(_, v)| v) 546 } 547 548 #[inline] size_hint(&self) -> (usize, Option<usize>)549 fn size_hint(&self) -> (usize, Option<usize>) { 550 self.inner.size_hint() 551 } 552 } 553 554 /// provide the rbtree all values and it can be modify 555 /// # Examples 556 /// ``` 557 /// use rbtree::RBTree; 558 /// let mut m = RBTree::new(); 559 /// for i in 0..32 { 560 /// m.insert(i, i); 561 /// } 562 /// assert_eq!(m.len(), 32); 563 /// for v in m.values_mut() { 564 /// *v *= 2; 565 /// } 566 /// for i in 0..32 { 567 /// assert_eq!(m.get(&i).unwrap(), &(i * 2)); 568 /// } 569 /// ``` 570 pub struct ValuesMut<'a, K: 'a + Ord, V: 'a> { 571 inner: IterMut<'a, K, V>, 572 } 573 574 impl<'a, K: Ord, V> Clone for ValuesMut<'a, K, V> { clone(&self) -> ValuesMut<'a, K, V>575 fn clone(&self) -> ValuesMut<'a, K, V> { 576 ValuesMut { 577 inner: self.inner.clone(), 578 } 579 } 580 } 581 582 impl<'a, K: Ord + Debug, V: Debug> fmt::Debug for ValuesMut<'a, K, V> { fmt(&self, f: &mut fmt::Formatter) -> fmt::Result583 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { 584 f.debug_list().entries(self.clone()).finish() 585 } 586 } 587 588 impl<'a, K: Ord, V> Iterator for ValuesMut<'a, K, V> { 589 type Item = &'a mut V; 590 591 #[inline] next(&mut self) -> Option<&'a mut V>592 fn next(&mut self) -> Option<&'a mut V> { 593 self.inner.next().map(|(_, v)| v) 594 } 595 596 #[inline] size_hint(&self) -> (usize, Option<usize>)597 fn size_hint(&self) -> (usize, Option<usize>) { 598 self.inner.size_hint() 599 } 600 } 601 602 /// Convert RBTree to iter, move out the tree. 603 pub struct IntoIter<K: Ord, V> { 604 head: NodePtr<K, V>, 605 tail: NodePtr<K, V>, 606 len: usize, 607 } 608 609 // Drop all owned pointers if the collection is dropped 610 impl<K: Ord, V> Drop for IntoIter<K, V> { 611 #[inline] drop(&mut self)612 fn drop(&mut self) { 613 for (_, _) in self {} 614 } 615 } 616 617 impl<K: Ord, V> Iterator for IntoIter<K, V> { 618 type Item = (K, V); 619 next(&mut self) -> Option<(K, V)>620 fn next(&mut self) -> Option<(K, V)> { 621 if self.len == 0 { 622 return None; 623 } 624 625 if self.head.is_null() { 626 return None; 627 } 628 629 let next = self.head.next(); 630 let (k, v) = unsafe { 631 ( 632 core::ptr::read(&(*self.head.0).key), 633 core::ptr::read(&(*self.head.0).value), 634 ) 635 }; 636 self.head = next; 637 self.len -= 1; 638 Some((k, v)) 639 } 640 size_hint(&self) -> (usize, Option<usize>)641 fn size_hint(&self) -> (usize, Option<usize>) { 642 (self.len, Some(self.len)) 643 } 644 } 645 646 impl<K: Ord, V> DoubleEndedIterator for IntoIter<K, V> { 647 #[inline] next_back(&mut self) -> Option<(K, V)>648 fn next_back(&mut self) -> Option<(K, V)> { 649 if self.len == 0 { 650 return None; 651 } 652 653 if self.tail.is_null() { 654 return None; 655 } 656 657 let prev = self.tail.prev(); 658 let obj = unsafe { Box::from_raw(self.tail.0) }; 659 let (k, v) = obj.pair(); 660 self.tail = prev; 661 self.len -= 1; 662 Some((k, v)) 663 } 664 } 665 666 /// provide iter ref for RBTree 667 /// # Examples 668 /// ``` 669 /// use rbtree::RBTree; 670 /// let mut m = RBTree::new(); 671 /// for i in 0..32 { 672 /// m.insert(i, i * 2); 673 /// } 674 /// assert_eq!(m.len(), 32); 675 /// let mut observed: u32 = 0; 676 /// for (k, v) in m.iter() { 677 /// assert_eq!(*v, *k * 2); 678 /// observed |= 1 << *k; 679 /// } 680 /// assert_eq!(observed, 0xFFFF_FFFF); 681 /// ``` 682 pub struct Iter<'a, K: Ord + 'a, V: 'a> { 683 head: NodePtr<K, V>, 684 tail: NodePtr<K, V>, 685 len: usize, 686 _marker: marker::PhantomData<&'a ()>, 687 } 688 689 impl<'a, K: Ord + 'a, V: 'a> Clone for Iter<'a, K, V> { clone(&self) -> Iter<'a, K, V>690 fn clone(&self) -> Iter<'a, K, V> { 691 Iter { 692 head: self.head, 693 tail: self.tail, 694 len: self.len, 695 _marker: self._marker, 696 } 697 } 698 } 699 700 impl<'a, K: Ord + 'a, V: 'a> Iterator for Iter<'a, K, V> { 701 type Item = (&'a K, &'a V); 702 next(&mut self) -> Option<(&'a K, &'a V)>703 fn next(&mut self) -> Option<(&'a K, &'a V)> { 704 if self.len == 0 { 705 return None; 706 } 707 708 if self.head.is_null() { 709 return None; 710 } 711 712 let (k, v) = unsafe { (&(*self.head.0).key, &(*self.head.0).value) }; 713 self.head = self.head.next(); 714 self.len -= 1; 715 Some((k, v)) 716 } 717 size_hint(&self) -> (usize, Option<usize>)718 fn size_hint(&self) -> (usize, Option<usize>) { 719 (self.len, Some(self.len)) 720 } 721 } 722 723 impl<'a, K: Ord + 'a, V: 'a> DoubleEndedIterator for Iter<'a, K, V> { 724 #[inline] next_back(&mut self) -> Option<(&'a K, &'a V)>725 fn next_back(&mut self) -> Option<(&'a K, &'a V)> { 726 // kdebug!("len = {:?}", self.len); 727 if self.len == 0 { 728 return None; 729 } 730 731 if self.tail == self.head { 732 return None; 733 } 734 735 let (k, v) = unsafe { (&(*self.tail.0).key, &(*self.tail.0).value) }; 736 self.tail = self.tail.prev(); 737 self.len -= 1; 738 Some((k, v)) 739 } 740 } 741 742 /// provide iter mut ref for RBTree 743 /// # Examples 744 /// ``` 745 /// use rbtree::RBTree; 746 /// let mut m = RBTree::new(); 747 /// for i in 0..32 { 748 /// m.insert(i, i); 749 /// } 750 /// assert_eq!(m.len(), 32); 751 /// for (_, v) in m.iter_mut() { 752 /// *v *= 2; 753 /// } 754 /// for i in 0..32 { 755 /// assert_eq!(m.get(&i).unwrap(), &(i * 2)); 756 /// } 757 /// ``` 758 pub struct IterMut<'a, K: Ord + 'a, V: 'a> { 759 head: NodePtr<K, V>, 760 tail: NodePtr<K, V>, 761 len: usize, 762 _marker: marker::PhantomData<&'a ()>, 763 } 764 765 impl<'a, K: Ord + 'a, V: 'a> Clone for IterMut<'a, K, V> { clone(&self) -> IterMut<'a, K, V>766 fn clone(&self) -> IterMut<'a, K, V> { 767 IterMut { 768 head: self.head, 769 tail: self.tail, 770 len: self.len, 771 _marker: self._marker, 772 } 773 } 774 } 775 776 impl<'a, K: Ord + 'a, V: 'a> Iterator for IterMut<'a, K, V> { 777 type Item = (&'a K, &'a mut V); 778 next(&mut self) -> Option<(&'a K, &'a mut V)>779 fn next(&mut self) -> Option<(&'a K, &'a mut V)> { 780 if self.len == 0 { 781 return None; 782 } 783 784 if self.head.is_null() { 785 return None; 786 } 787 788 let (k, v) = unsafe { (&(*self.head.0).key, &mut (*self.head.0).value) }; 789 self.head = self.head.next(); 790 self.len -= 1; 791 Some((k, v)) 792 } 793 size_hint(&self) -> (usize, Option<usize>)794 fn size_hint(&self) -> (usize, Option<usize>) { 795 (self.len, Some(self.len)) 796 } 797 } 798 799 impl<'a, K: Ord + 'a, V: 'a> DoubleEndedIterator for IterMut<'a, K, V> { 800 #[inline] next_back(&mut self) -> Option<(&'a K, &'a mut V)>801 fn next_back(&mut self) -> Option<(&'a K, &'a mut V)> { 802 if self.len == 0 { 803 return None; 804 } 805 806 if self.tail == self.head { 807 return None; 808 } 809 810 let (k, v) = unsafe { (&(*self.tail.0).key, &mut (*self.tail.0).value) }; 811 self.tail = self.tail.prev(); 812 self.len -= 1; 813 Some((k, v)) 814 } 815 } 816 817 impl<K: Ord, V> IntoIterator for RBTree<K, V> { 818 type Item = (K, V); 819 type IntoIter = IntoIter<K, V>; 820 821 #[inline] into_iter(mut self) -> IntoIter<K, V>822 fn into_iter(mut self) -> IntoIter<K, V> { 823 let iter = if self.root.is_null() { 824 IntoIter { 825 head: NodePtr::null(), 826 tail: NodePtr::null(), 827 len: self.len, 828 } 829 } else { 830 IntoIter { 831 head: self.first_child(), 832 tail: self.last_child(), 833 len: self.len, 834 } 835 }; 836 self.fast_clear(); 837 iter 838 } 839 } 840 841 impl<K: Ord, V> RBTree<K, V> { 842 /// Creates an empty `RBTree`. new() -> RBTree<K, V>843 pub fn new() -> RBTree<K, V> { 844 RBTree { 845 root: NodePtr::null(), 846 len: 0, 847 } 848 } 849 850 /// Returns the len of `RBTree`. 851 #[inline] len(&self) -> usize852 pub fn len(&self) -> usize { 853 self.len 854 } 855 856 /// Returns `true` if the `RBTree` is empty. 857 #[inline] is_empty(&self) -> bool858 pub fn is_empty(&self) -> bool { 859 self.root.is_null() 860 } 861 862 /* 863 * 对红黑树的节点(x)进行左旋转 864 * 865 * 左旋示意图(对节点x进行左旋)： 866 * px px 867 * / / 868 * x y 869 * / \ --(左旋)--> / \ # 870 * lx y x ry 871 * / \ / \ 872 * ly ry lx ly 873 * 874 * 875 */ 876 #[inline] left_rotate(&mut self, mut node: NodePtr<K, V>)877 unsafe fn left_rotate(&mut self, mut node: NodePtr<K, V>) { 878 let mut temp = node.right(); 879 node.set_right(temp.left()); 880 881 if !temp.left().is_null() { 882 temp.left().set_parent(node.clone()); 883 } 884 885 temp.set_parent(node.parent()); 886 if node == self.root { 887 self.root = temp.clone(); 888 } else if node == node.parent().left() { 889 node.parent().set_left(temp.clone()); 890 } else { 891 node.parent().set_right(temp.clone()); 892 } 893 894 temp.set_left(node.clone()); 895 node.set_parent(temp.clone()); 896 } 897 898 /* 899 * 对红黑树的节点(y)进行右旋转 900 * 901 * 右旋示意图(对节点y进行左旋)： 902 * py py 903 * / / 904 * y x 905 * / \ --(右旋)--> / \ # 906 * x ry lx y 907 * / \ / \ # 908 * lx rx rx ry 909 * 910 */ 911 #[inline] right_rotate(&mut self, mut node: NodePtr<K, V>)912 unsafe fn right_rotate(&mut self, mut node: NodePtr<K, V>) { 913 let mut temp = node.left(); 914 node.set_left(temp.right()); 915 916 if !temp.right().is_null() { 917 temp.right().set_parent(node.clone()); 918 } 919 920 temp.set_parent(node.parent()); 921 if node == self.root { 922 self.root = temp.clone(); 923 } else if node == node.parent().right() { 924 node.parent().set_right(temp.clone()); 925 } else { 926 node.parent().set_left(temp.clone()); 927 } 928 929 temp.set_right(node.clone()); 930 node.set_parent(temp.clone()); 931 } 932 933 /// replace value if key exist, if not exist insert it. 934 /// # Examples 935 /// ``` 936 /// use rbtree::RBTree; 937 /// let mut m = RBTree::new(); 938 /// assert_eq!(m.len(), 0); 939 /// m.insert(2, 4); 940 /// assert_eq!(m.len(), 1); 941 /// assert_eq!(m.replace_or_insert(2, 6).unwrap(), 4); 942 /// assert_eq!(m.len(), 1); 943 /// assert_eq!(*m.get(&2).unwrap(), 6); 944 /// ``` 945 #[inline] replace_or_insert(&mut self, k: K, mut v: V) -> Option<V>946 pub fn replace_or_insert(&mut self, k: K, mut v: V) -> Option<V> { 947 let node = self.find_node(&k); 948 if node.is_null() { 949 self.insert(k, v); 950 return None; 951 } 952 953 unsafe { 954 mem::swap(&mut v, &mut (*node.0).value); 955 } 956 957 Some(v) 958 } 959 960 #[inline] insert_fixup(&mut self, mut node: NodePtr<K, V>)961 unsafe fn insert_fixup(&mut self, mut node: NodePtr<K, V>) { 962 let mut parent; 963 let mut gparent; 964 965 while node.parent().is_red_color() { 966 parent = node.parent(); 967 gparent = parent.parent(); 968 //若“父节点”是“祖父节点的左孩子” 969 if parent == gparent.left() { 970 // Case 1条件：叔叔节点是红色 971 let mut uncle = gparent.right(); 972 if !uncle.is_null() && uncle.is_red_color() { 973 uncle.set_black_color(); 974 parent.set_black_color(); 975 gparent.set_red_color(); 976 node = gparent; 977 continue; 978 } 979 980 // Case 2条件：叔叔是黑色，且当前节点是右孩子 981 if parent.right() == node { 982 self.left_rotate(parent); 983 let temp = parent; 984 parent = node; 985 node = temp; 986 } 987 988 // Case 3条件：叔叔是黑色，且当前节点是左孩子。 989 parent.set_black_color(); 990 gparent.set_red_color(); 991 self.right_rotate(gparent); 992 } else { 993 // Case 1条件：叔叔节点是红色 994 let mut uncle = gparent.left(); 995 if !uncle.is_null() && uncle.is_red_color() { 996 uncle.set_black_color(); 997 parent.set_black_color(); 998 gparent.set_red_color(); 999 node = gparent; 1000 continue; 1001 } 1002 1003 // Case 2条件：叔叔是黑色，且当前节点是右孩子 1004 if parent.left() == node { 1005 self.right_rotate(parent); 1006 let temp = parent; 1007 parent = node; 1008 node = temp; 1009 } 1010 1011 // Case 3条件：叔叔是黑色，且当前节点是左孩子。 1012 parent.set_black_color(); 1013 gparent.set_red_color(); 1014 self.left_rotate(gparent); 1015 } 1016 } 1017 self.root.set_black_color(); 1018 } 1019 1020 #[inline] insert(&mut self, k: K, v: V)1021 pub fn insert(&mut self, k: K, v: V) { 1022 self.len += 1; 1023 let mut node = NodePtr::new(k, v); 1024 let mut y = NodePtr::null(); 1025 let mut x = self.root; 1026 1027 while !x.is_null() { 1028 y = x; 1029 match node.cmp(&&mut x) { 1030 Ordering::Less => { 1031 x = x.left(); 1032 } 1033 _ => { 1034 x = x.right(); 1035 } 1036 }; 1037 } 1038 node.set_parent(y); 1039 1040 if y.is_null() { 1041 self.root = node; 1042 } else { 1043 match node.cmp(&&mut y) { 1044 Ordering::Less => { 1045 y.set_left(node); 1046 } 1047 _ => { 1048 y.set_right(node); 1049 } 1050 }; 1051 } 1052 1053 node.set_red_color(); 1054 unsafe { 1055 self.insert_fixup(node); 1056 } 1057 } 1058 1059 #[inline] find_node(&self, k: &K) -> NodePtr<K, V>1060 fn find_node(&self, k: &K) -> NodePtr<K, V> { 1061 if self.root.is_null() { 1062 return NodePtr::null(); 1063 } 1064 let mut temp = &self.root; 1065 unsafe { 1066 loop { 1067 let next = match k.cmp(&(*temp.0).key) { 1068 Ordering::Less => &mut (*temp.0).left, 1069 Ordering::Greater => &mut (*temp.0).right, 1070 Ordering::Equal => return *temp, 1071 }; 1072 if next.is_null() { 1073 break; 1074 } 1075 temp = next; 1076 } 1077 } 1078 NodePtr::null() 1079 } 1080 1081 #[inline] first_child(&self) -> NodePtr<K, V>1082 fn first_child(&self) -> NodePtr<K, V> { 1083 if self.root.is_null() { 1084 NodePtr::null() 1085 } else { 1086 let mut temp = self.root; 1087 while !temp.left().is_null() { 1088 temp = temp.left(); 1089 } 1090 return temp; 1091 } 1092 } 1093 1094 #[inline] last_child(&self) -> NodePtr<K, V>1095 fn last_child(&self) -> NodePtr<K, V> { 1096 if self.root.is_null() { 1097 NodePtr::null() 1098 } else { 1099 let mut temp = self.root; 1100 while !temp.right().is_null() { 1101 temp = temp.right(); 1102 } 1103 return temp; 1104 } 1105 } 1106 1107 #[inline] get_first(&self) -> Option<(&K, &V)>1108 pub fn get_first(&self) -> Option<(&K, &V)> { 1109 let first = self.first_child(); 1110 if first.is_null() { 1111 return None; 1112 } 1113 unsafe { Some((&(*first.0).key, &(*first.0).value)) } 1114 } 1115 1116 #[inline] get_last(&self) -> Option<(&K, &V)>1117 pub fn get_last(&self) -> Option<(&K, &V)> { 1118 let last = self.last_child(); 1119 if last.is_null() { 1120 return None; 1121 } 1122 unsafe { Some((&(*last.0).key, &(*last.0).value)) } 1123 } 1124 1125 #[inline] pop_first(&mut self) -> Option<(K, V)>1126 pub fn pop_first(&mut self) -> Option<(K, V)> { 1127 let first = self.first_child(); 1128 if first.is_null() { 1129 return None; 1130 } 1131 unsafe { Some(self.delete(first)) } 1132 } 1133 1134 #[inline] pop_last(&mut self) -> Option<(K, V)>1135 pub fn pop_last(&mut self) -> Option<(K, V)> { 1136 let last = self.last_child(); 1137 if last.is_null() { 1138 return None; 1139 } 1140 unsafe { Some(self.delete(last)) } 1141 } 1142 1143 #[inline] get_first_mut(&mut self) -> Option<(&K, &mut V)>1144 pub fn get_first_mut(&mut self) -> Option<(&K, &mut V)> { 1145 let first = self.first_child(); 1146 if first.is_null() { 1147 return None; 1148 } 1149 unsafe { Some((&(*first.0).key, &mut (*first.0).value)) } 1150 } 1151 1152 #[inline] get_last_mut(&mut self) -> Option<(&K, &mut V)>1153 pub fn get_last_mut(&mut self) -> Option<(&K, &mut V)> { 1154 let last = self.last_child(); 1155 if last.is_null() { 1156 return None; 1157 } 1158 unsafe { Some((&(*last.0).key, &mut (*last.0).value)) } 1159 } 1160 1161 #[inline] get(&self, k: &K) -> Option<&V>1162 pub fn get(&self, k: &K) -> Option<&V> { 1163 let node = self.find_node(k); 1164 if node.is_null() { 1165 return None; 1166 } 1167 1168 unsafe { Some(&(*node.0).value) } 1169 } 1170 1171 #[inline] get_mut(&mut self, k: &K) -> Option<&mut V>1172 pub fn get_mut(&mut self, k: &K) -> Option<&mut V> { 1173 let node = self.find_node(k); 1174 if node.is_null() { 1175 return None; 1176 } 1177 1178 unsafe { Some(&mut (*node.0).value) } 1179 } 1180 1181 #[inline] contains_key(&self, k: &K) -> bool1182 pub fn contains_key(&self, k: &K) -> bool { 1183 let node = self.find_node(k); 1184 if node.is_null() { 1185 return false; 1186 } 1187 true 1188 } 1189 1190 #[inline] clear_recurse(&mut self, current: NodePtr<K, V>)1191 fn clear_recurse(&mut self, current: NodePtr<K, V>) { 1192 if !current.is_null() { 1193 unsafe { 1194 self.clear_recurse(current.left()); 1195 self.clear_recurse(current.right()); 1196 drop(Box::from_raw(current.0)); 1197 } 1198 } 1199 } 1200 1201 #[inline] clear(&mut self)1202 pub fn clear(&mut self) { 1203 let root = self.root; 1204 self.root = NodePtr::null(); 1205 self.clear_recurse(root); 1206 } 1207 1208 /// Empties the `RBTree` without freeing objects in it. 1209 #[inline] fast_clear(&mut self)1210 fn fast_clear(&mut self) { 1211 self.root = NodePtr::null(); 1212 } 1213 1214 #[inline] remove(&mut self, k: &K) -> Option<V>1215 pub fn remove(&mut self, k: &K) -> Option<V> { 1216 let node = self.find_node(k); 1217 if node.is_null() { 1218 return None; 1219 } 1220 unsafe { Some(self.delete(node).1) } 1221 } 1222 1223 #[inline] delete_fixup(&mut self, mut node: NodePtr<K, V>, mut parent: NodePtr<K, V>)1224 unsafe fn delete_fixup(&mut self, mut node: NodePtr<K, V>, mut parent: NodePtr<K, V>) { 1225 let mut other; 1226 while node != self.root && node.is_black_color() { 1227 if parent.left() == node { 1228 other = parent.right(); 1229 //x的兄弟w是红色的 1230 if other.is_red_color() { 1231 other.set_black_color(); 1232 parent.set_red_color(); 1233 self.left_rotate(parent); 1234 other = parent.right(); 1235 } 1236 1237 //x的兄弟w是黑色，且w的俩个孩子也都是黑色的 1238 if other.left().is_black_color() && other.right().is_black_color() { 1239 other.set_red_color(); 1240 node = parent; 1241 parent = node.parent(); 1242 } else { 1243 //x的兄弟w是黑色的，并且w的左孩子是红色，右孩子为黑色。 1244 if other.right().is_black_color() { 1245 other.left().set_black_color(); 1246 other.set_red_color(); 1247 self.right_rotate(other); 1248 other = parent.right(); 1249 } 1250 //x的兄弟w是黑色的；并且w的右孩子是红色的，左孩子任意颜色。 1251 other.set_color(parent.get_color()); 1252 parent.set_black_color(); 1253 other.right().set_black_color(); 1254 self.left_rotate(parent); 1255 node = self.root; 1256 break; 1257 } 1258 } else { 1259 other = parent.left(); 1260 //x的兄弟w是红色的 1261 if other.is_red_color() { 1262 other.set_black_color(); 1263 parent.set_red_color(); 1264 self.right_rotate(parent); 1265 other = parent.left(); 1266 } 1267 1268 //x的兄弟w是黑色，且w的俩个孩子也都是黑色的 1269 if other.left().is_black_color() && other.right().is_black_color() { 1270 other.set_red_color(); 1271 node = parent; 1272 parent = node.parent(); 1273 } else { 1274 //x的兄弟w是黑色的，并且w的左孩子是红色，右孩子为黑色。 1275 if other.left().is_black_color() { 1276 other.right().set_black_color(); 1277 other.set_red_color(); 1278 self.left_rotate(other); 1279 other = parent.left(); 1280 } 1281 //x的兄弟w是黑色的；并且w的右孩子是红色的，左孩子任意颜色。 1282 other.set_color(parent.get_color()); 1283 parent.set_black_color(); 1284 other.left().set_black_color(); 1285 self.right_rotate(parent); 1286 node = self.root; 1287 break; 1288 } 1289 } 1290 } 1291 1292 node.set_black_color(); 1293 } 1294 1295 #[inline] delete(&mut self, node: NodePtr<K, V>) -> (K, V)1296 unsafe fn delete(&mut self, node: NodePtr<K, V>) -> (K, V) { 1297 let mut child; 1298 let mut parent; 1299 let color; 1300 1301 self.len -= 1; 1302 // 被删除节点的"左右孩子都不为空"的情况。 1303 if !node.left().is_null() && !node.right().is_null() { 1304 // 被删节点的后继节点。(称为"取代节点") 1305 // 用它来取代"被删节点"的位置，然后再将"被删节点"去掉。 1306 let mut replace = node.right().min_node(); 1307 if node == self.root { 1308 self.root = replace; 1309 } else { 1310 if node.parent().left() == node { 1311 node.parent().set_left(replace); 1312 } else { 1313 node.parent().set_right(replace); 1314 } 1315 } 1316 1317 // child是"取代节点"的右孩子，也是需要"调整的节点"。 1318 // "取代节点"肯定不存在左孩子！因为它是一个后继节点。 1319 child = replace.right(); 1320 parent = replace.parent(); 1321 color = replace.get_color(); 1322 if parent == node { 1323 parent = replace; 1324 } else { 1325 if !child.is_null() { 1326 child.set_parent(parent); 1327 } 1328 parent.set_left(child); 1329 replace.set_right(node.right()); 1330 node.right().set_parent(replace); 1331 } 1332 1333 replace.set_parent(node.parent()); 1334 replace.set_color(node.get_color()); 1335 replace.set_left(node.left()); 1336 node.left().set_parent(replace); 1337 1338 if color == Color::Black { 1339 self.delete_fixup(child, parent); 1340 } 1341 1342 let obj = Box::from_raw(node.0); 1343 return obj.pair(); 1344 } 1345 1346 if !node.left().is_null() { 1347 child = node.left(); 1348 } else { 1349 child = node.right(); 1350 } 1351 1352 parent = node.parent(); 1353 color = node.get_color(); 1354 if !child.is_null() { 1355 child.set_parent(parent); 1356 } 1357 1358 if self.root == node { 1359 self.root = child 1360 } else { 1361 if parent.left() == node { 1362 parent.set_left(child); 1363 } else { 1364 parent.set_right(child); 1365 } 1366 } 1367 1368 if color == Color::Black { 1369 self.delete_fixup(child, parent); 1370 } 1371 1372 let obj = Box::from_raw(node.0); 1373 return obj.pair(); 1374 } 1375 1376 /// Return the keys iter 1377 #[inline] keys(&self) -> Keys<K, V>1378 pub fn keys(&self) -> Keys<K, V> { 1379 Keys { inner: self.iter() } 1380 } 1381 1382 /// Return the value iter 1383 #[inline] values(&self) -> Values<K, V>1384 pub fn values(&self) -> Values<K, V> { 1385 Values { inner: self.iter() } 1386 } 1387 1388 /// Return the value iter mut 1389 #[inline] values_mut(&mut self) -> ValuesMut<K, V>1390 pub fn values_mut(&mut self) -> ValuesMut<K, V> { 1391 ValuesMut { 1392 inner: self.iter_mut(), 1393 } 1394 } 1395 1396 /// Return the key and value iter 1397 #[inline] iter(&self) -> Iter<K, V>1398 pub fn iter(&self) -> Iter<K, V> { 1399 Iter { 1400 head: self.first_child(), 1401 tail: self.last_child(), 1402 len: self.len, 1403 _marker: marker::PhantomData, 1404 } 1405 } 1406 1407 /// Return the key and mut value iter 1408 #[inline] iter_mut(&mut self) -> IterMut<K, V>1409 pub fn iter_mut(&mut self) -> IterMut<K, V> { 1410 IterMut { 1411 head: self.first_child(), 1412 tail: self.last_child(), 1413 len: self.len, 1414 _marker: marker::PhantomData, 1415 } 1416 } 1417 } 1418 1419 mod tests { 1420 1421 #[test] test_insert()1422 fn test_insert() { 1423 let mut m = RBTree::new(); 1424 assert_eq!(m.len(), 0); 1425 m.insert(1, 2); 1426 assert_eq!(m.len(), 1); 1427 m.insert(2, 4); 1428 assert_eq!(m.len(), 2); 1429 m.insert(2, 6); 1430 assert_eq!(m.len(), 3); 1431 assert_eq!(*m.get(&1).unwrap(), 2); 1432 assert_eq!(*m.get(&2).unwrap(), 4); 1433 assert_eq!(*m.get(&2).unwrap(), 4); 1434 } 1435 1436 #[test] test_replace()1437 fn test_replace() { 1438 let mut m = RBTree::new(); 1439 assert_eq!(m.len(), 0); 1440 m.insert(2, 4); 1441 assert_eq!(m.len(), 1); 1442 assert_eq!(m.replace_or_insert(2, 6).unwrap(), 4); 1443 assert_eq!(m.len(), 1); 1444 assert_eq!(*m.get(&2).unwrap(), 6); 1445 } 1446 1447 #[test] test_clone()1448 fn test_clone() { 1449 let mut m = RBTree::new(); 1450 assert_eq!(m.len(), 0); 1451 m.insert(1, 2); 1452 assert_eq!(m.len(), 1); 1453 m.insert(2, 4); 1454 assert_eq!(m.len(), 2); 1455 let m2 = m.clone(); 1456 m.clear(); 1457 assert_eq!(*m2.get(&1).unwrap(), 2); 1458 assert_eq!(*m2.get(&2).unwrap(), 4); 1459 assert_eq!(m2.len(), 2); 1460 } 1461 1462 #[test] test_empty_remove()1463 fn test_empty_remove() { 1464 let mut m: RBTree<isize, bool> = RBTree::new(); 1465 assert_eq!(m.remove(&0), None); 1466 } 1467 1468 #[test] test_empty_iter()1469 fn test_empty_iter() { 1470 let mut m: RBTree<isize, bool> = RBTree::new(); 1471 assert_eq!(m.iter().next(), None); 1472 assert_eq!(m.iter_mut().next(), None); 1473 assert_eq!(m.len(), 0); 1474 assert!(m.is_empty()); 1475 assert_eq!(m.into_iter().next(), None); 1476 } 1477 1478 #[test] test_lots_of_insertions()1479 fn test_lots_of_insertions() { 1480 let mut m = RBTree::new(); 1481 1482 // Try this a few times to make sure we never screw up the hashmap's 1483 // internal state. 1484 for _ in 0..10 { 1485 assert!(m.is_empty()); 1486 1487 for i in 1..101 { 1488 m.insert(i, i); 1489 1490 for j in 1..i + 1 { 1491 let r = m.get(&j); 1492 assert_eq!(r, Some(&j)); 1493 } 1494 1495 for j in i + 1..101 { 1496 let r = m.get(&j); 1497 assert_eq!(r, None); 1498 } 1499 } 1500 1501 for i in 101..201 { 1502 assert!(!m.contains_key(&i)); 1503 } 1504 1505 // remove forwards 1506 for i in 1..101 { 1507 assert!(m.remove(&i).is_some()); 1508 1509 for j in 1..i + 1 { 1510 assert!(!m.contains_key(&j)); 1511 } 1512 1513 for j in i + 1..101 { 1514 assert!(m.contains_key(&j)); 1515 } 1516 } 1517 1518 for i in 1..101 { 1519 assert!(!m.contains_key(&i)); 1520 } 1521 1522 for i in 1..101 { 1523 m.insert(i, i); 1524 } 1525 1526 // remove backwards 1527 for i in (1..101).rev() { 1528 assert!(m.remove(&i).is_some()); 1529 1530 for j in i..101 { 1531 assert!(!m.contains_key(&j)); 1532 } 1533 1534 for j in 1..i { 1535 assert!(m.contains_key(&j)); 1536 } 1537 } 1538 } 1539 } 1540 1541 #[test] test_find_mut()1542 fn test_find_mut() { 1543 let mut m = RBTree::new(); 1544 m.insert(1, 12); 1545 m.insert(2, 8); 1546 m.insert(5, 14); 1547 let new = 100; 1548 match m.get_mut(&5) { 1549 None => panic!(), 1550 Some(x) => *x = new, 1551 } 1552 assert_eq!(m.get(&5), Some(&new)); 1553 } 1554 1555 #[test] test_remove()1556 fn test_remove() { 1557 let mut m = RBTree::new(); 1558 m.insert(1, 2); 1559 assert_eq!(*m.get(&1).unwrap(), 2); 1560 m.insert(5, 3); 1561 assert_eq!(*m.get(&5).unwrap(), 3); 1562 m.insert(9, 4); 1563 assert_eq!(*m.get(&1).unwrap(), 2); 1564 assert_eq!(*m.get(&5).unwrap(), 3); 1565 assert_eq!(*m.get(&9).unwrap(), 4); 1566 assert_eq!(m.remove(&1).unwrap(), 2); 1567 assert_eq!(m.remove(&5).unwrap(), 3); 1568 assert_eq!(m.remove(&9).unwrap(), 4); 1569 assert_eq!(m.len(), 0); 1570 } 1571 1572 #[test] test_is_empty()1573 fn test_is_empty() { 1574 let mut m = RBTree::new(); 1575 m.insert(1, 2); 1576 assert!(!m.is_empty()); 1577 assert!(m.remove(&1).is_some()); 1578 assert!(m.is_empty()); 1579 } 1580 1581 #[test] test_pop()1582 fn test_pop() { 1583 let mut m = RBTree::new(); 1584 m.insert(2, 4); 1585 m.insert(1, 2); 1586 m.insert(3, 6); 1587 assert_eq!(m.len(), 3); 1588 assert_eq!(m.pop_first(), Some((1, 2))); 1589 assert_eq!(m.len(), 2); 1590 assert_eq!(m.pop_last(), Some((3, 6))); 1591 assert_eq!(m.len(), 1); 1592 assert_eq!(m.get_first(), Some((&2, &4))); 1593 assert_eq!(m.get_last(), Some((&2, &4))); 1594 } 1595 1596 #[test] test_iterate()1597 fn test_iterate() { 1598 let mut m = RBTree::new(); 1599 for i in 0..32 { 1600 m.insert(i, i * 2); 1601 } 1602 assert_eq!(m.len(), 32); 1603 1604 let mut observed: u32 = 0; 1605 1606 for (k, v) in m.iter() { 1607 assert_eq!(*v, *k * 2); 1608 observed |= 1 << *k; 1609 } 1610 assert_eq!(observed, 0xFFFF_FFFF); 1611 } 1612 1613 #[test] test_keys()1614 fn test_keys() { 1615 let vec = vec![(1, 'a'), (2, 'b'), (3, 'c')]; 1616 let map: RBTree<_, _> = vec.into_iter().collect(); 1617 let keys: Vec<_> = map.keys().cloned().collect(); 1618 assert_eq!(keys.len(), 3); 1619 assert!(keys.contains(&1)); 1620 assert!(keys.contains(&2)); 1621 assert!(keys.contains(&3)); 1622 } 1623 1624 #[test] test_values()1625 fn test_values() { 1626 let vec = vec![(1, 'a'), (2, 'b'), (3, 'c')]; 1627 let map: RBTree<_, _> = vec.into_iter().collect(); 1628 let values: Vec<_> = map.values().cloned().collect(); 1629 assert_eq!(values.len(), 3); 1630 assert!(values.contains(&'a')); 1631 assert!(values.contains(&'b')); 1632 assert!(values.contains(&'c')); 1633 } 1634 1635 #[test] test_values_mut()1636 fn test_values_mut() { 1637 let vec = vec![(1, 1), (2, 2), (3, 3)]; 1638 let mut map: RBTree<_, _> = vec.into_iter().collect(); 1639 for value in map.values_mut() { 1640 *value = (*value) * 2 1641 } 1642 let values: Vec<_> = map.values().cloned().collect(); 1643 assert_eq!(values.len(), 3); 1644 assert!(values.contains(&2)); 1645 assert!(values.contains(&4)); 1646 assert!(values.contains(&6)); 1647 } 1648 1649 #[test] test_find()1650 fn test_find() { 1651 let mut m = RBTree::new(); 1652 assert!(m.get(&1).is_none()); 1653 m.insert(1, 2); 1654 match m.get(&1) { 1655 None => panic!(), 1656 Some(v) => assert_eq!(*v, 2), 1657 } 1658 } 1659 1660 #[test] test_eq()1661 fn test_eq() { 1662 let mut m1 = RBTree::new(); 1663 m1.insert(1, 2); 1664 m1.insert(2, 3); 1665 m1.insert(3, 4); 1666 1667 let mut m2 = RBTree::new(); 1668 m2.insert(1, 2); 1669 m2.insert(2, 3); 1670 1671 assert!(m1 != m2); 1672 1673 m2.insert(3, 4); 1674 1675 assert_eq!(m1, m2); 1676 } 1677 1678 #[test] test_show()1679 fn test_show() { 1680 let mut map = RBTree::new(); 1681 let empty: RBTree<i32, i32> = RBTree::new(); 1682 1683 map.insert(1, 2); 1684 map.insert(3, 4); 1685 1686 let map_str = format!("{:?}", map); 1687 1688 assert!(map_str == "{1: 2, 3: 4}" || map_str == "{3: 4, 1: 2}"); 1689 assert_eq!(format!("{:?}", empty), "{}"); 1690 } 1691 1692 #[test] test_from_iter()1693 fn test_from_iter() { 1694 let xs = [(1, 1), (2, 2), (3, 3), (4, 4), (5, 5), (6, 6)]; 1695 1696 let map: RBTree<_, _> = xs.iter().cloned().collect(); 1697 1698 for &(k, v) in &xs { 1699 assert_eq!(map.get(&k), Some(&v)); 1700 } 1701 } 1702 1703 #[test] test_size_hint()1704 fn test_size_hint() { 1705 let xs = [(1, 1), (2, 2), (3, 3), (4, 4), (5, 5), (6, 6)]; 1706 1707 let map: RBTree<_, _> = xs.iter().cloned().collect(); 1708 1709 let mut iter = map.iter(); 1710 1711 for _ in iter.by_ref().take(3) {} 1712 1713 assert_eq!(iter.size_hint(), (3, Some(3))); 1714 } 1715 1716 #[test] test_iter_len()1717 fn test_iter_len() { 1718 let xs = [(1, 1), (2, 2), (3, 3), (4, 4), (5, 5), (6, 6)]; 1719 1720 let map: RBTree<_, _> = xs.iter().cloned().collect(); 1721 1722 let mut iter = map.iter(); 1723 1724 for _ in iter.by_ref().take(3) {} 1725 1726 assert_eq!(iter.count(), 3); 1727 } 1728 1729 #[test] test_mut_size_hint()1730 fn test_mut_size_hint() { 1731 let xs = [(1, 1), (2, 2), (3, 3), (4, 4), (5, 5), (6, 6)]; 1732 1733 let mut map: RBTree<_, _> = xs.iter().cloned().collect(); 1734 1735 let mut iter = map.iter_mut(); 1736 1737 for _ in iter.by_ref().take(3) {} 1738 1739 assert_eq!(iter.size_hint(), (3, Some(3))); 1740 } 1741 1742 #[test] test_iter_mut_len()1743 fn test_iter_mut_len() { 1744 let xs = [(1, 1), (2, 2), (3, 3), (4, 4), (5, 5), (6, 6)]; 1745 1746 let mut map: RBTree<_, _> = xs.iter().cloned().collect(); 1747 1748 let mut iter = map.iter_mut(); 1749 1750 for _ in iter.by_ref().take(3) {} 1751 1752 assert_eq!(iter.count(), 3); 1753 } 1754 1755 #[test] test_index()1756 fn test_index() { 1757 let mut map = RBTree::new(); 1758 1759 map.insert(1, 2); 1760 map.insert(2, 1); 1761 map.insert(3, 4); 1762 1763 assert_eq!(map[&2], 1); 1764 } 1765 1766 #[test] 1767 #[should_panic] test_index_nonexistent()1768 fn test_index_nonexistent() { 1769 let mut map = RBTree::new(); 1770 1771 map.insert(1, 2); 1772 map.insert(2, 1); 1773 map.insert(3, 4); 1774 1775 map[&4]; 1776 } 1777 1778 #[test] test_extend_iter()1779 fn test_extend_iter() { 1780 let mut a = RBTree::new(); 1781 a.insert(1, "one"); 1782 let mut b = RBTree::new(); 1783 b.insert(2, "two"); 1784 b.insert(3, "three"); 1785 1786 a.extend(b.into_iter()); 1787 1788 assert_eq!(a.len(), 3); 1789 assert_eq!(a[&1], "one"); 1790 assert_eq!(a[&2], "two"); 1791 assert_eq!(a[&3], "three"); 1792 } 1793 } 1794