diff --git a/library/alloc/src/collections/btree/map.rs b/library/alloc/src/collections/btree/map.rs index 20c6ebd22928b..b45b6f80f11c8 100644 --- a/library/alloc/src/collections/btree/map.rs +++ b/library/alloc/src/collections/btree/map.rs @@ -2,13 +2,14 @@ use core::borrow::Borrow; use core::cmp::Ordering; use core::fmt::{self, Debug}; use core::hash::{Hash, Hasher}; -use core::iter::{FromIterator, FusedIterator, Peekable}; +use core::iter::{FromIterator, FusedIterator}; use core::marker::PhantomData; use core::mem::{self, ManuallyDrop}; use core::ops::{Index, RangeBounds}; use core::ptr; use super::borrow::DormantMutRef; +use super::merge_iter::MergeIterInner; use super::node::{self, marker, ForceResult::*, Handle, NodeRef}; use super::search::{self, SearchResult::*}; use super::unwrap_unchecked; @@ -454,10 +455,7 @@ impl fmt::Debug for RangeMut<'_, K, V> { } // An iterator for merging two sorted sequences into one -struct MergeIter> { - left: Peekable, - right: Peekable, -} +struct MergeIter>(MergeIterInner); impl BTreeMap { /// Makes a new empty BTreeMap. @@ -909,7 +907,7 @@ impl BTreeMap { // First, we merge `self` and `other` into a sorted sequence in linear time. let self_iter = mem::take(self).into_iter(); let other_iter = mem::take(other).into_iter(); - let iter = MergeIter { left: self_iter.peekable(), right: other_iter.peekable() }; + let iter = MergeIter(MergeIterInner::new(self_iter, other_iter)); // Second, we build a tree from the sorted sequence in linear time. self.from_sorted_iter(iter); @@ -2216,27 +2214,16 @@ impl BTreeMap { } } -impl> Iterator for MergeIter { +impl Iterator for MergeIter +where + I: Iterator + ExactSizeIterator + FusedIterator, +{ type Item = (K, V); + /// If two keys are equal, returns the key/value-pair from the right source. fn next(&mut self) -> Option<(K, V)> { - let res = match (self.left.peek(), self.right.peek()) { - (Some(&(ref left_key, _)), Some(&(ref right_key, _))) => left_key.cmp(right_key), - (Some(_), None) => Ordering::Less, - (None, Some(_)) => Ordering::Greater, - (None, None) => return None, - }; - - // Check which elements comes first and only advance the corresponding iterator. - // If two keys are equal, take the value from `right`. - match res { - Ordering::Less => self.left.next(), - Ordering::Greater => self.right.next(), - Ordering::Equal => { - self.left.next(); - self.right.next() - } - } + let (a_next, b_next) = self.0.nexts(|a: &(K, V), b: &(K, V)| K::cmp(&a.0, &b.0)); + b_next.or(a_next) } } diff --git a/library/alloc/src/collections/btree/merge_iter.rs b/library/alloc/src/collections/btree/merge_iter.rs new file mode 100644 index 0000000000000..88e6f86c2c6f5 --- /dev/null +++ b/library/alloc/src/collections/btree/merge_iter.rs @@ -0,0 +1,96 @@ +use core::cmp::Ordering; +use core::fmt::{self, Debug}; +use core::iter::FusedIterator; + +/// Core of an iterator that merges the output of two ascending iterators, +/// for instance a union or a symmetric difference. +pub struct MergeIterInner +where + I: Iterator, +{ + a: I, + b: I, + peeked: Option>, +} + +/// Benchmarks faster than wrapping both iterators in a Peekable. +#[derive(Clone, Debug)] +enum Peeked { + A(I::Item), + B(I::Item), +} + +impl Clone for MergeIterInner +where + I: Clone + Iterator, + I::Item: Clone, +{ + fn clone(&self) -> Self { + Self { a: self.a.clone(), b: self.b.clone(), peeked: self.peeked.clone() } + } +} + +impl Debug for MergeIterInner +where + I: Iterator + Debug, + I::Item: Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_tuple("MergeIterInner").field(&self.a).field(&self.b).finish() + } +} + +impl MergeIterInner +where + I: ExactSizeIterator + FusedIterator, +{ + /// Creates a new core for an iterator merging a pair of sources. + pub fn new(a: I, b: I) -> Self { + MergeIterInner { a, b, peeked: None } + } + + /// Returns the next pair of items stemming from the pair of sources + /// being merged. If both returned options contain a value, that value + /// is equal and occurs in both sources. If one of the returned options + /// contains a value, that value doesn't occur in the other source. + /// If neither returned option contains a value, iteration has finished + /// and subsequent calls will return the same empty pair. + pub fn nexts Ordering>( + &mut self, + cmp: Cmp, + ) -> (Option, Option) { + let mut a_next; + let mut b_next; + match self.peeked.take() { + Some(Peeked::A(next)) => { + a_next = Some(next); + b_next = self.b.next(); + } + Some(Peeked::B(next)) => { + b_next = Some(next); + a_next = self.a.next(); + } + None => { + a_next = self.a.next(); + b_next = self.b.next(); + } + } + if let (Some(ref a1), Some(ref b1)) = (&a_next, &b_next) { + match cmp(a1, b1) { + Ordering::Less => self.peeked = b_next.take().map(Peeked::B), + Ordering::Greater => self.peeked = a_next.take().map(Peeked::A), + Ordering::Equal => (), + } + } + (a_next, b_next) + } + + /// Returns a pair of upper bounds for the `size_hint` of the final iterator. + pub fn lens(&self) -> (usize, usize) { + match self.peeked { + Some(Peeked::A(_)) => (1 + self.a.len(), self.b.len()), + Some(Peeked::B(_)) => (self.a.len(), 1 + self.b.len()), + _ => (self.a.len(), self.b.len()), + } + } +} diff --git a/library/alloc/src/collections/btree/mod.rs b/library/alloc/src/collections/btree/mod.rs index bcc50ed561587..3f475898e0fbf 100644 --- a/library/alloc/src/collections/btree/mod.rs +++ b/library/alloc/src/collections/btree/mod.rs @@ -1,5 +1,6 @@ mod borrow; pub mod map; +mod merge_iter; mod navigate; mod node; mod remove; diff --git a/library/alloc/src/collections/btree/set.rs b/library/alloc/src/collections/btree/set.rs index a559e87e4e298..3ad74969bec33 100644 --- a/library/alloc/src/collections/btree/set.rs +++ b/library/alloc/src/collections/btree/set.rs @@ -9,6 +9,7 @@ use core::iter::{FromIterator, FusedIterator, Peekable}; use core::ops::{BitAnd, BitOr, BitXor, RangeBounds, Sub}; use super::map::{BTreeMap, Keys}; +use super::merge_iter::MergeIterInner; use super::Recover; // FIXME(conventions): implement bounded iterators @@ -114,77 +115,6 @@ pub struct Range<'a, T: 'a> { iter: super::map::Range<'a, T, ()>, } -/// Core of SymmetricDifference and Union. -/// More efficient than btree.map.MergeIter, -/// and crucially for SymmetricDifference, nexts() reports on both sides. -#[derive(Clone)] -struct MergeIterInner -where - I: Iterator, - I::Item: Copy, -{ - a: I, - b: I, - peeked: Option>, -} - -#[derive(Copy, Clone, Debug)] -enum MergeIterPeeked { - A(I::Item), - B(I::Item), -} - -impl MergeIterInner -where - I: ExactSizeIterator + FusedIterator, - I::Item: Copy + Ord, -{ - fn new(a: I, b: I) -> Self { - MergeIterInner { a, b, peeked: None } - } - - fn nexts(&mut self) -> (Option, Option) { - let mut a_next = match self.peeked { - Some(MergeIterPeeked::A(next)) => Some(next), - _ => self.a.next(), - }; - let mut b_next = match self.peeked { - Some(MergeIterPeeked::B(next)) => Some(next), - _ => self.b.next(), - }; - let ord = match (a_next, b_next) { - (None, None) => Equal, - (_, None) => Less, - (None, _) => Greater, - (Some(a1), Some(b1)) => a1.cmp(&b1), - }; - self.peeked = match ord { - Less => b_next.take().map(MergeIterPeeked::B), - Equal => None, - Greater => a_next.take().map(MergeIterPeeked::A), - }; - (a_next, b_next) - } - - fn lens(&self) -> (usize, usize) { - match self.peeked { - Some(MergeIterPeeked::A(_)) => (1 + self.a.len(), self.b.len()), - Some(MergeIterPeeked::B(_)) => (self.a.len(), 1 + self.b.len()), - _ => (self.a.len(), self.b.len()), - } - } -} - -impl Debug for MergeIterInner -where - I: Iterator + Debug, - I::Item: Copy + Debug, -{ - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_tuple("MergeIterInner").field(&self.a).field(&self.b).finish() - } -} - /// A lazy iterator producing elements in the difference of `BTreeSet`s. /// /// This `struct` is created by the [`difference`] method on [`BTreeSet`]. @@ -1461,7 +1391,7 @@ impl<'a, T: Ord> Iterator for SymmetricDifference<'a, T> { fn next(&mut self) -> Option<&'a T> { loop { - let (a_next, b_next) = self.0.nexts(); + let (a_next, b_next) = self.0.nexts(Self::Item::cmp); if a_next.and(b_next).is_none() { return a_next.or(b_next); } @@ -1555,7 +1485,7 @@ impl<'a, T: Ord> Iterator for Union<'a, T> { type Item = &'a T; fn next(&mut self) -> Option<&'a T> { - let (a_next, b_next) = self.0.nexts(); + let (a_next, b_next) = self.0.nexts(Self::Item::cmp); a_next.or(b_next) }