BTreeMap/Set: merge the implementations of MergeIter

4 files changed, 111 insertions, 97 deletions

diff --git a/library/alloc/src/collections/btree/map.rs b/library/alloc/src/collections/btree/map.rs
index 20c6ebd2292..b45b6f80f11 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<K: fmt::Debug, V: fmt::Debug> fmt::Debug for RangeMut<'_, K, V> {
 }
 
 // An iterator for merging two sorted sequences into one
-struct MergeIter<K, V, I: Iterator<Item = (K, V)>> {
-    left: Peekable<I>,
-    right: Peekable<I>,
-}
+struct MergeIter<K, V, I: Iterator<Item = (K, V)>>(MergeIterInner<I>);
 
 impl<K: Ord, V> BTreeMap<K, V> {
     /// Makes a new empty BTreeMap.
@@ -909,7 +907,7 @@ impl<K: Ord, V> BTreeMap<K, V> {
         // 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<K, V> BTreeMap<K, V> {
     }
 }
 
-impl<K: Ord, V, I: Iterator<Item = (K, V)>> Iterator for MergeIter<K, V, I> {
+impl<K: Ord, V, I> Iterator for MergeIter<K, V, I>
+where
+    I: Iterator<Item = (K, V)> + 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 00000000000..88e6f86c2c6
--- /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<I>
+where
+    I: Iterator,
+{
+    a: I,
+    b: I,
+    peeked: Option<Peeked<I>>,
+}
+
+/// Benchmarks faster than wrapping both iterators in a Peekable.
+#[derive(Clone, Debug)]
+enum Peeked<I: Iterator> {
+    A(I::Item),
+    B(I::Item),
+}
+
+impl<I> Clone for MergeIterInner<I>
+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<I> Debug for MergeIterInner<I>
+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<I> MergeIterInner<I>
+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<Cmp: Fn(&I::Item, &I::Item) -> Ordering>(
+        &mut self,
+        cmp: Cmp,
+    ) -> (Option<I::Item>, Option<I::Item>) {
+        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 bcc50ed5615..3f475898e0f 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 a559e87e4e2..3ad74969bec 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<I>
-where
-    I: Iterator,
-    I::Item: Copy,
-{
-    a: I,
-    b: I,
-    peeked: Option<MergeIterPeeked<I>>,
-}
-
-#[derive(Copy, Clone, Debug)]
-enum MergeIterPeeked<I: Iterator> {
-    A(I::Item),
-    B(I::Item),
-}
-
-impl<I> MergeIterInner<I>
-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<I::Item>, Option<I::Item>) {
-        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<I> Debug for MergeIterInner<I>
-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)
     }