xsl/collections/rbtree/map.rs
1mod values;
2use super::{
3 entry::{Entry, OccupiedEntry, VacantEntry},
4 flag::Color,
5 iter::{Iter, IterMut},
6 node::{Node, NodeRef, SearchResult},
7};
8use crate::{
9 alloc::{Allocator, Global},
10 collections::rbtree::{
11 flag::{toggle_rela, LEFT, RIGHT, ROOT},
12 node::OwnedNodeRef,
13 },
14};
15
16use core::{
17 alloc::Layout,
18 borrow::Borrow,
19 cmp::Ordering,
20 fmt::{Debug, Display},
21 ops::Index,
22};
23use values::{Values, ValuesMut};
24
25pub(super) enum NodeDesc<K, V> {
26 Found(OwnedNodeRef<K, V>),
27 NotFound(NdNotFound<K, V>),
28}
29
30pub(super) enum NdNotFound<K, V> {
31 Normal(OwnedNodeRef<K, V>, u8),
32 Root,
33}
34
35pub struct RBTreeMap<K, V, A = Global>
36where
37 A: Allocator + Clone,
38{
39 pub(super) root: NodeRef<K, V>,
40 pub(super) alloc: A,
41 pub(super) length: usize,
42}
43impl<K, V, const N: usize> From<[(K, V); N]> for RBTreeMap<K, V>
44where
45 K: Ord,
46{
47 /// Converts a `[(K, V); N]` into a `BTreeMap<(K, V)>`.
48 ///
49 /// ```
50 /// use xsl::collections::RBTreeMap;
51 ///
52 /// let map1 = RBTreeMap::from([(1, 2), (3, 4), (5, 6), (7, 8), (9, 10)]);
53 /// let map2: RBTreeMap<_, _> = [(1, 2), (3, 4), (5, 6), (7, 8), (9, 10)].into();
54 /// print!("{:?}", map1);
55 /// assert_eq!(map1, map2);
56 /// ```
57 fn from(mut arr: [(K, V); N]) -> Self {
58 if N == 0 {
59 return RBTreeMap::new();
60 }
61
62 // use stable sort to preserve the insertion order.
63 arr.sort_by(|a, b| a.0.cmp(&b.0));
64 RBTreeMap::bulk_build_from_sorted_iter(arr.into_iter(), Global::default())
65 }
66}
67impl<K, V, A> Debug for RBTreeMap<K, V, A>
68where
69 K: Debug,
70 V: Debug,
71 A: Allocator + Clone,
72{
73 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
74 // write!(f, "RBTreeMap {{")?;
75 // for (k, v) in self.iter() {
76 // write!(f, "{:?}: {:?},", k, v)?;
77 // }
78 // write!(f, "}}")
79 use crate::alloc::Vec;
80
81 let mut cur_stack = Vec::new();
82 let mut next_stack = Vec::new();
83 cur_stack.push(self.root.clone());
84 while !cur_stack.is_empty() {
85 for node in cur_stack.iter() {
86 if node.is_none() {
87 write!(f, "[]")?;
88 continue;
89 }
90 write!(f, "{:?}", &**node)?;
91 next_stack.push(node.next[0].clone());
92 next_stack.push(node.next[1].clone());
93 }
94 cur_stack.clear();
95 writeln!(f)?;
96 core::mem::swap(&mut cur_stack, &mut next_stack);
97 }
98 Ok(())
99 }
100}
101
102impl<'a, K, V, A> IntoIterator for &'a RBTreeMap<K, V, A>
103where
104 A: Allocator + Clone,
105{
106 type Item = (&'a K, &'a V);
107 type IntoIter = Iter<'a, K, V>;
108
109 fn into_iter(self) -> Iter<'a, K, V> {
110 self.iter()
111 }
112}
113impl<K, V, A> PartialEq for RBTreeMap<K, V, A>
114where
115 K: PartialEq,
116 V: PartialEq,
117 A: Allocator + Clone,
118{
119 fn eq(&self, other: &RBTreeMap<K, V, A>) -> bool {
120 self.len() == other.len() && self.iter().zip(other).all(|(a, b)| a == b)
121 }
122}
123
124impl<K, V, A> Eq for RBTreeMap<K, V, A>
125where
126 K: Eq,
127 V: Eq,
128 A: Allocator + Clone,
129{
130}
131impl<K, V, A> PartialOrd for RBTreeMap<K, V, A>
132where
133 K: PartialOrd,
134 V: PartialOrd,
135 A: Allocator + Clone,
136{
137 #[inline]
138 fn partial_cmp(&self, other: &RBTreeMap<K, V, A>) -> Option<Ordering> {
139 self.iter().partial_cmp(other.iter())
140 }
141}
142impl<K: Ord, V, A: Allocator + Clone> Extend<(K, V)> for RBTreeMap<K, V, A> {
143 #[inline]
144 fn extend<T: IntoIterator<Item = (K, V)>>(&mut self, iter: T) {
145 iter.into_iter().for_each(move |(k, v)| {
146 self.insert(k, v);
147 });
148 }
149}
150impl<K, V, Q> Index<&Q> for RBTreeMap<K, V>
151where
152 K: Borrow<Q> + Ord,
153 Q: ?Sized + Ord,
154{
155 type Output = V;
156 /// Returns a reference to the value corresponding to the supplied key.
157 ///
158 /// # Panics
159 ///
160 /// Panics if the key is not present in the `BTreeMap`.
161 fn index(&self, index: &Q) -> &Self::Output {
162 self.get(index).expect("no entry found for key")
163 }
164}
165
166impl<'a, K: Ord + Copy, V: Copy, A: Allocator + Clone> Extend<(&'a K, &'a V)>
167 for RBTreeMap<K, V, A>
168{
169 fn extend<I: IntoIterator<Item = (&'a K, &'a V)>>(&mut self, iter: I) {
170 self.extend(iter.into_iter().map(|(&key, &value)| (key, value)));
171 }
172}
173
174impl<K, V, A> Clone for RBTreeMap<K, V, A>
175where
176 K: Clone,
177 V: Clone,
178 A: Allocator + Clone,
179{
180 fn clone(&self) -> Self {
181 fn new_branch<K, V, A>(
182 alloc: &A,
183 src: OwnedNodeRef<K, V>,
184 mut dst: OwnedNodeRef<K, V>,
185 rela: u8,
186 ) -> (OwnedNodeRef<K, V>, OwnedNodeRef<K, V>)
187 where
188 K: Clone,
189 V: Clone,
190 A: Allocator,
191 {
192 let src_child_ptr = src.next[rela as usize].get_owned();
193 let src_child = src_child_ptr.clone();
194 let mut new_node = {
195 #[cfg(debug_assertions)]
196 {
197 OwnedNodeRef::new_in(&alloc)
198 }
199 #[cfg(not(debug_assertions))]
200 {
201 OwnedNodeRef::new_in(alloc.clone())
202 }
203 };
204 new_node.init_from(&*src_child);
205 dst.next[rela as usize] = new_node.get_node_ref();
206 new_node.set_parent(dst, rela);
207 (src_child_ptr, new_node)
208 }
209 use crate::alloc::Vec;
210 let mut new_tree = Self::new_in(self.alloc.clone());
211 if self.is_empty() {
212 return new_tree;
213 }
214 let mut stack = Vec::new();
215 stack.push(new_branch(
216 &new_tree.alloc,
217 self.root.get_owned(),
218 new_tree.root.get_owned(),
219 ROOT,
220 ));
221 while let Some((src, dst)) = stack.pop() {
222 //First determine whether the child node is empty
223 //Simple performance test shows that the following code is faster than the next code
224 if !src.next[0].is_none() {
225 stack.push(new_branch(&new_tree.alloc, src.clone(), dst.clone(), LEFT));
226 }
227 if !src.next[1].is_none() {
228 stack.push(new_branch(&new_tree.alloc, src, dst, RIGHT));
229 }
230 }
231 new_tree.length = self.length;
232 new_tree
233 }
234}
235
236impl<K, V> Display for RBTreeMap<K, V>
237where
238 K: Display,
239 V: Display,
240{
241 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
242 use crate::alloc::Vec;
243
244 let mut cur_stack = Vec::new();
245 let mut next_stack = Vec::new();
246 cur_stack.push(self.root.clone());
247 while !cur_stack.is_empty() {
248 for node in cur_stack.iter() {
249 if node.is_none() {
250 write!(f, "[]")?;
251 continue;
252 }
253 write!(f, "{}", &**node)?;
254 next_stack.push(node.next[0].clone());
255 next_stack.push(node.next[1].clone());
256 }
257 cur_stack.clear();
258 writeln!(f)?;
259 core::mem::swap(&mut cur_stack, &mut next_stack);
260 }
261 Ok(())
262 }
263}
264
265impl<K, V, A> Drop for RBTreeMap<K, V, A>
266where
267 A: Allocator + Clone,
268{
269 fn drop(&mut self) {
270 self.clear();
271 }
272}
273
274impl<K, V, A> RBTreeMap<K, V, A>
275where
276 A: Allocator + Clone,
277{
278 /// Clears the map, removing all elements.
279 ///
280 /// # Examples
281 ///
282 /// ```
283 /// use xsl::collections::RBTreeMap;
284 ///
285 /// let mut map = RBTreeMap::new();
286 /// map.insert(1, "a");
287 /// map.clear();
288 /// assert!(map.is_empty());
289 /// ```
290 pub fn clear(&mut self) {
291 if self.is_empty() {
292 return;
293 }
294 use crate::alloc::Vec;
295 let mut stack = Vec::new();
296 stack.push(self.root.clone());
297 while !stack.is_empty() {
298 let mut node = stack.pop().unwrap();
299 if node.is_none() {
300 continue;
301 }
302 stack.push(node.next[0].clone());
303 stack.push(node.next[1].clone());
304 unsafe {
305 core::ptr::drop_in_place(node.ptr.as_mut().unwrap().as_mut());
306 self.alloc
307 .deallocate(node.ptr.unwrap().cast(), Layout::new::<Node<K, V>>());
308 }
309 }
310 self.length = 0;
311 }
312 /// Returns `true` if the map contains no elements.
313 ///
314 /// # Examples
315 ///
316 /// ```
317 /// use xsl::collections::RBTreeMap;
318 ///
319 /// let mut map = RBTreeMap::new();
320 /// assert!(map.is_empty());
321 /// map.insert(1, "a");
322 /// assert!(!map.is_empty());
323 /// ```
324 #[must_use]
325 pub const fn is_empty(&self) -> bool {
326 self.len() == 0
327 }
328 /// Returns the number of elements in the map.
329 ///
330 /// # Examples
331 ///
332 /// ```
333 /// use xsl::collections::RBTreeMap;
334 ///
335 /// let mut map = RBTreeMap::new();
336 /// assert_eq!(map.len(), 0);
337 /// map.insert(1, "a");
338 /// assert_eq!(map.len(), 1);
339 /// ```
340 #[must_use]
341 pub const fn len(&self) -> usize {
342 self.length
343 }
344 /// Returns the first entry in the map for in-place manipulation.
345 /// The key of this entry is the minimum key in the map.
346 ///
347 /// # Examples
348 ///
349 /// ```
350 /// use xsl::collections::RBTreeMap;
351 ///
352 /// let mut map = RBTreeMap::new();
353 /// map.insert(1, "a");
354 /// map.insert(2, "b");
355 /// if let Some(mut entry) = map.first_entry() {
356 /// if *entry.key() > 0 {
357 /// entry.insert("first");
358 /// }
359 /// }
360 /// assert_eq!(*map.get(&1).unwrap(), "first");
361 /// assert_eq!(*map.get(&2).unwrap(), "b");
362 /// ```
363 pub fn first_entry(&mut self) -> Option<OccupiedEntry<'_, K, V, A>> {
364 self.raw_first().map(|node| OccupiedEntry::new(node, self))
365 }
366 /// Returns the first key-value pair in the map.
367 /// The key in this pair is the minimum key in the map.
368 ///
369 /// # Examples
370 ///
371 /// ```
372 /// use xsl::collections::RBTreeMap;
373 ///
374 /// let mut map = RBTreeMap::new();
375 /// assert_eq!(map.first_key_value(), None);
376 /// map.insert(1, "b");
377 /// map.insert(2, "a");
378 /// assert_eq!(map.first_key_value(), Some((&1, &"b")));
379 /// ```
380 pub fn first_key_value(&self) -> Option<(&K, &V)> {
381 self.raw_first().map(|node| node.into_ref_key_value())
382 }
383 /// Returns the last key-value pair in the map.
384 /// The key in this pair is the maximum key in the map.
385 ///
386 /// # Examples
387 ///
388 /// ```
389 /// use xsl::collections::RBTreeMap;
390 ///
391 /// let mut map = RBTreeMap::new();
392 /// map.insert(1, "b");
393 /// map.insert(2, "a");
394 /// assert_eq!(map.last_key_value(), Some((&2, &"a")));
395 /// ```
396 pub fn last_key_value(&self) -> Option<(&K, &V)> {
397 self.raw_last().map(|node| node.into_ref_key_value())
398 }
399
400 /// Returns the last entry in the map for in-place manipulation.
401 /// The key of this entry is the maximum key in the map.
402 ///
403 /// # Examples
404 ///
405 /// ```
406 /// use xsl::collections::RBTreeMap;
407 ///
408 /// let mut map = RBTreeMap::new();
409 /// map.insert(1, "a");
410 /// map.insert(2, "b");
411 /// if let Some(mut entry) = map.last_entry() {
412 /// if *entry.key() > 0 {
413 /// entry.insert("last");
414 /// }
415 /// }
416 /// assert_eq!(*map.get(&1).unwrap(), "a");
417 /// assert_eq!(*map.get(&2).unwrap(), "last");
418 /// ```
419 pub fn last_entry(&mut self) -> Option<OccupiedEntry<'_, K, V, A>> {
420 self.raw_last().map(|node| OccupiedEntry::new(node, self))
421 }
422 /// Removes and returns the first element in the map.
423 /// The key of this element is the minimum key that was in the map.
424 ///
425 /// # Examples
426 ///
427 /// Draining elements in ascending order, while keeping a usable map each iteration.
428 ///
429 /// ```
430 /// use xsl::collections::RBTreeMap;
431 ///
432 /// let mut map = RBTreeMap::new();
433 /// map.insert(1, "a");
434 /// map.insert(2, "b");
435 /// while let Some((key, _val)) = map.pop_first() {
436 /// assert!(map.iter().all(|(k, _v)| *k > key));
437 /// }
438 /// assert!(map.is_empty());
439 /// ```
440 pub fn pop_first(&mut self) -> Option<(K, V)> {
441 self.first_entry().map(|entry| entry.remove_entry())
442 }
443 /// Removes and returns the last element in the map.
444 /// The key of this element is the maximum key that was in the map.
445 ///
446 /// # Examples
447 ///
448 /// Draining elements in descending order, while keeping a usable map each iteration.
449 ///
450 /// ```
451 /// use xsl::collections::RBTreeMap;
452 ///
453 /// let mut map = RBTreeMap::new();
454 /// map.insert(1, "a");
455 /// map.insert(2, "b");
456 /// while let Some((key, _val)) = map.pop_last() {
457 /// assert!(map.iter().all(|(k, _v)| *k < key));
458 /// }
459 /// assert!(map.is_empty());
460 /// ```
461 pub fn pop_last(&mut self) -> Option<(K, V)> {
462 self.last_entry().map(|entry| entry.remove_entry())
463 }
464 /// Gets an iterator over the entries of the map, sorted by key.
465 ///
466 /// # Examples
467 ///
468 /// ```
469 /// use xsl::collections::RBTreeMap;
470 ///
471 /// let mut map = RBTreeMap::new();
472 /// map.insert(3, "c");
473 /// map.insert(2, "b");
474 /// map.insert(1, "a");
475 ///
476 /// for (key, value) in map.iter() {
477 /// println!("{key}: {value}");
478 /// }
479 ///
480 /// let (first_key, first_value) = map.iter().next().unwrap();
481 /// assert_eq!((*first_key, *first_value), (1, "a"));
482 /// ```
483 pub fn iter(&self) -> Iter<'_, K, V> {
484 if self.is_empty() {
485 Iter::new_empty()
486 } else {
487 Iter::new(self.root.get_owned(), self.length)
488 }
489 }
490 /// Gets a mutable iterator over the entries of the map, sorted by key.
491 ///
492 /// # Examples
493 ///
494 /// ```
495 /// use xsl::collections::RBTreeMap;
496 ///
497 /// let mut map = RBTreeMap::from([
498 /// ("a", 1),
499 /// ("b", 2),
500 /// ("c", 3),
501 /// ]);
502 ///
503 /// // add 10 to the value if the key isn't "a"
504 /// for (key, value) in map.iter_mut() {
505 /// if key != &"a" {
506 /// *value += 10;
507 /// }
508 /// }
509 /// ```
510 pub fn iter_mut(&mut self) -> IterMut<'_, K, V> {
511 if self.is_empty() {
512 IterMut::new_empty()
513 } else {
514 IterMut::new(self.root.get_owned(), self.length)
515 }
516 }
517}
518impl<K, V, A> RBTreeMap<K, V, A>
519where
520 K: Ord,
521 A: Allocator + Clone,
522{
523 /// Inserts a key-value pair into the map.
524 ///
525 /// If the map did not have this key present, `None` is returned.
526 ///
527 /// If the map did have this key present, the value is updated, and the old
528 /// value is returned. The key is not updated, though; this matters for
529 /// types that can be `==` without being identical. See the [module-level
530 /// documentation] for more.
531 ///
532 /// [module-level documentation]: index.html#insert-and-complex-keys
533 ///
534 /// # Examples
535 ///
536 /// ```
537 /// use xsl::collections::RBTreeMap;
538 ///
539 /// let mut map = RBTreeMap::new();
540 /// assert_eq!(map.insert(37, "a"), None);
541 /// assert_eq!(map.is_empty(), false);
542 ///
543 /// map.insert(37, "b");
544 /// assert_eq!(map.insert(37, "c"), Some("b"));
545 /// assert_eq!(map[&37], "c");
546 /// ```
547 pub fn insert(&mut self, key: K, value: V) -> Option<V> {
548 match self.entry(key) {
549 Entry::Occupied(mut e) => Some(e.insert(value)),
550 Entry::Vacant(e) => {
551 e.insert(value);
552 None
553 }
554 }
555 }
556 /// Gets the given key's corresponding entry in the map for in-place manipulation.
557 ///
558 /// # Examples
559 ///
560 /// ```
561 /// use xsl::collections::RBTreeMap;
562 ///
563 /// let mut count: RBTreeMap<&str, usize> = RBTreeMap::new();
564 ///
565 /// // count the number of occurrences of letters in the vec
566 /// for x in ["a", "b", "a", "c", "a", "b"] {
567 /// count.entry(x).and_modify(|curr| *curr += 1).or_insert(1);
568 /// }
569 ///
570 /// assert_eq!(count["a"], 3);
571 /// assert_eq!(count["b"], 2);
572 /// assert_eq!(count["c"], 1);
573 /// ```
574 pub fn entry(&mut self, key: K) -> Entry<'_, K, V, A> {
575 match self.raw_search(&key) {
576 NodeDesc::Found(node) => Entry::Occupied(OccupiedEntry::new(node, self)),
577 NodeDesc::NotFound(nd) => Entry::Vacant(VacantEntry::new(key, nd, self)),
578 }
579 }
580}
581impl<K, V, A> RBTreeMap<K, V, A>
582where
583 A: Allocator + Clone,
584{
585 /// Returns a reference to the value corresponding to the key.
586 ///
587 /// The key may be any borrowed form of the map's key type, but the ordering
588 /// on the borrowed form *must* match the ordering on the key type.
589 ///
590 /// # Examples
591 ///
592 /// ```
593 /// use xsl::collections::RBTreeMap;
594 ///
595 /// let mut map = RBTreeMap::new();
596 /// map.insert(1, "a");
597 /// assert_eq!(map.get(&1), Some(&"a"));
598 /// assert_eq!(map.get(&2), None);
599 /// ```
600 pub fn get<Q>(&self, key: &Q) -> Option<&V>
601 where
602 K: Borrow<Q>,
603 Q: ?Sized + Ord,
604 {
605 match self.raw_search(key) {
606 NodeDesc::Found(node) => Some(&node.into_ref().key_value.1),
607 NodeDesc::NotFound(_) => None,
608 }
609 }
610 /// Returns a mutable reference to the value corresponding to the key.
611 ///
612 /// The key may be any borrowed form of the map's key type, but the ordering
613 /// on the borrowed form *must* match the ordering on the key type.
614 ///
615 /// # Examples
616 ///
617 /// ```
618 /// use xsl::collections::RBTreeMap;
619 ///
620 /// let mut map = RBTreeMap::new();
621 /// map.insert(1, "a");
622 /// if let Some(x) = map.get_mut(&1) {
623 /// *x = "b";
624 /// }
625 /// assert_eq!(map[&1], "b");
626 /// ```
627 // See `get` for implementation notes, this is basically a copy-paste with mut's added
628 pub fn get_mut<Q>(&mut self, key: &Q) -> Option<&mut V>
629 where
630 K: Borrow<Q>,
631 Q: ?Sized + Ord,
632 {
633 match self.raw_search(key) {
634 NodeDesc::Found(node) => Some(&mut node.into_mut().key_value.1),
635 NodeDesc::NotFound(_) => None,
636 }
637 }
638
639 /// Returns `true` if the map contains a value for the specified key.
640 ///
641 /// The key may be any borrowed form of the map's key type, but the ordering
642 /// on the borrowed form *must* match the ordering on the key type.
643 ///
644 /// # Examples
645 ///
646 /// ```
647 /// use xsl::collections::RBTreeMap;
648 ///
649 /// let mut map = RBTreeMap::new();
650 /// map.insert(1, "a");
651 /// assert_eq!(map.contains_key(&1), true);
652 /// assert_eq!(map.contains_key(&2), false);
653 /// ```
654 pub fn contains_key<Q>(&self, key: &Q) -> bool
655 where
656 K: Borrow<Q>,
657 Q: ?Sized + Ord,
658 {
659 self.get(key).is_some()
660 }
661 /// Returns the key-value pair corresponding to the supplied key.
662 ///
663 /// The supplied key may be any borrowed form of the map's key type, but the ordering
664 /// on the borrowed form *must* match the ordering on the key type.
665 ///
666 /// # Examples
667 ///
668 /// ```
669 /// use xsl::collections::RBTreeMap;
670 ///
671 /// let mut map = RBTreeMap::new();
672 /// map.insert(1, "a");
673 /// assert_eq!(map.get_key_value(&1), Some((&1, &"a")));
674 /// assert_eq!(map.get_key_value(&2), None);
675 /// ```
676 pub fn get_key_value<Q>(&self, key: &Q) -> Option<(&K, &V)>
677 where
678 K: Borrow<Q>,
679 Q: ?Sized + Ord,
680 {
681 match self.raw_search(key) {
682 NodeDesc::Found(node) => Some(node.into_ref_key_value()),
683 NodeDesc::NotFound(_) => None,
684 }
685 }
686
687 /// Removes a key from the map, returning the stored key and value if the key
688 /// was previously in the map.
689 ///
690 /// The key may be any borrowed form of the map's key type, but the ordering
691 /// on the borrowed form *must* match the ordering on the key type.
692 ///
693 /// # Examples
694 ///
695 /// ```
696 /// use xsl::collections::RBTreeMap;
697 ///
698 /// let mut map = RBTreeMap::new();
699 /// map.insert(1, "a");
700 /// assert_eq!(map.remove_entry(&1), Some((1, "a")));
701 /// assert_eq!(map.remove_entry(&1), None);
702 /// ```
703 pub fn remove_entry<Q>(&mut self, key: &Q) -> Option<(K, V)>
704 where
705 K: Borrow<Q>,
706 Q: ?Sized + Ord,
707 {
708 match self.raw_search(key) {
709 NodeDesc::Found(node) => Some(self.raw_remove(node)),
710 NodeDesc::NotFound(_) => None,
711 }
712 }
713 /// Removes a key from the map, returning the value at the key if the key
714 /// was previously in the map.
715 ///
716 /// The key may be any borrowed form of the map's key type, but the ordering
717 /// on the borrowed form *must* match the ordering on the key type.
718 ///
719 /// # Examples
720 ///
721 /// ```
722 /// use xsl::collections::RBTreeMap;
723 ///
724 /// let mut map = RBTreeMap::new();
725 /// map.insert(1, "a");
726 /// assert_eq!(map.remove(&1), Some("a"));
727 /// assert_eq!(map.remove(&1), None);
728 /// ```
729 pub fn remove<Q>(&mut self, key: &Q) -> Option<V>
730 where
731 K: Borrow<Q>,
732 Q: ?Sized + Ord,
733 {
734 self.remove_entry(key).map(|(_, v)| v)
735 }
736 /// Gets an iterator over the values of the map, in order by key.
737 ///
738 /// # Examples
739 ///
740 /// ```
741 /// use xsl::collections::RBTreeMap;
742 ///
743 /// let mut a = RBTreeMap::new();
744 /// a.insert(1, "hello");
745 /// a.insert(2, "goodbye");
746 ///
747 /// let values: Vec<&str> = a.values().cloned().collect();
748 /// assert_eq!(values, ["hello", "goodbye"]);
749 /// ```
750 pub fn values(&self) -> Values<'_, K, V> {
751 Values::new(self.iter())
752 }
753 /// Gets a mutable iterator over the values of the map, in order by key.
754 ///
755 /// # Examples
756 ///
757 /// ```
758 /// use xsl::collections::RBTreeMap;
759 ///
760 /// let mut a = RBTreeMap::new();
761 /// a.insert(1, String::from("hello"));
762 /// a.insert(2, String::from("goodbye"));
763 ///
764 /// for value in a.values_mut() {
765 /// value.push_str("!");
766 /// }
767 ///
768 /// let values: Vec<String> = a.values().cloned().collect();
769 /// assert_eq!(values, [String::from("hello!"),
770 /// String::from("goodbye!")]);
771 /// ```
772 pub fn values_mut(&mut self) -> ValuesMut<'_, K, V> {
773 ValuesMut::new(self.iter_mut())
774 }
775}
776impl<K, V, A> RBTreeMap<K, V, A>
777where
778 A: Allocator + Clone,
779{
780 pub(super) fn new_in(alloc: A) -> Self {
781 RBTreeMap {
782 root: NodeRef::none(),
783 alloc,
784 length: 0,
785 }
786 }
787}
788impl<K, V> RBTreeMap<K, V> {
789 pub fn new() -> Self {
790 let alloc = Global::default();
791 RBTreeMap {
792 root: NodeRef::none(),
793 alloc,
794 length: 0,
795 }
796 }
797}
798impl<K, V, A> RBTreeMap<K, V, A>
799where
800 A: Allocator + Clone,
801{
802 pub fn bulk_build_from_sorted_iter<I>(iter: I, alloc: A) -> Self
803 where
804 I: IntoIterator<Item = (K, V)>,
805 K: Ord,
806 A: Allocator + Clone,
807 {
808 let mut tree = Self::new_in(alloc);
809 for (k, v) in iter {
810 tree.insert(k, v);
811 }
812 tree
813 }
814}
815
816impl<K, V, A> RBTreeMap<K, V, A>
817where
818 A: Allocator + Clone,
819{
820 pub(super) fn raw_remove(&mut self, node: OwnedNodeRef<K, V>) -> (K, V) {
821 let kv = unsafe { core::mem::transmute_copy(&node.key_value) };
822 fn replace<K, V>(mut node: OwnedNodeRef<K, V>) -> OwnedNodeRef<K, V> {
823 if node.next[1].is_none() {
824 if node.next[0].is_none() {
825 return node;
826 }
827 unsafe {
828 core::ptr::copy_nonoverlapping(&node.next[0].key_value, &mut node.key_value, 1);
829 }
830 return node.next[0].get_owned();
831 }
832 let repl_node = unsafe { node.next[1].get_owned().min() };
833 unsafe {
834 core::ptr::copy_nonoverlapping(&repl_node.key_value, &mut node.key_value, 1);
835 }
836 return replace(repl_node);
837 }
838 let repl_node = replace(node);
839 let mut parent = repl_node.parent.clone();
840 let rela = repl_node.flag.rela();
841 let color = repl_node.flag.color();
842 self.length -= 1;
843 unsafe {
844 self.alloc
845 .deallocate(repl_node.unwrap().cast(), Layout::new::<Node<K, V>>());
846 }
847 if color == Color::RED {
848 parent.next[rela as usize] = NodeRef::none();
849 return kv;
850 }
851 if rela == ROOT {
852 self.root = NodeRef::none();
853 return kv;
854 }
855 let toggle_rela = toggle_rela(rela);
856 parent.next[rela as usize] = NodeRef::none();
857 let mut brother = parent.next[toggle_rela as usize].get_owned();
858 let prela = parent.flag.rela();
859 if brother.flag.is_red() {
860 if parent.flag.is_root() {
861 brother.flag.set_root();
862 self.root = brother.get_node_ref();
863 } else {
864 parent.parent.set_child(brother.clone(), prela);
865 brother.flag.set_black();
866 }
867 parent.next[toggle_rela as usize] = NodeRef::none();
868 let mut nephew = brother.next[rela as usize].get_owned();
869 match (
870 nephew.next[rela as usize].clone().into_owned(),
871 nephew.next[toggle_rela as usize].clone().into_owned(),
872 ) {
873 (None, _) => {
874 parent.flag.set_red();
875 nephew.set_child(parent, rela);
876 }
877 (Some(mut lgnephew), Some(mut rgnephew)) => {
878 nephew.flag.set_red();
879 lgnephew.flag.set_black();
880 rgnephew.flag.set_black();
881 parent.flag.set_red();
882 lgnephew.set_child(parent, rela);
883 }
884 (Some(mut lgnephew), None) => {
885 parent.next[toggle_rela as usize] = NodeRef::none();
886 nephew.next[rela as usize] = NodeRef::none();
887
888 brother.set_child(lgnephew.clone(), rela);
889 lgnephew.set_child(nephew, toggle_rela);
890 lgnephew.set_child(parent, rela);
891 }
892 };
893 } else {
894 match (
895 brother.next[rela as usize].clone().into_owned(),
896 brother.next[toggle_rela as usize].clone().into_owned(),
897 ) {
898 (None, None) => {
899 brother.flag.set_red();
900 if !parent.flag.is_root() {
901 if parent.flag.is_black() {
902 if let Some(new_root) = parent.parent.rasie(prela) {
903 self.root = new_root.get_node_ref();
904 }
905 } else {
906 parent.flag.set_black();
907 }
908 }
909 }
910 (Some(mut lnephew), None) => {
911 if parent.flag.is_root() {
912 lnephew.flag.set_root();
913 self.root = lnephew.get_node_ref();
914 } else {
915 parent.parent.set_child(lnephew.clone(), prela);
916 }
917 if parent.flag.is_black() {
918 lnephew.flag.set_black();
919 } else {
920 parent.flag.set_black();
921 }
922 parent.next[toggle_rela as usize] = NodeRef::none();
923 brother.next[rela as usize] = NodeRef::none();
924 lnephew.set_child(brother, toggle_rela);
925 lnephew.set_child(parent, rela);
926 }
927 (lnephew, Some(mut rnephew)) => {
928 if parent.flag.is_root() {
929 brother.flag.set_root();
930 self.root = brother.get_node_ref();
931 } else {
932 parent.parent.set_child(brother.clone(), prela);
933 }
934 parent.next[toggle_rela as usize] = NodeRef::none();
935 match lnephew {
936 Some(mut lnephew) => {
937 brother.flag.set_color(parent.flag.color());
938 lnephew.flag.set_black();
939 rnephew.flag.set_black();
940 parent.flag.set_red();
941 lnephew.set_child(parent, rela);
942 }
943 None => {
944 rnephew.flag.set_color(parent.flag.color());
945 brother.set_child(parent, rela);
946 }
947 }
948 }
949 }
950 }
951 kv
952 }
953 pub(super) fn raw_search<Q>(&self, key: &Q) -> NodeDesc<K, V>
954 where
955 K: Borrow<Q>,
956 Q: ?Sized + Ord,
957 {
958 if self.len() == 0 {
959 return NodeDesc::NotFound(NdNotFound::Root);
960 }
961 match self.root.get_owned().search(key) {
962 SearchResult::Found(node) => NodeDesc::Found(node),
963 SearchResult::NotFound(node, rela) => {
964 NodeDesc::NotFound(NdNotFound::Normal(node, rela))
965 }
966 }
967 }
968 pub fn raw_first(&self) -> Option<OwnedNodeRef<K, V>> {
969 if self.is_empty() {
970 return None;
971 }
972 Some(unsafe { self.root.get_owned().min() })
973 }
974 pub fn raw_last(&self) -> Option<OwnedNodeRef<K, V>> {
975 if self.is_empty() {
976 return None;
977 }
978 Some(unsafe { self.root.get_owned().max() })
979 }
980}
981mod tests;