Skip to content

Commit b5e21db

Browse files
borsbors
committed
Auto merge of #68499 - ssomers:btree_search_tidying, r=Mark-Simulacrum
BtreeMap range_search spruced up #39457 created a lower level entry point for `range_search` to operate on, but it's really not hard to move it up a level of abstraction, making it somewhat shorter and reusing existing unsafe code (`new_edge` is unsafe although it is currently not tagged as such). Benchmark added. Comparison says there's no real difference: ``` >cargo benchcmp old3.txt new3.txt --threshold 5 name old3.txt ns/iter new3.txt ns/iter diff ns/iter diff % speedup btree::map::find_seq_100 19 21 2 10.53% x 0.90 btree::map::range_excluded_unbounded 3,117 2,838 -279 -8.95% x 1.10 btree::map::range_included_unbounded 1,768 1,871 103 5.83% x 0.94 btree::set::intersection_10k_neg_vs_10k_pos 35 37 2 5.71% x 0.95 btree::set::intersection_staggered_100_vs_10k 2,488 2,314 -174 -6.99% x 1.08 btree::set::is_subset_10k_vs_100 3 2 -1 -33.33% x 1.50 ``` r? @Mark-Simulacrum
2 parents f8fd462 + ae03e16 commit b5e21db

File tree

4 files changed

+101
-42
lines changed

4 files changed

+101
-42
lines changed

‎src/liballoc/benches/btree/map.rs

+56
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
use std::collections::BTreeMap;
22
use std::iter::Iterator;
3+
use std::ops::Bound::{Excluded, Unbounded};
34
use std::vec::Vec;
45

