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 unsafe impl<K: Ord, V> Send for RBTree<K, V> {} 347 unsafe impl<K: Ord, V> Sync for RBTree<K, V> {} 348 349 // Drop all owned pointers if the tree is dropped 350 impl<K: Ord, V> Drop for RBTree<K, V> { 351 #[inline] drop(&mut self)352 fn drop(&mut self) { 353 self.clear(); 354 } 355 } 356 357 /// If key and value are both impl Clone, we can call clone to get a copy. 358 impl<K: Ord + Clone, V: Clone> Clone for RBTree<K, V> { clone(&self) -> RBTree<K, V>359 fn clone(&self) -> RBTree<K, V> { 360 unsafe { 361 let mut new = RBTree::new(); 362 new.root = self.root.deep_clone(); 363 new.len = self.len; 364 new 365 } 366 } 367 } 368 369 impl<K, V> Debug for RBTree<K, V> 370 where 371 K: Ord + Debug, 372 V: Debug, 373 { fmt(&self, f: &mut fmt::Formatter) -> fmt::Result374 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { 375 f.debug_map().entries(self.iter()).finish() 376 } 377 } 378 379 /// This is a method to help us to get inner struct. 380 impl<K: Ord + Debug, V: Debug> RBTree<K, V> { tree_print(&self, node: NodePtr<K, V>, direction: i32)381 fn tree_print(&self, node: NodePtr<K, V>, direction: i32) { 382 if node.is_null() { 383 return; 384 } 385 if direction == 0 { 386 unsafe { 387 kdebug!("'{:?}' is root node", (*node.0)); 388 } 389 } else { 390 let direct = if direction == -1 { "left" } else { "right" }; 391 unsafe { 392 kdebug!( 393 "{:?} is {:?}'s {:?} child ", 394 (*node.0), 395 *node.parent().0, 396 direct 397 ); 398 } 399 } 400 self.tree_print(node.left(), -1); 401 self.tree_print(node.right(), 1); 402 } 403 print_tree(&self)404 pub fn print_tree(&self) { 405 if self.root.is_null() { 406 kdebug!("This is a empty tree"); 407 return; 408 } 409 kdebug!("This tree size = {:?}, begin:-------------", self.len()); 410 self.tree_print(self.root, 0); 411 kdebug!("end--------------------------"); 412 } 413 } 414 415 /// all key be same, but it has multi key, if has multi key, it perhaps no correct 416 impl<K, V> PartialEq for RBTree<K, V> 417 where 418 K: Eq + Ord, 419 V: PartialEq, 420 { eq(&self, other: &RBTree<K, V>) -> bool421 fn eq(&self, other: &RBTree<K, V>) -> bool { 422 if self.len() != other.len() { 423 return false; 424 } 425 426 self.iter() 427 .all(|(key, value)| other.get(key).map_or(false, |v| *value == *v)) 428 } 429 } 430 431 impl<K, V> Eq for RBTree<K, V> 432 where 433 K: Eq + Ord, 434 V: Eq, 435 { 436 } 437 438 impl<'a, K, V> Index<&'a K> for RBTree<K, V> 439 where 440 K: Ord, 441 { 442 type Output = V; 443 444 #[inline] index(&self, index: &K) -> &V445 fn index(&self, index: &K) -> &V { 446 self.get(index).expect("no entry found for key") 447 } 448 } 449 450 impl<K: Ord, V> FromIterator<(K, V)> for RBTree<K, V> { from_iter<T: IntoIterator<Item = (K, V)>>(iter: T) -> RBTree<K, V>451 fn from_iter<T: IntoIterator<Item = (K, V)>>(iter: T) -> RBTree<K, V> { 452 let mut tree = RBTree::new(); 453 tree.extend(iter); 454 tree 455 } 456 } 457 458 /// RBTree into iter 459 impl<K: Ord, V> Extend<(K, V)> for RBTree<K, V> { extend<T: IntoIterator<Item = (K, V)>>(&mut self, iter: T)460 fn extend<T: IntoIterator<Item = (K, V)>>(&mut self, iter: T) { 461 let iter = iter.into_iter(); 462 for (k, v) in iter { 463 self.insert(k, v); 464 } 465 } 466 } 467 468 /// provide the rbtree all keys 469 /// # Examples 470 /// ``` 471 /// use rbtree::RBTree; 472 /// let mut m = RBTree::new(); 473 /// for i in 1..6 { 474 /// m.insert(i, i); 475 /// } 476 /// let vec = vec![1, 2, 3, 4, 5]; 477 /// let key_vec: Vec<_> = m.keys().cloned().collect(); 478 /// assert_eq!(vec, key_vec); 479 /// ``` 480 pub struct Keys<'a, K: Ord + 'a, V: 'a> { 481 inner: Iter<'a, K, V>, 482 } 483 484 impl<'a, K: Ord, V> Clone for Keys<'a, K, V> { clone(&self) -> Keys<'a, K, V>485 fn clone(&self) -> Keys<'a, K, V> { 486 Keys { 487 inner: self.inner.clone(), 488 } 489 } 490 } 491 492 impl<'a, K: Ord + Debug, V> fmt::Debug for Keys<'a, K, V> { fmt(&self, f: &mut fmt::Formatter) -> fmt::Result493 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { 494 f.debug_list().entries(self.clone()).finish() 495 } 496 } 497 498 impl<'a, K: Ord, V> Iterator for Keys<'a, K, V> { 499 type Item = &'a K; 500 501 #[inline] next(&mut self) -> Option<&'a K>502 fn next(&mut self) -> Option<&'a K> { 503 self.inner.next().map(|(k, _)| k) 504 } 505 506 #[inline] size_hint(&self) -> (usize, Option<usize>)507 fn size_hint(&self) -> (usize, Option<usize>) { 508 self.inner.size_hint() 509 } 510 } 511 512 /// provide the rbtree all values order by key 513 /// # Examples 514 /// ``` 515 /// use rbtree::RBTree; 516 /// let mut m = RBTree::new(); 517 /// m.insert(2, 5); 518 /// m.insert(1, 6); 519 /// m.insert(3, 8); 520 /// m.insert(4, 9); 521 /// let vec = vec![6, 5, 8, 9]; 522 /// let key_vec: Vec<_> = m.values().cloned().collect(); 523 /// assert_eq!(vec, key_vec); 524 /// ``` 525 pub struct Values<'a, K: 'a + Ord, V: 'a> { 526 inner: Iter<'a, K, V>, 527 } 528 529 impl<'a, K: Ord, V> Clone for Values<'a, K, V> { clone(&self) -> Values<'a, K, V>530 fn clone(&self) -> Values<'a, K, V> { 531 Values { 532 inner: self.inner.clone(), 533 } 534 } 535 } 536 537 impl<'a, K: Ord + Debug, V: Debug> fmt::Debug for Values<'a, K, V> { fmt(&self, f: &mut fmt::Formatter) -> fmt::Result538 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { 539 f.debug_list().entries(self.clone()).finish() 540 } 541 } 542 543 impl<'a, K: Ord, V> Iterator for Values<'a, K, V> { 544 type Item = &'a V; 545 546 #[inline] next(&mut self) -> Option<&'a V>547 fn next(&mut self) -> Option<&'a V> { 548 self.inner.next().map(|(_, v)| v) 549 } 550 551 #[inline] size_hint(&self) -> (usize, Option<usize>)552 fn size_hint(&self) -> (usize, Option<usize>) { 553 self.inner.size_hint() 554 } 555 } 556 557 /// provide the rbtree all values and it can be modify 558 /// # Examples 559 /// ``` 560 /// use rbtree::RBTree; 561 /// let mut m = RBTree::new(); 562 /// for i in 0..32 { 563 /// m.insert(i, i); 564 /// } 565 /// assert_eq!(m.len(), 32); 566 /// for v in m.values_mut() { 567 /// *v *= 2; 568 /// } 569 /// for i in 0..32 { 570 /// assert_eq!(m.get(&i).unwrap(), &(i * 2)); 571 /// } 572 /// ``` 573 pub struct ValuesMut<'a, K: 'a + Ord, V: 'a> { 574 inner: IterMut<'a, K, V>, 575 } 576 577 impl<'a, K: Ord, V> Clone for ValuesMut<'a, K, V> { clone(&self) -> ValuesMut<'a, K, V>578 fn clone(&self) -> ValuesMut<'a, K, V> { 579 ValuesMut { 580 inner: self.inner.clone(), 581 } 582 } 583 } 584 585 impl<'a, K: Ord + Debug, V: Debug> fmt::Debug for ValuesMut<'a, K, V> { fmt(&self, f: &mut fmt::Formatter) -> fmt::Result586 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { 587 f.debug_list().entries(self.clone()).finish() 588 } 589 } 590 591 impl<'a, K: Ord, V> Iterator for ValuesMut<'a, K, V> { 592 type Item = &'a mut V; 593 594 #[inline] next(&mut self) -> Option<&'a mut V>595 fn next(&mut self) -> Option<&'a mut V> { 596 self.inner.next().map(|(_, v)| v) 597 } 598 599 #[inline] size_hint(&self) -> (usize, Option<usize>)600 fn size_hint(&self) -> (usize, Option<usize>) { 601 self.inner.size_hint() 602 } 603 } 604 605 /// Convert RBTree to iter, move out the tree. 606 pub struct IntoIter<K: Ord, V> { 607 head: NodePtr<K, V>, 608 tail: NodePtr<K, V>, 609 len: usize, 610 } 611 612 // Drop all owned pointers if the collection is dropped 613 impl<K: Ord, V> Drop for IntoIter<K, V> { 614 #[inline] drop(&mut self)615 fn drop(&mut self) { 616 for (_, _) in self {} 617 } 618 } 619 620 impl<K: Ord, V> Iterator for IntoIter<K, V> { 621 type Item = (K, V); 622 next(&mut self) -> Option<(K, V)>623 fn next(&mut self) -> Option<(K, V)> { 624 if self.len == 0 { 625 return None; 626 } 627 628 if self.head.is_null() { 629 return None; 630 } 631 632 let next = self.head.next(); 633 let (k, v) = unsafe { 634 ( 635 core::ptr::read(&(*self.head.0).key), 636 core::ptr::read(&(*self.head.0).value), 637 ) 638 }; 639 self.head = next; 640 self.len -= 1; 641 Some((k, v)) 642 } 643 size_hint(&self) -> (usize, Option<usize>)644 fn size_hint(&self) -> (usize, Option<usize>) { 645 (self.len, Some(self.len)) 646 } 647 } 648 649 impl<K: Ord, V> DoubleEndedIterator for IntoIter<K, V> { 650 #[inline] next_back(&mut self) -> Option<(K, V)>651 fn next_back(&mut self) -> Option<(K, V)> { 652 if self.len == 0 { 653 return None; 654 } 655 656 if self.tail.is_null() { 657 return None; 658 } 659 660 let prev = self.tail.prev(); 661 let obj = unsafe { Box::from_raw(self.tail.0) }; 662 let (k, v) = obj.pair(); 663 self.tail = prev; 664 self.len -= 1; 665 Some((k, v)) 666 } 667 } 668 669 /// provide iter ref for RBTree 670 /// # Examples 671 /// ``` 672 /// use rbtree::RBTree; 673 /// let mut m = RBTree::new(); 674 /// for i in 0..32 { 675 /// m.insert(i, i * 2); 676 /// } 677 /// assert_eq!(m.len(), 32); 678 /// let mut observed: u32 = 0; 679 /// for (k, v) in m.iter() { 680 /// assert_eq!(*v, *k * 2); 681 /// observed |= 1 << *k; 682 /// } 683 /// assert_eq!(observed, 0xFFFF_FFFF); 684 /// ``` 685 pub struct Iter<'a, K: Ord + 'a, V: 'a> { 686 head: NodePtr<K, V>, 687 tail: NodePtr<K, V>, 688 len: usize, 689 _marker: marker::PhantomData<&'a ()>, 690 } 691 692 impl<'a, K: Ord + 'a, V: 'a> Clone for Iter<'a, K, V> { clone(&self) -> Iter<'a, K, V>693 fn clone(&self) -> Iter<'a, K, V> { 694 Iter { 695 head: self.head, 696 tail: self.tail, 697 len: self.len, 698 _marker: self._marker, 699 } 700 } 701 } 702 703 impl<'a, K: Ord + 'a, V: 'a> Iterator for Iter<'a, K, V> { 704 type Item = (&'a K, &'a V); 705 next(&mut self) -> Option<(&'a K, &'a V)>706 fn next(&mut self) -> Option<(&'a K, &'a V)> { 707 if self.len == 0 { 708 return None; 709 } 710 711 if self.head.is_null() { 712 return None; 713 } 714 715 let (k, v) = unsafe { (&(*self.head.0).key, &(*self.head.0).value) }; 716 self.head = self.head.next(); 717 self.len -= 1; 718 Some((k, v)) 719 } 720 size_hint(&self) -> (usize, Option<usize>)721 fn size_hint(&self) -> (usize, Option<usize>) { 722 (self.len, Some(self.len)) 723 } 724 } 725 726 impl<'a, K: Ord + 'a, V: 'a> DoubleEndedIterator for Iter<'a, K, V> { 727 #[inline] next_back(&mut self) -> Option<(&'a K, &'a V)>728 fn next_back(&mut self) -> Option<(&'a K, &'a V)> { 729 // kdebug!("len = {:?}", self.len); 730 if self.len == 0 { 731 return None; 732 } 733 734 if self.tail == self.head { 735 return None; 736 } 737 738 let (k, v) = unsafe { (&(*self.tail.0).key, &(*self.tail.0).value) }; 739 self.tail = self.tail.prev(); 740 self.len -= 1; 741 Some((k, v)) 742 } 743 } 744 745 /// provide iter mut ref for RBTree 746 /// # Examples 747 /// ``` 748 /// use rbtree::RBTree; 749 /// let mut m = RBTree::new(); 750 /// for i in 0..32 { 751 /// m.insert(i, i); 752 /// } 753 /// assert_eq!(m.len(), 32); 754 /// for (_, v) in m.iter_mut() { 755 /// *v *= 2; 756 /// } 757 /// for i in 0..32 { 758 /// assert_eq!(m.get(&i).unwrap(), &(i * 2)); 759 /// } 760 /// ``` 761 pub struct IterMut<'a, K: Ord + 'a, V: 'a> { 762 head: NodePtr<K, V>, 763 tail: NodePtr<K, V>, 764 len: usize, 765 _marker: marker::PhantomData<&'a ()>, 766 } 767 768 impl<'a, K: Ord + 'a, V: 'a> Clone for IterMut<'a, K, V> { clone(&self) -> IterMut<'a, K, V>769 fn clone(&self) -> IterMut<'a, K, V> { 770 IterMut { 771 head: self.head, 772 tail: self.tail, 773 len: self.len, 774 _marker: self._marker, 775 } 776 } 777 } 778 779 impl<'a, K: Ord + 'a, V: 'a> Iterator for IterMut<'a, K, V> { 780 type Item = (&'a K, &'a mut V); 781 next(&mut self) -> Option<(&'a K, &'a mut V)>782 fn next(&mut self) -> Option<(&'a K, &'a mut V)> { 783 if self.len == 0 { 784 return None; 785 } 786 787 if self.head.is_null() { 788 return None; 789 } 790 791 let (k, v) = unsafe { (&(*self.head.0).key, &mut (*self.head.0).value) }; 792 self.head = self.head.next(); 793 self.len -= 1; 794 Some((k, v)) 795 } 796 size_hint(&self) -> (usize, Option<usize>)797 fn size_hint(&self) -> (usize, Option<usize>) { 798 (self.len, Some(self.len)) 799 } 800 } 801 802 impl<'a, K: Ord + 'a, V: 'a> DoubleEndedIterator for IterMut<'a, K, V> { 803 #[inline] next_back(&mut self) -> Option<(&'a K, &'a mut V)>804 fn next_back(&mut self) -> Option<(&'a K, &'a mut V)> { 805 if self.len == 0 { 806 return None; 807 } 808 809 if self.tail == self.head { 810 return None; 811 } 812 813 let (k, v) = unsafe { (&(*self.tail.0).key, &mut (*self.tail.0).value) }; 814 self.tail = self.tail.prev(); 815 self.len -= 1; 816 Some((k, v)) 817 } 818 } 819 820 impl<K: Ord, V> IntoIterator for RBTree<K, V> { 821 type Item = (K, V); 822 type IntoIter = IntoIter<K, V>; 823 824 #[inline] into_iter(mut self) -> IntoIter<K, V>825 fn into_iter(mut self) -> IntoIter<K, V> { 826 let iter = if self.root.is_null() { 827 IntoIter { 828 head: NodePtr::null(), 829 tail: NodePtr::null(), 830 len: self.len, 831 } 832 } else { 833 IntoIter { 834 head: self.first_child(), 835 tail: self.last_child(), 836 len: self.len, 837 } 838 }; 839 self.fast_clear(); 840 iter 841 } 842 } 843 844 impl<K: Ord, V> RBTree<K, V> { 845 /// Creates an empty `RBTree`. new() -> RBTree<K, V>846 pub fn new() -> RBTree<K, V> { 847 RBTree { 848 root: NodePtr::null(), 849 len: 0, 850 } 851 } 852 853 /// Returns the len of `RBTree`. 854 #[inline] len(&self) -> usize855 pub fn len(&self) -> usize { 856 self.len 857 } 858 859 /// Returns `true` if the `RBTree` is empty. 860 #[inline] is_empty(&self) -> bool861 pub fn is_empty(&self) -> bool { 862 self.root.is_null() 863 } 864 865 /* 866 * 对红黑树的节点(x)进行左旋转 867 * 868 * 左旋示意图(对节点x进行左旋)： 869 * px px 870 * / / 871 * x y 872 * / \ --(左旋)--> / \ # 873 * lx y x ry 874 * / \ / \ 875 * ly ry lx ly 876 * 877 * 878 */ 879 #[inline] left_rotate(&mut self, mut node: NodePtr<K, V>)880 unsafe fn left_rotate(&mut self, mut node: NodePtr<K, V>) { 881 let mut temp = node.right(); 882 node.set_right(temp.left()); 883 884 if !temp.left().is_null() { 885 temp.left().set_parent(node.clone()); 886 } 887 888 temp.set_parent(node.parent()); 889 if node == self.root { 890 self.root = temp.clone(); 891 } else if node == node.parent().left() { 892 node.parent().set_left(temp.clone()); 893 } else { 894 node.parent().set_right(temp.clone()); 895 } 896 897 temp.set_left(node.clone()); 898 node.set_parent(temp.clone()); 899 } 900 901 /* 902 * 对红黑树的节点(y)进行右旋转 903 * 904 * 右旋示意图(对节点y进行左旋)： 905 * py py 906 * / / 907 * y x 908 * / \ --(右旋)--> / \ # 909 * x ry lx y 910 * / \ / \ # 911 * lx rx rx ry 912 * 913 */ 914 #[inline] right_rotate(&mut self, mut node: NodePtr<K, V>)915 unsafe fn right_rotate(&mut self, mut node: NodePtr<K, V>) { 916 let mut temp = node.left(); 917 node.set_left(temp.right()); 918 919 if !temp.right().is_null() { 920 temp.right().set_parent(node.clone()); 921 } 922 923 temp.set_parent(node.parent()); 924 if node == self.root { 925 self.root = temp.clone(); 926 } else if node == node.parent().right() { 927 node.parent().set_right(temp.clone()); 928 } else { 929 node.parent().set_left(temp.clone()); 930 } 931 932 temp.set_right(node.clone()); 933 node.set_parent(temp.clone()); 934 } 935 936 /// replace value if key exist, if not exist insert it. 937 /// # Examples 938 /// ``` 939 /// use rbtree::RBTree; 940 /// let mut m = RBTree::new(); 941 /// assert_eq!(m.len(), 0); 942 /// m.insert(2, 4); 943 /// assert_eq!(m.len(), 1); 944 /// assert_eq!(m.replace_or_insert(2, 6).unwrap(), 4); 945 /// assert_eq!(m.len(), 1); 946 /// assert_eq!(*m.get(&2).unwrap(), 6); 947 /// ``` 948 #[inline] replace_or_insert(&mut self, k: K, mut v: V) -> Option<V>949 pub fn replace_or_insert(&mut self, k: K, mut v: V) -> Option<V> { 950 let node = self.find_node(&k); 951 if node.is_null() { 952 self.insert(k, v); 953 return None; 954 } 955 956 unsafe { 957 mem::swap(&mut v, &mut (*node.0).value); 958 } 959 960 Some(v) 961 } 962 963 #[inline] insert_fixup(&mut self, mut node: NodePtr<K, V>)964 unsafe fn insert_fixup(&mut self, mut node: NodePtr<K, V>) { 965 let mut parent; 966 let mut gparent; 967 968 while node.parent().is_red_color() { 969 parent = node.parent(); 970 gparent = parent.parent(); 971 //若“父节点”是“祖父节点的左孩子” 972 if parent == gparent.left() { 973 // Case 1条件：叔叔节点是红色 974 let mut uncle = gparent.right(); 975 if !uncle.is_null() && uncle.is_red_color() { 976 uncle.set_black_color(); 977 parent.set_black_color(); 978 gparent.set_red_color(); 979 node = gparent; 980 continue; 981 } 982 983 // Case 2条件：叔叔是黑色，且当前节点是右孩子 984 if parent.right() == node { 985 self.left_rotate(parent); 986 let temp = parent; 987 parent = node; 988 node = temp; 989 } 990 991 // Case 3条件：叔叔是黑色，且当前节点是左孩子。 992 parent.set_black_color(); 993 gparent.set_red_color(); 994 self.right_rotate(gparent); 995 } else { 996 // Case 1条件：叔叔节点是红色 997 let mut uncle = gparent.left(); 998 if !uncle.is_null() && uncle.is_red_color() { 999 uncle.set_black_color(); 1000 parent.set_black_color(); 1001 gparent.set_red_color(); 1002 node = gparent; 1003 continue; 1004 } 1005 1006 // Case 2条件：叔叔是黑色，且当前节点是右孩子 1007 if parent.left() == node { 1008 self.right_rotate(parent); 1009 let temp = parent; 1010 parent = node; 1011 node = temp; 1012 } 1013 1014 // Case 3条件：叔叔是黑色，且当前节点是左孩子。 1015 parent.set_black_color(); 1016 gparent.set_red_color(); 1017 self.left_rotate(gparent); 1018 } 1019 } 1020 self.root.set_black_color(); 1021 } 1022 1023 #[inline] insert(&mut self, k: K, v: V)1024 pub fn insert(&mut self, k: K, v: V) { 1025 self.len += 1; 1026 let mut node = NodePtr::new(k, v); 1027 let mut y = NodePtr::null(); 1028 let mut x = self.root; 1029 1030 while !x.is_null() { 1031 y = x; 1032 match node.cmp(&&mut x) { 1033 Ordering::Less => { 1034 x = x.left(); 1035 } 1036 _ => { 1037 x = x.right(); 1038 } 1039 }; 1040 } 1041 node.set_parent(y); 1042 1043 if y.is_null() { 1044 self.root = node; 1045 } else { 1046 match node.cmp(&&mut y) { 1047 Ordering::Less => { 1048 y.set_left(node); 1049 } 1050 _ => { 1051 y.set_right(node); 1052 } 1053 }; 1054 } 1055 1056 node.set_red_color(); 1057 unsafe { 1058 self.insert_fixup(node); 1059 } 1060 } 1061 1062 #[inline] find_node(&self, k: &K) -> NodePtr<K, V>1063 fn find_node(&self, k: &K) -> NodePtr<K, V> { 1064 if self.root.is_null() { 1065 return NodePtr::null(); 1066 } 1067 let mut temp = &self.root; 1068 unsafe { 1069 loop { 1070 let next = match k.cmp(&(*temp.0).key) { 1071 Ordering::Less => &mut (*temp.0).left, 1072 Ordering::Greater => &mut (*temp.0).right, 1073 Ordering::Equal => return *temp, 1074 }; 1075 if next.is_null() { 1076 break; 1077 } 1078 temp = next; 1079 } 1080 } 1081 NodePtr::null() 1082 } 1083 1084 #[inline] first_child(&self) -> NodePtr<K, V>1085 fn first_child(&self) -> NodePtr<K, V> { 1086 if self.root.is_null() { 1087 NodePtr::null() 1088 } else { 1089 let mut temp = self.root; 1090 while !temp.left().is_null() { 1091 temp = temp.left(); 1092 } 1093 return temp; 1094 } 1095 } 1096 1097 #[inline] last_child(&self) -> NodePtr<K, V>1098 fn last_child(&self) -> NodePtr<K, V> { 1099 if self.root.is_null() { 1100 NodePtr::null() 1101 } else { 1102 let mut temp = self.root; 1103 while !temp.right().is_null() { 1104 temp = temp.right(); 1105 } 1106 return temp; 1107 } 1108 } 1109 1110 #[inline] get_first(&self) -> Option<(&K, &V)>1111 pub fn get_first(&self) -> Option<(&K, &V)> { 1112 let first = self.first_child(); 1113 if first.is_null() { 1114 return None; 1115 } 1116 unsafe { Some((&(*first.0).key, &(*first.0).value)) } 1117 } 1118 1119 #[inline] get_last(&self) -> Option<(&K, &V)>1120 pub fn get_last(&self) -> Option<(&K, &V)> { 1121 let last = self.last_child(); 1122 if last.is_null() { 1123 return None; 1124 } 1125 unsafe { Some((&(*last.0).key, &(*last.0).value)) } 1126 } 1127 1128 #[inline] pop_first(&mut self) -> Option<(K, V)>1129 pub fn pop_first(&mut self) -> Option<(K, V)> { 1130 let first = self.first_child(); 1131 if first.is_null() { 1132 return None; 1133 } 1134 unsafe { Some(self.delete(first)) } 1135 } 1136 1137 #[inline] pop_last(&mut self) -> Option<(K, V)>1138 pub fn pop_last(&mut self) -> Option<(K, V)> { 1139 let last = self.last_child(); 1140 if last.is_null() { 1141 return None; 1142 } 1143 unsafe { Some(self.delete(last)) } 1144 } 1145 1146 #[inline] get_first_mut(&mut self) -> Option<(&K, &mut V)>1147 pub fn get_first_mut(&mut self) -> Option<(&K, &mut V)> { 1148 let first = self.first_child(); 1149 if first.is_null() { 1150 return None; 1151 } 1152 unsafe { Some((&(*first.0).key, &mut (*first.0).value)) } 1153 } 1154 1155 #[inline] get_last_mut(&mut self) -> Option<(&K, &mut V)>1156 pub fn get_last_mut(&mut self) -> Option<(&K, &mut V)> { 1157 let last = self.last_child(); 1158 if last.is_null() { 1159 return None; 1160 } 1161 unsafe { Some((&(*last.0).key, &mut (*last.0).value)) } 1162 } 1163 1164 #[inline] get(&self, k: &K) -> Option<&V>1165 pub fn get(&self, k: &K) -> Option<&V> { 1166 let node = self.find_node(k); 1167 if node.is_null() { 1168 return None; 1169 } 1170 1171 unsafe { Some(&(*node.0).value) } 1172 } 1173 1174 #[inline] get_mut(&mut self, k: &K) -> Option<&mut V>1175 pub fn get_mut(&mut self, k: &K) -> Option<&mut V> { 1176 let node = self.find_node(k); 1177 if node.is_null() { 1178 return None; 1179 } 1180 1181 unsafe { Some(&mut (*node.0).value) } 1182 } 1183 1184 #[inline] contains_key(&self, k: &K) -> bool1185 pub fn contains_key(&self, k: &K) -> bool { 1186 let node = self.find_node(k); 1187 if node.is_null() { 1188 return false; 1189 } 1190 true 1191 } 1192 1193 #[inline] clear_recurse(&mut self, current: NodePtr<K, V>)1194 fn clear_recurse(&mut self, current: NodePtr<K, V>) { 1195 if !current.is_null() { 1196 unsafe { 1197 self.clear_recurse(current.left()); 1198 self.clear_recurse(current.right()); 1199 drop(Box::from_raw(current.0)); 1200 } 1201 } 1202 } 1203 1204 #[inline] clear(&mut self)1205 pub fn clear(&mut self) { 1206 let root = self.root; 1207 self.root = NodePtr::null(); 1208 self.clear_recurse(root); 1209 } 1210 1211 /// Empties the `RBTree` without freeing objects in it. 1212 #[inline] fast_clear(&mut self)1213 fn fast_clear(&mut self) { 1214 self.root = NodePtr::null(); 1215 } 1216 1217 #[inline] remove(&mut self, k: &K) -> Option<V>1218 pub fn remove(&mut self, k: &K) -> Option<V> { 1219 let node = self.find_node(k); 1220 if node.is_null() { 1221 return None; 1222 } 1223 unsafe { Some(self.delete(node).1) } 1224 } 1225 1226 #[inline] delete_fixup(&mut self, mut node: NodePtr<K, V>, mut parent: NodePtr<K, V>)1227 unsafe fn delete_fixup(&mut self, mut node: NodePtr<K, V>, mut parent: NodePtr<K, V>) { 1228 let mut other; 1229 while node != self.root && node.is_black_color() { 1230 if parent.left() == node { 1231 other = parent.right(); 1232 //x的兄弟w是红色的 1233 if other.is_red_color() { 1234 other.set_black_color(); 1235 parent.set_red_color(); 1236 self.left_rotate(parent); 1237 other = parent.right(); 1238 } 1239 1240 //x的兄弟w是黑色，且w的俩个孩子也都是黑色的 1241 if other.left().is_black_color() && other.right().is_black_color() { 1242 other.set_red_color(); 1243 node = parent; 1244 parent = node.parent(); 1245 } else { 1246 //x的兄弟w是黑色的，并且w的左孩子是红色，右孩子为黑色。 1247 if other.right().is_black_color() { 1248 other.left().set_black_color(); 1249 other.set_red_color(); 1250 self.right_rotate(other); 1251 other = parent.right(); 1252 } 1253 //x的兄弟w是黑色的；并且w的右孩子是红色的，左孩子任意颜色。 1254 other.set_color(parent.get_color()); 1255 parent.set_black_color(); 1256 other.right().set_black_color(); 1257 self.left_rotate(parent); 1258 node = self.root; 1259 break; 1260 } 1261 } else { 1262 other = parent.left(); 1263 //x的兄弟w是红色的 1264 if other.is_red_color() { 1265 other.set_black_color(); 1266 parent.set_red_color(); 1267 self.right_rotate(parent); 1268 other = parent.left(); 1269 } 1270 1271 //x的兄弟w是黑色，且w的俩个孩子也都是黑色的 1272 if other.left().is_black_color() && other.right().is_black_color() { 1273 other.set_red_color(); 1274 node = parent; 1275 parent = node.parent(); 1276 } else { 1277 //x的兄弟w是黑色的，并且w的左孩子是红色，右孩子为黑色。 1278 if other.left().is_black_color() { 1279 other.right().set_black_color(); 1280 other.set_red_color(); 1281 self.left_rotate(other); 1282 other = parent.left(); 1283 } 1284 //x的兄弟w是黑色的；并且w的右孩子是红色的，左孩子任意颜色。 1285 other.set_color(parent.get_color()); 1286 parent.set_black_color(); 1287 other.left().set_black_color(); 1288 self.right_rotate(parent); 1289 node = self.root; 1290 break; 1291 } 1292 } 1293 } 1294 1295 node.set_black_color(); 1296 } 1297 1298 #[inline] delete(&mut self, node: NodePtr<K, V>) -> (K, V)1299 unsafe fn delete(&mut self, node: NodePtr<K, V>) -> (K, V) { 1300 let mut child; 1301 let mut parent; 1302 let color; 1303 1304 self.len -= 1; 1305 // 被删除节点的"左右孩子都不为空"的情况。 1306 if !node.left().is_null() && !node.right().is_null() { 1307 // 被删节点的后继节点。(称为"取代节点") 1308 // 用它来取代"被删节点"的位置，然后再将"被删节点"去掉。 1309 let mut replace = node.right().min_node(); 1310 if node == self.root { 1311 self.root = replace; 1312 } else { 1313 if node.parent().left() == node { 1314 node.parent().set_left(replace); 1315 } else { 1316 node.parent().set_right(replace); 1317 } 1318 } 1319 1320 // child是"取代节点"的右孩子，也是需要"调整的节点"。 1321 // "取代节点"肯定不存在左孩子！因为它是一个后继节点。 1322 child = replace.right(); 1323 parent = replace.parent(); 1324 color = replace.get_color(); 1325 if parent == node { 1326 parent = replace; 1327 } else { 1328 if !child.is_null() { 1329 child.set_parent(parent); 1330 } 1331 parent.set_left(child); 1332 replace.set_right(node.right()); 1333 node.right().set_parent(replace); 1334 } 1335 1336 replace.set_parent(node.parent()); 1337 replace.set_color(node.get_color()); 1338 replace.set_left(node.left()); 1339 node.left().set_parent(replace); 1340 1341 if color == Color::Black { 1342 self.delete_fixup(child, parent); 1343 } 1344 1345 let obj = Box::from_raw(node.0); 1346 return obj.pair(); 1347 } 1348 1349 if !node.left().is_null() { 1350 child = node.left(); 1351 } else { 1352 child = node.right(); 1353 } 1354 1355 parent = node.parent(); 1356 color = node.get_color(); 1357 if !child.is_null() { 1358 child.set_parent(parent); 1359 } 1360 1361 if self.root == node { 1362 self.root = child 1363 } else { 1364 if parent.left() == node { 1365 parent.set_left(child); 1366 } else { 1367 parent.set_right(child); 1368 } 1369 } 1370 1371 if color == Color::Black { 1372 self.delete_fixup(child, parent); 1373 } 1374 1375 let obj = Box::from_raw(node.0); 1376 return obj.pair(); 1377 } 1378 1379 /// Return the keys iter 1380 #[inline] keys(&self) -> Keys<K, V>1381 pub fn keys(&self) -> Keys<K, V> { 1382 Keys { inner: self.iter() } 1383 } 1384 1385 /// Return the value iter 1386 #[inline] values(&self) -> Values<K, V>1387 pub fn values(&self) -> Values<K, V> { 1388 Values { inner: self.iter() } 1389 } 1390 1391 /// Return the value iter mut 1392 #[inline] values_mut(&mut self) -> ValuesMut<K, V>1393 pub fn values_mut(&mut self) -> ValuesMut<K, V> { 1394 ValuesMut { 1395 inner: self.iter_mut(), 1396 } 1397 } 1398 1399 /// Return the key and value iter 1400 #[inline] iter(&self) -> Iter<K, V>1401 pub fn iter(&self) -> Iter<K, V> { 1402 Iter { 1403 head: self.first_child(), 1404 tail: self.last_child(), 1405 len: self.len, 1406 _marker: marker::PhantomData, 1407 } 1408 } 1409 1410 /// Return the key and mut value iter 1411 #[inline] iter_mut(&mut self) -> IterMut<K, V>1412 pub fn iter_mut(&mut self) -> IterMut<K, V> { 1413 IterMut { 1414 head: self.first_child(), 1415 tail: self.last_child(), 1416 len: self.len, 1417 _marker: marker::PhantomData, 1418 } 1419 } 1420 } 1421 1422 #[cfg(test)] 1423 mod tests { 1424 1425 #[test] test_insert()1426 fn test_insert() { 1427 use crate::libs::rbtree::RBTree; 1428 let mut m = RBTree::new(); 1429 assert_eq!(m.len(), 0); 1430 m.insert(1, 2); 1431 assert_eq!(m.len(), 1); 1432 m.insert(2, 4); 1433 assert_eq!(m.len(), 2); 1434 m.insert(2, 6); 1435 assert_eq!(m.len(), 3); 1436 assert_eq!(*m.get(&1).unwrap(), 2); 1437 assert_eq!(*m.get(&2).unwrap(), 4); 1438 assert_eq!(*m.get(&2).unwrap(), 4); 1439 } 1440 1441 #[test] test_replace()1442 fn test_replace() { 1443 use crate::libs::rbtree::RBTree; 1444 let mut m = RBTree::new(); 1445 assert_eq!(m.len(), 0); 1446 m.insert(2, 4); 1447 assert_eq!(m.len(), 1); 1448 assert_eq!(m.replace_or_insert(2, 6).unwrap(), 4); 1449 assert_eq!(m.len(), 1); 1450 assert_eq!(*m.get(&2).unwrap(), 6); 1451 } 1452 1453 #[test] test_clone()1454 fn test_clone() { 1455 use crate::libs::rbtree::RBTree; 1456 let mut m = RBTree::new(); 1457 assert_eq!(m.len(), 0); 1458 m.insert(1, 2); 1459 assert_eq!(m.len(), 1); 1460 m.insert(2, 4); 1461 assert_eq!(m.len(), 2); 1462 let m2 = m.clone(); 1463 m.clear(); 1464 assert_eq!(*m2.get(&1).unwrap(), 2); 1465 assert_eq!(*m2.get(&2).unwrap(), 4); 1466 assert_eq!(m2.len(), 2); 1467 } 1468 1469 #[test] test_empty_remove()1470 fn test_empty_remove() { 1471 use crate::libs::rbtree::RBTree; 1472 let mut m: RBTree<isize, bool> = RBTree::new(); 1473 assert_eq!(m.remove(&0), None); 1474 } 1475 1476 #[test] test_empty_iter()1477 fn test_empty_iter() { 1478 use crate::libs::rbtree::RBTree; 1479 let mut m: RBTree<isize, bool> = RBTree::new(); 1480 assert_eq!(m.iter().next(), None); 1481 assert_eq!(m.iter_mut().next(), None); 1482 assert_eq!(m.len(), 0); 1483 assert!(m.is_empty()); 1484 assert_eq!(m.into_iter().next(), None); 1485 } 1486 1487 #[test] test_lots_of_insertions()1488 fn test_lots_of_insertions() { 1489 use crate::libs::rbtree::RBTree; 1490 let mut m = RBTree::new(); 1491 1492 // Try this a few times to make sure we never screw up the hashmap's 1493 // internal state. 1494 for _ in 0..10 { 1495 assert!(m.is_empty()); 1496 1497 for i in 1..101 { 1498 m.insert(i, i); 1499 1500 for j in 1..i + 1 { 1501 let r = m.get(&j); 1502 assert_eq!(r, Some(&j)); 1503 } 1504 1505 for j in i + 1..101 { 1506 let r = m.get(&j); 1507 assert_eq!(r, None); 1508 } 1509 } 1510 1511 for i in 101..201 { 1512 assert!(!m.contains_key(&i)); 1513 } 1514 1515 // remove forwards 1516 for i in 1..101 { 1517 assert!(m.remove(&i).is_some()); 1518 1519 for j in 1..i + 1 { 1520 assert!(!m.contains_key(&j)); 1521 } 1522 1523 for j in i + 1..101 { 1524 assert!(m.contains_key(&j)); 1525 } 1526 } 1527 1528 for i in 1..101 { 1529 assert!(!m.contains_key(&i)); 1530 } 1531 1532 for i in 1..101 { 1533 m.insert(i, i); 1534 } 1535 1536 // remove backwards 1537 for i in (1..101).rev() { 1538 assert!(m.remove(&i).is_some()); 1539 1540 for j in i..101 { 1541 assert!(!m.contains_key(&j)); 1542 } 1543 1544 for j in 1..i { 1545 assert!(m.contains_key(&j)); 1546 } 1547 } 1548 } 1549 } 1550 1551 #[test] test_find_mut()1552 fn test_find_mut() { 1553 use crate::libs::rbtree::RBTree; 1554 let mut m = RBTree::new(); 1555 m.insert(1, 12); 1556 m.insert(2, 8); 1557 m.insert(5, 14); 1558 let new = 100; 1559 match m.get_mut(&5) { 1560 None => panic!(), 1561 Some(x) => *x = new, 1562 } 1563 assert_eq!(m.get(&5), Some(&new)); 1564 } 1565 1566 #[test] test_remove()1567 fn test_remove() { 1568 use crate::libs::rbtree::RBTree; 1569 let mut m = RBTree::new(); 1570 m.insert(1, 2); 1571 assert_eq!(*m.get(&1).unwrap(), 2); 1572 m.insert(5, 3); 1573 assert_eq!(*m.get(&5).unwrap(), 3); 1574 m.insert(9, 4); 1575 assert_eq!(*m.get(&1).unwrap(), 2); 1576 assert_eq!(*m.get(&5).unwrap(), 3); 1577 assert_eq!(*m.get(&9).unwrap(), 4); 1578 assert_eq!(m.remove(&1).unwrap(), 2); 1579 assert_eq!(m.remove(&5).unwrap(), 3); 1580 assert_eq!(m.remove(&9).unwrap(), 4); 1581 assert_eq!(m.len(), 0); 1582 } 1583 1584 #[test] test_is_empty()1585 fn test_is_empty() { 1586 use crate::libs::rbtree::RBTree; 1587 let mut m = RBTree::new(); 1588 m.insert(1, 2); 1589 assert!(!m.is_empty()); 1590 assert!(m.remove(&1).is_some()); 1591 assert!(m.is_empty()); 1592 } 1593 1594 #[test] test_pop()1595 fn test_pop() { 1596 use crate::libs::rbtree::RBTree; 1597 let mut m = RBTree::new(); 1598 m.insert(2, 4); 1599 m.insert(1, 2); 1600 m.insert(3, 6); 1601 assert_eq!(m.len(), 3); 1602 assert_eq!(m.pop_first(), Some((1, 2))); 1603 assert_eq!(m.len(), 2); 1604 assert_eq!(m.pop_last(), Some((3, 6))); 1605 assert_eq!(m.len(), 1); 1606 assert_eq!(m.get_first(), Some((&2, &4))); 1607 assert_eq!(m.get_last(), Some((&2, &4))); 1608 } 1609 1610 #[test] test_iterate()1611 fn test_iterate() { 1612 use crate::libs::rbtree::RBTree; 1613 let mut m = RBTree::new(); 1614 for i in 0..32 { 1615 m.insert(i, i * 2); 1616 } 1617 assert_eq!(m.len(), 32); 1618 1619 let mut observed: u32 = 0; 1620 1621 for (k, v) in m.iter() { 1622 assert_eq!(*v, *k * 2); 1623 observed |= 1 << *k; 1624 } 1625 assert_eq!(observed, 0xFFFF_FFFF); 1626 } 1627 1628 #[test] test_keys()1629 fn test_keys() { 1630 use crate::libs::rbtree::RBTree; 1631 let vec = vec![(1, 'a'), (2, 'b'), (3, 'c')]; 1632 let map: RBTree<_, _> = vec.into_iter().collect(); 1633 let keys: Vec<_> = map.keys().cloned().collect(); 1634 assert_eq!(keys.len(), 3); 1635 assert!(keys.contains(&1)); 1636 assert!(keys.contains(&2)); 1637 assert!(keys.contains(&3)); 1638 } 1639 1640 #[test] test_values()1641 fn test_values() { 1642 use crate::libs::rbtree::RBTree; 1643 let vec = vec![(1, 'a'), (2, 'b'), (3, 'c')]; 1644 let map: RBTree<_, _> = vec.into_iter().collect(); 1645 let values: Vec<_> = map.values().cloned().collect(); 1646 assert_eq!(values.len(), 3); 1647 assert!(values.contains(&'a')); 1648 assert!(values.contains(&'b')); 1649 assert!(values.contains(&'c')); 1650 } 1651 1652 #[test] test_values_mut()1653 fn test_values_mut() { 1654 use crate::libs::rbtree::RBTree; 1655 let vec = vec![(1, 1), (2, 2), (3, 3)]; 1656 let mut map: RBTree<_, _> = vec.into_iter().collect(); 1657 for value in map.values_mut() { 1658 *value = (*value) * 2 1659 } 1660 let values: Vec<_> = map.values().cloned().collect(); 1661 assert_eq!(values.len(), 3); 1662 assert!(values.contains(&2)); 1663 assert!(values.contains(&4)); 1664 assert!(values.contains(&6)); 1665 } 1666 1667 #[test] test_find()1668 fn test_find() { 1669 use crate::libs::rbtree::RBTree; 1670 let mut m = RBTree::new(); 1671 assert!(m.get(&1).is_none()); 1672 m.insert(1, 2); 1673 match m.get(&1) { 1674 None => panic!(), 1675 Some(v) => assert_eq!(*v, 2), 1676 } 1677 } 1678 1679 #[test] test_eq()1680 fn test_eq() { 1681 use crate::libs::rbtree::RBTree; 1682 let mut m1 = RBTree::new(); 1683 m1.insert(1, 2); 1684 m1.insert(2, 3); 1685 m1.insert(3, 4); 1686 1687 let mut m2 = RBTree::new(); 1688 m2.insert(1, 2); 1689 m2.insert(2, 3); 1690 1691 assert!(m1 != m2); 1692 1693 m2.insert(3, 4); 1694 1695 assert_eq!(m1, m2); 1696 } 1697 1698 #[test] test_show()1699 fn test_show() { 1700 use crate::libs::rbtree::RBTree; 1701 let mut map = RBTree::new(); 1702 let empty: RBTree<i32, i32> = RBTree::new(); 1703 1704 map.insert(1, 2); 1705 map.insert(3, 4); 1706 1707 let map_str = format!("{:?}", map); 1708 1709 assert!(map_str == "{1: 2, 3: 4}" || map_str == "{3: 4, 1: 2}"); 1710 assert_eq!(format!("{:?}", empty), "{}"); 1711 } 1712 1713 #[test] test_from_iter()1714 fn test_from_iter() { 1715 use crate::libs::rbtree::RBTree; 1716 let xs = [(1, 1), (2, 2), (3, 3), (4, 4), (5, 5), (6, 6)]; 1717 1718 let map: RBTree<_, _> = xs.iter().cloned().collect(); 1719 1720 for &(k, v) in &xs { 1721 assert_eq!(map.get(&k), Some(&v)); 1722 } 1723 } 1724 1725 #[test] test_size_hint()1726 fn test_size_hint() { 1727 use crate::libs::rbtree::RBTree; 1728 let xs = [(1, 1), (2, 2), (3, 3), (4, 4), (5, 5), (6, 6)]; 1729 1730 let map: RBTree<_, _> = xs.iter().cloned().collect(); 1731 1732 let mut iter = map.iter(); 1733 1734 for _ in iter.by_ref().take(3) {} 1735 1736 assert_eq!(iter.size_hint(), (3, Some(3))); 1737 } 1738 1739 #[test] test_iter_len()1740 fn test_iter_len() { 1741 use crate::libs::rbtree::RBTree; 1742 let xs = [(1, 1), (2, 2), (3, 3), (4, 4), (5, 5), (6, 6)]; 1743 1744 let map: RBTree<_, _> = xs.iter().cloned().collect(); 1745 1746 let mut iter = map.iter(); 1747 1748 for _ in iter.by_ref().take(3) {} 1749 1750 assert_eq!(iter.count(), 3); 1751 } 1752 1753 #[test] test_mut_size_hint()1754 fn test_mut_size_hint() { 1755 use crate::libs::rbtree::RBTree; 1756 let xs = [(1, 1), (2, 2), (3, 3), (4, 4), (5, 5), (6, 6)]; 1757 1758 let mut map: RBTree<_, _> = xs.iter().cloned().collect(); 1759 1760 let mut iter = map.iter_mut(); 1761 1762 for _ in iter.by_ref().take(3) {} 1763 1764 assert_eq!(iter.size_hint(), (3, Some(3))); 1765 } 1766 1767 #[test] test_iter_mut_len()1768 fn test_iter_mut_len() { 1769 use crate::libs::rbtree::RBTree; 1770 let xs = [(1, 1), (2, 2), (3, 3), (4, 4), (5, 5), (6, 6)]; 1771 1772 let mut map: RBTree<_, _> = xs.iter().cloned().collect(); 1773 1774 let mut iter = map.iter_mut(); 1775 1776 for _ in iter.by_ref().take(3) {} 1777 1778 assert_eq!(iter.count(), 3); 1779 } 1780 1781 #[test] test_index()1782 fn test_index() { 1783 use crate::libs::rbtree::RBTree; 1784 let mut map = RBTree::new(); 1785 1786 map.insert(1, 2); 1787 map.insert(2, 1); 1788 map.insert(3, 4); 1789 1790 assert_eq!(map[&2], 1); 1791 } 1792 1793 #[test] 1794 #[should_panic] test_index_nonexistent()1795 fn test_index_nonexistent() { 1796 use crate::libs::rbtree::RBTree; 1797 let mut map = RBTree::new(); 1798 1799 map.insert(1, 2); 1800 map.insert(2, 1); 1801 map.insert(3, 4); 1802 1803 map[&4]; 1804 } 1805 1806 #[test] test_extend_iter()1807 fn test_extend_iter() { 1808 use crate::libs::rbtree::RBTree; 1809 let mut a = RBTree::new(); 1810 a.insert(1, "one"); 1811 let mut b = RBTree::new(); 1812 b.insert(2, "two"); 1813 b.insert(3, "three"); 1814 1815 a.extend(b.into_iter()); 1816 1817 assert_eq!(a.len(), 3); 1818 assert_eq!(a[&1], "one"); 1819 assert_eq!(a[&2], "two"); 1820 assert_eq!(a[&3], "three"); 1821 } 1822 } 1823