Skip to content

Commit

Permalink
Auto merge of #39457 - bvinc:master, r=alexcrichton
Browse files Browse the repository at this point in the history 
Dont segfault if btree range is not in order

This is a first attempt to fix issue #33197.  The issue is that the BTree iterator uses next_unchecked for fast iteration, but it can be tricked into running off the end of the tree and segfaulting if range is called with a maximum that is less than the minimum.

Since a user defined Ord should not determine the safety of BTreeMap, and we still want fast iteration, I've implemented the idea of @gereeter and walk the tree simultaneously searching for both keys to make sure that if our keys diverge, the min key is to the left of our max key.  I currently panic if that is not the case.

Open questions:

1.  Do we want to panic in this error case or do we want to return an empty iterator?  The drain API panics if the range is bad, but drain is given a range of index values, while this is a generic key type.  Panicking is brittle and returning an empty iterator is probably the most flexible and matches what people would want it to do... but artificially returning a BTreeMap::Range with start==end seems like a pretty weird and unnatural thing to do, although it's doable since those fields are not accessible.

The same question for other weird cases:
2.  (Included(101), Excluded(100)) on a map that contains [1,2,3].  Both BTree edges end up on the same part of the map, but comparing the keys shows the range is backwards.
3.  (Excluded(5), Excluded(5)).  The keys are equal but BTree edges end up backwards if the map contains 5.
4.  (Included(5), Excluded(5)).  Should naturally produce an empty iterator, right?
  • Loading branch information