56
use rand::{seq::SliceRandom, thread_rng, Rng};
@@ -200,3 +201,58 @@ pub fn first_and_last_100(b: &mut Bencher) {
200201
pub fn first_and_last_10k(b: &mut Bencher) {
201202
bench_first_and_last(b, 10_000);
202203
}
204+
205+
#[bench]
206+
pub fn range_excluded_excluded(b: &mut Bencher) {
207+
let size = 144;
208+
let map: BTreeMap<_, _> = (0..size).map(|i| (i, i)).collect();
209+
b.iter(|| {
210+
for first in 0..size {
211+
for last in first + 1..size {
212+
black_box(map.range((Excluded(first), Excluded(last))));
213+
}
214+
}
215+
});
216+
}
217+
218+
#[bench]
219+
pub fn range_excluded_unbounded(b: &mut Bencher) {
220+
let size = 144;
221+
let map: BTreeMap<_, _> = (0..size).map(|i| (i, i)).collect();
222+
b.iter(|| {
223+
for first in 0..size {
224+
black_box(map.range((Excluded(first), Unbounded)));
225+
}
226+
});
227+
}
228+
229+
#[bench]
230+
pub fn range_included_included(b: &mut Bencher) {
231+
let size = 144;
232+
let map: BTreeMap<_, _> = (0..size).map(|i| (i, i)).collect();
233+
b.iter(|| {
234+
for first in 0..size {
235+
for last in first..size {
236+
black_box(map.range(first..=last));
237+
}
238+
}
239+
});
240+
}
241+
242+
#[bench]
243+
pub fn range_included_unbounded(b: &mut Bencher) {
244+
let size = 144;
245+
let map: BTreeMap<_, _> = (0..size).map(|i| (i, i)).collect();
246+
b.iter(|| {
247+
for first in 0..size {
248+
black_box(map.range(first..));
249+
}
250+
});
251+
}
252+
253+
#[bench]
254+
pub fn range_unbounded_unbounded(b: &mut Bencher) {
255+
let size = 144;
256+
let map: BTreeMap<_, _> = (0..size).map(|i| (i, i)).collect();
257+
b.iter(|| map.range(..));
258+
}

‎src/liballoc/collections/btree/map.rs

+26-40
Original file line numberDiff line numberDiff line change
@@ -1861,65 +1861,51 @@ where
18611861
let mut max_node = root2;
18621862
let mut min_found = false;
18631863
let mut max_found = false;
1864-
let mut diverged = false;
18651864

18661865
loop {
1867-
let min_edge = match (min_found, range.start_bound()) {
1868-
(false, Included(key)) => match search::search_linear(&min_node, key) {
1869-
(i, true) => {
1866+
let front = match (min_found, range.start_bound()) {
1867+
(false, Included(key)) => match search::search_node(min_node, key) {
1868+
Found(kv) => {
18701869
min_found = true;
1871-
i
1870+
kv.left_edge()
18721871
}
1873-
(i, false) => i,
1872+
GoDown(edge) => edge,
18741873
},
1875-
(false, Excluded(key)) => match search::search_linear(&min_node, key) {
1876-
(i, true) => {
1874+
(false, Excluded(key)) => match search::search_node(min_node, key) {
1875+
Found(kv) => {
18771876
min_found = true;
1878-
i + 1
1877+
kv.right_edge()
18791878
}
1880-
(i, false) => i,
1879+
GoDown(edge) => edge,
18811880
},
1882-
(_, Unbounded) => 0,
1883-
(true, Included(_)) => min_node.len(),
1884-
(true, Excluded(_)) => 0,
1881+
(true, Included(_)) => min_node.last_edge(),
1882+
(true, Excluded(_)) => min_node.first_edge(),
1883+
(_, Unbounded) => min_node.first_edge(),
18851884
};
18861885

1887-
let max_edge = match (max_found, range.end_bound()) {
1888-
(false, Included(key)) => match search::search_linear(&max_node, key) {
1889-
(i, true) => {
1886+
let back = match (max_found, range.end_bound()) {
1887+
(false, Included(key)) => match search::search_node(max_node, key) {
1888+
Found(kv) => {
18901889
max_found = true;
1891-
i + 1
1890+
kv.right_edge()
18921891
}
1893-
(i, false) => i,
1892+
GoDown(edge) => edge,
18941893
},
1895-
(false, Excluded(key)) => match search::search_linear(&max_node, key) {
1896-
(i, true) => {
1894+
(false, Excluded(key)) => match search::search_node(max_node, key) {
1895+
Found(kv) => {
18971896
max_found = true;
1898-
i
1897+
kv.left_edge()
18991898
}
1900-
(i, false) => i,
1899+
GoDown(edge) => edge,
19011900
},
1902-
(_, Unbounded) => max_node.len(),
1903-
(true, Included(_)) => 0,
1904-
(true, Excluded(_)) => max_node.len(),
1901+
(true, Included(_)) => max_node.first_edge(),
1902+
(true, Excluded(_)) => max_node.last_edge(),
1903+
(_, Unbounded) => max_node.last_edge(),
19051904
};
19061905

1907-
if !diverged {
1908-
if max_edge < min_edge {
1909-
panic!("Ord is ill-defined in BTreeMap range")
1910-
}
1911-
if min_edge != max_edge {
1912-
diverged = true;
1913-
}
1906+
if front.partial_cmp(&back) == Some(Ordering::Greater) {
1907+
panic!("Ord is ill-defined in BTreeMap range");
19141908
}
1915-
1916-
// Safety guarantee: `min_edge` is always in range for `min_node`, because
1917-
// `min_edge` is unconditionally calculated for each iteration's value of `min_node`,
1918-
// either (if not found) as the edge index returned by `search_linear`,
1919-
// or (if found) as the KV index returned by `search_linear`, possibly + 1.
1920-
// Likewise for `max_node` versus `max_edge`.
1921-
let front = unsafe { Handle::new_edge(min_node, min_edge) };
1922-
let back = unsafe { Handle::new_edge(max_node, max_edge) };
19231909
match (front.force(), back.force()) {
19241910
(Leaf(f), Leaf(b)) => {
19251911
return (f, b);

‎src/liballoc/collections/btree/node.rs

+9
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@
3131
// - A node of length `n` has `n` keys, `n` values, and (in an internal node) `n + 1` edges.
3232
// This implies that even an empty internal node has at least one edge.
3333

34+
use core::cmp::Ordering;
3435
use core::marker::PhantomData;
3536
use core::mem::{self, MaybeUninit};
3637
use core::ptr::{self, NonNull, Unique};
@@ -832,6 +833,14 @@ impl<BorrowType, K, V, NodeType, HandleType> PartialEq
832833
}
833834
}
834835

836+
impl<BorrowType, K, V, NodeType, HandleType> PartialOrd
837+
for Handle<NodeRef<BorrowType, K, V, NodeType>, HandleType>
838+
{
839+
fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
840+
if self.node.node == other.node.node { Some(self.idx.cmp(&other.idx)) } else { None }
841+
}
842+
}
843+
835844
impl<BorrowType, K, V, NodeType, HandleType>
836845
Handle<NodeRef<BorrowType, K, V, NodeType>, HandleType>
837846
{

‎src/liballoc/collections/btree/search.rs

+10-2
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,10 @@ pub enum SearchResult<BorrowType, K, V, FoundType, GoDownType> {
1010
GoDown(Handle<NodeRef<BorrowType, K, V, GoDownType>, marker::Edge>),
1111
}
1212

13+
/// Looks up a given key in a (sub)tree headed by the given node, recursively.
14+
/// Returns a `Found` with the handle of the matching KV, if any. Otherwise,
15+
/// returns a `GoDown` with the handle of the possible leaf edge where the key
16+
/// belongs.
1317
pub fn search_tree<BorrowType, K, V, Q: ?Sized>(
1418
mut node: NodeRef<BorrowType, K, V, marker::LeafOrInternal>,
1519
key: &Q,
@@ -32,6 +36,10 @@ where
3236
}
3337
}
3438

39+
/// Looks up a given key in a given node, without recursion.
40+
/// Returns a `Found` with the handle of the matching KV, if any. Otherwise,
41+
/// returns a `GoDown` with the handle of the edge where the key might be found.
42+
/// If the node is a leaf, a `GoDown` edge is not an actual edge but a possible edge.
3543
pub fn search_node<BorrowType, K, V, Type, Q: ?Sized>(
3644
node: NodeRef<BorrowType, K, V, Type>,
3745
key: &Q,
@@ -50,8 +58,8 @@ where
5058
/// or could exist, and whether it exists in the node itself. If it doesn't
5159
/// exist in the node itself, it may exist in the subtree with that index
5260
/// (if the node has subtrees). If the key doesn't exist in node or subtree,
53-
/// the returned index is the position or subtree to insert at.
54-
pub fn search_linear<BorrowType, K, V, Type, Q: ?Sized>(
61+
/// the returned index is the position or subtree where the key belongs.
62+
fn search_linear<BorrowType, K, V, Type, Q: ?Sized>(
5563
node: &NodeRef<BorrowType, K, V, Type>,
5664
key: &Q,
5765
) -> (usize, bool)

0 commit comments

Comments
 (0)