Skip to content

Commit cf814d6

Browse files
committed
BTreeMap::from_iter: use bulk building to improve the performance
Bulk building is a common technique to increase the performance of building a fresh btree map. Instead of inserting items one-by-one, we sort all the items beforehand then create the BtreeMap in bulk.
1 parent 6a6885c commit cf814d6

File tree

3 files changed

+79
-5
lines changed

3 files changed

+79
-5
lines changed
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
1+
use core::iter::Peekable;
2+
3+
/// A iterator for deduping the key of a sorted iterator.
4+
/// When encountering the duplicated key, only the last key-value pair is yielded.
5+
///
6+
/// Used by [`BTreeMap::bulk_build_from_sorted_iter`].
7+
pub struct DedupSortedIter<K, V, I>
8+
where
9+
I: Iterator<Item = (K, V)>,
10+
{
11+
iter: Peekable<I>,
12+
}
13+
14+
impl<K, V, I> DedupSortedIter<K, V, I>
15+
where
16+
I: Iterator<Item = (K, V)>,
17+
{
18+
pub fn new(iter: I) -> Self {
19+
Self { iter: iter.peekable() }
20+
}
21+
}
22+
23+
impl<K, V, I> Iterator for DedupSortedIter<K, V, I>
24+
where
25+
K: Eq,
26+
I: Iterator<Item = (K, V)>,
27+
{
28+
type Item = (K, V);
29+
30+
fn next(&mut self) -> Option<(K, V)> {
31+
loop {
32+
let next = match self.iter.next() {
33+
Some(next) => next,
34+
None => return None,
35+
};
36+
37+
let peeked = match self.iter.peek() {
38+
Some(peeked) => peeked,
39+
None => return Some(next),
40+
};
41+
42+
if next.0 != peeked.0 {
43+
return Some(next);
44+
}
45+
}
46+
}
47+
}

library/alloc/src/collections/btree/map.rs

+31-5
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
use crate::vec::Vec;
12
use core::borrow::Borrow;
23
use core::cmp::Ordering;
34
use core::fmt::{self, Debug};
@@ -9,6 +10,7 @@ use core::ops::{Index, RangeBounds};
910
use core::ptr;
1011

1112
use super::borrow::DormantMutRef;
13+
use super::dedup_sorted_iter::DedupSortedIter;
1214
use super::navigate::{LazyLeafRange, LeafRange};
1315
use super::node::{self, marker, ForceResult::*, Handle, NodeRef, Root};
1416
use super::search::SearchResult::*;
@@ -1290,6 +1292,18 @@ impl<K, V> BTreeMap<K, V> {
12901292
pub fn into_values(self) -> IntoValues<K, V> {
12911293
IntoValues { inner: self.into_iter() }
12921294
}
1295+
1296+
/// Makes a `BTreeMap` from a sorted iterator.
1297+
pub(crate) fn bulk_build_from_sorted_iter<I>(iter: I) -> Self
1298+
where
1299+
K: Ord,
1300+
I: Iterator<Item = (K, V)>,
1301+
{
1302+
let mut root = Root::new();
1303+
let mut length = 0;
1304+
root.bulk_push(DedupSortedIter::new(iter), &mut length);
1305+
BTreeMap { root: Some(root), length }
1306+
}
12931307
}
12941308

12951309
#[stable(feature = "rust1", since = "1.0.0")]
@@ -1914,9 +1928,15 @@ impl<K, V> FusedIterator for RangeMut<'_, K, V> {}
19141928
#[stable(feature = "rust1", since = "1.0.0")]
19151929
impl<K: Ord, V> FromIterator<(K, V)> for BTreeMap<K, V> {
19161930
fn from_iter<T: IntoIterator<Item = (K, V)>>(iter: T) -> BTreeMap<K, V> {
1917-
let mut map = BTreeMap::new();
1918-
map.extend(iter);
1919-
map
1931+
let mut inputs: Vec<_> = iter.into_iter().collect();
1932+
1933+
if inputs.is_empty() {
1934+
return BTreeMap::new();
1935+
}
1936+
1937+
// use stable sort to preserve the insertion order.
1938+
inputs.sort_by(|a, b| a.0.cmp(&b.0));
1939+
BTreeMap::bulk_build_from_sorted_iter(inputs.into_iter())
19201940
}
19211941
}
19221942

@@ -2025,8 +2045,14 @@ impl<K: Ord, V, const N: usize> From<[(K, V); N]> for BTreeMap<K, V> {
20252045
/// let map2: BTreeMap<_, _> = [(1, 2), (3, 4)].into();
20262046
/// assert_eq!(map1, map2);
20272047
/// ```
2028-
fn from(arr: [(K, V); N]) -> Self {
2029-
core::array::IntoIter::new(arr).collect()
2048+
fn from(mut arr: [(K, V); N]) -> Self {
2049+
if N == 0 {
2050+
return BTreeMap::new();
2051+
}
2052+
2053+
// use stable sort to preserve the insertion order.
2054+
arr.sort_by(|a, b| a.0.cmp(&b.0));
2055+
BTreeMap::bulk_build_from_sorted_iter(core::array::IntoIter::new(arr))
20302056
}
20312057
}
20322058

library/alloc/src/collections/btree/mod.rs

+1
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
mod append;
22
mod borrow;
3+
mod dedup_sorted_iter;
34
mod fix;
45
pub mod map;
56
mod mem;

0 commit comments

Comments
 (0)