@bors
bors committed Feb 15, 2017
2 parents 62eb605 + fb91047 commit 4c4925f
Show file tree
Hide file tree
Showing 3 changed files with 134 additions and 115 deletions.
204 changes: 91 additions & 113 deletions src/libcollections/btree/map.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -714,6 +714,11 @@ impl<K: Ord, V> BTreeMap<K, V> {
/// `range((Excluded(4), Included(10)))` will yield a left-exclusive, right-inclusive
/// range from 4 to 10.
///
/// # Panics
///
/// Panics if range `start > end`.
/// Panics if range `start == end` and both bounds are `Excluded`.
///
/// # Examples
///
/// Basic usage:
Expand All @@ -739,64 +744,11 @@ impl<K: Ord, V> BTreeMap<K, V> {
pub fn range<T: ?Sized, R>(&self, range: R) -> Range<K, V>
where T: Ord, K: Borrow<T>, R: RangeArgument<T>
{
let min = range.start();
let max = range.end();
let front = match min {
Included(key) => {
match search::search_tree(self.root.as_ref(), key) {
Found(kv_handle) => {
match kv_handle.left_edge().force() {
Leaf(bottom) => bottom,
Internal(internal) => last_leaf_edge(internal.descend()),
}
}
GoDown(bottom) => bottom,
}
}
Excluded(key) => {
match search::search_tree(self.root.as_ref(), key) {
Found(kv_handle) => {
match kv_handle.right_edge().force() {
Leaf(bottom) => bottom,
Internal(internal) => first_leaf_edge(internal.descend()),
}
}
GoDown(bottom) => bottom,
}
}
Unbounded => first_leaf_edge(self.root.as_ref()),
};
let root1 = self.root.as_ref();
let root2 = self.root.as_ref();
let (f, b) = range_search(root1, root2, range);

let back = match max {
Included(key) => {
match search::search_tree(self.root.as_ref(), key) {
Found(kv_handle) => {
match kv_handle.right_edge().force() {
Leaf(bottom) => bottom,
Internal(internal) => first_leaf_edge(internal.descend()),
}
}
GoDown(bottom) => bottom,
}
}
Excluded(key) => {
match search::search_tree(self.root.as_ref(), key) {
Found(kv_handle) => {
match kv_handle.left_edge().force() {
Leaf(bottom) => bottom,
Internal(internal) => last_leaf_edge(internal.descend()),
}
}
GoDown(bottom) => bottom,
}
}
Unbounded => last_leaf_edge(self.root.as_ref()),
};

Range {
front: front,
back: back,
}
Range { front: f, back: b}
}

/// Constructs a mutable double-ended iterator over a sub-range of elements in the map.
Expand All @@ -806,6 +758,11 @@ impl<K: Ord, V> BTreeMap<K, V> {
/// `range((Excluded(4), Included(10)))` will yield a left-exclusive, right-inclusive
/// range from 4 to 10.
///
/// # Panics
///
/// Panics if range `start > end`.
/// Panics if range `start == end` and both bounds are `Excluded`.
///
/// # Examples
///
/// Basic usage:
Expand All @@ -831,66 +788,13 @@ impl<K: Ord, V> BTreeMap<K, V> {
pub fn range_mut<T: ?Sized, R>(&mut self, range: R) -> RangeMut<K, V>
where T: Ord, K: Borrow<T>, R: RangeArgument<T>
{
let min = range.start();
let max = range.end();
let root1 = self.root.as_mut();
let root2 = unsafe { ptr::read(&root1) };

let front = match min {
Included(key) => {
match search::search_tree(root1, key) {
Found(kv_handle) => {
match kv_handle.left_edge().force() {
Leaf(bottom) => bottom,
Internal(internal) => last_leaf_edge(internal.descend()),
}
}
GoDown(bottom) => bottom,
}
}
Excluded(key) => {
match search::search_tree(root1, key) {
Found(kv_handle) => {
match kv_handle.right_edge().force() {
Leaf(bottom) => bottom,
Internal(internal) => first_leaf_edge(internal.descend()),
}
}
GoDown(bottom) => bottom,
}
}
Unbounded => first_leaf_edge(root1),
};

let back = match max {
Included(key) => {
match search::search_tree(root2, key) {
Found(kv_handle) => {
match kv_handle.right_edge().force() {
Leaf(bottom) => bottom,
Internal(internal) => first_leaf_edge(internal.descend()),
}
}
GoDown(bottom) => bottom,
}
}
Excluded(key) => {
match search::search_tree(root2, key) {
Found(kv_handle) => {
match kv_handle.left_edge().force() {
Leaf(bottom) => bottom,
Internal(internal) => last_leaf_edge(internal.descend()),
}
}
GoDown(bottom) => bottom,
}
}
Unbounded => last_leaf_edge(root2),
};
let (f, b) = range_search(root1, root2, range);

RangeMut {
front: front,
back: back,
front: f,
back: b,
_marker: PhantomData,
}
}
Expand Down Expand Up @@ -1827,6 +1731,80 @@ fn last_leaf_edge<BorrowType, K, V>
}
}

fn range_search<BorrowType, K, V, Q: ?Sized, R: RangeArgument<Q>>(
root1: NodeRef<BorrowType, K, V, marker::LeafOrInternal>,
root2: NodeRef<BorrowType, K, V, marker::LeafOrInternal>,
range: R
)-> (Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge>,
Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge>)
where Q: Ord, K: Borrow<Q>
{
match (range.start(), range.end()) {
(Excluded(s), Excluded(e)) if s==e =>
panic!("range start and end are equal and excluded in BTreeMap"),
(Included(s), Included(e)) |
(Included(s), Excluded(e)) |
(Excluded(s), Included(e)) |
(Excluded(s), Excluded(e)) if s>e =>
panic!("range start is greater than range end in BTreeMap"),
_ => {},
};

let mut min_node = root1;
let mut max_node = root2;
let mut min_found = false;
let mut max_found = false;
let mut diverged = false;

loop {
let min_edge = match (min_found, range.start()) {
(false, Included(key)) => match search::search_linear(&min_node, key) {
(i, true) => { min_found = true; i },
(i, false) => i,
},
(false, Excluded(key)) => match search::search_linear(&min_node, key) {
(i, true) => { min_found = true; i+1 },
(i, false) => i,
},
(_, Unbounded) => 0,
(true, Included(_)) => min_node.keys().len(),
(true, Excluded(_)) => 0,
};

let max_edge = match (max_found, range.end()) {
(false, Included(key)) => match search::search_linear(&max_node, key) {
(i, true) => { max_found = true; i+1 },
(i, false) => i,
},
(false, Excluded(key)) => match search::search_linear(&max_node, key) {
(i, true) => { max_found = true; i },
(i, false) => i,
},
(_, Unbounded) => max_node.keys().len(),
(true, Included(_)) => 0,
(true, Excluded(_)) => max_node.keys().len(),
};

if !diverged {
if max_edge < min_edge { panic!("Ord is ill-defined in BTreeMap range") }
if min_edge != max_edge { diverged = true; }
}

let front = Handle::new_edge(min_node, min_edge);
let back = Handle::new_edge(max_node, max_edge);
match (front.force(), back.force()) {
(Leaf(f), Leaf(b)) => {
return (f, b);
},
(Internal(min_int), Internal(max_int)) => {
min_node = min_int.descend();
max_node = max_int.descend();
},
_ => unreachable!("BTreeMap has different depths"),
};
}
}

#[inline(always)]
unsafe fn unwrap_unchecked<T>(val: Option<T>) -> T {
val.unwrap_or_else(|| {
Expand Down
3 changes: 1 addition & 2 deletions src/libcollections/btree/search.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -58,7 +58,7 @@ pub fn search_node<BorrowType, K, V, Type, Q: ?Sized>(
}
}

fn search_linear<BorrowType, K, V, Type, Q: ?Sized>(
pub fn search_linear<BorrowType, K, V, Type, Q: ?Sized>(
node: &NodeRef<BorrowType, K, V, Type>,
key: &Q
) -> (usize, bool)
Expand All @@ -73,4 +73,3 @@ fn search_linear<BorrowType, K, V, Type, Q: ?Sized>(
}
(node.keys().len(), false)
}

42 changes: 42 additions & 0 deletions src/libcollectionstest/btree/map.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -178,6 +178,48 @@ fn test_range_small() {
assert_eq!(j, size - 2);
}

#[test]
fn test_range_equal_empty_cases() {
let map: BTreeMap<_, _> = (0..5).map(|i| (i, i)).collect();
assert_eq!(map.range((Included(2), Excluded(2))).next(), None);
assert_eq!(map.range((Excluded(2), Included(2))).next(), None);
}

#[test]
#[should_panic]
fn test_range_equal_excluded() {
let map: BTreeMap<_, _> = (0..5).map(|i| (i, i)).collect();
map.range((Excluded(2), Excluded(2)));
}

#[test]
#[should_panic]
fn test_range_backwards_1() {
let map: BTreeMap<_, _> = (0..5).map(|i| (i, i)).collect();
map.range((Included(3), Included(2)));
}

#[test]
#[should_panic]
fn test_range_backwards_2() {
let map: BTreeMap<_, _> = (0..5).map(|i| (i, i)).collect();
map.range((Included(3), Excluded(2)));
}

#[test]
#[should_panic]
fn test_range_backwards_3() {
let map: BTreeMap<_, _> = (0..5).map(|i| (i, i)).collect();
map.range((Excluded(3), Included(2)));
}

#[test]
#[should_panic]
fn test_range_backwards_4() {
let map: BTreeMap<_, _> = (0..5).map(|i| (i, i)).collect();
map.range((Excluded(3), Excluded(2)));
}

#[test]
fn test_range_1000() {
let size = 1000;
Expand Down

0 comments on commit 4c4925f

Please sign in to comment.