rbtree.rs 47 KB

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