Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Dont segfault if btree range is not in order #39457

Merged
merged 1 commit into from
Feb 16, 2017
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
204 changes: 91 additions & 113 deletions src/libcollections/btree/map.rs
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
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
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