Skip to content

Commit fb91047

Browse files
bvincbvinc
committed
Dont segfault if btree range is not in order
1 parent 24a70eb commit fb91047

File tree

3 files changed

+134
-115
lines changed

3 files changed

+134
-115
lines changed

‎src/libcollections/btree/map.rs‎

Lines changed: 91 additions & 113 deletions
Original file line numberDiff line numberDiff line change
@@ -714,6 +714,11 @@ impl<K: Ord, V> BTreeMap<K, V> {
714714
/// `range((Excluded(4), Included(10)))` will yield a left-exclusive, right-inclusive
715715
/// range from 4 to 10.
716716
///
717+
/// # Panics
718+
///
719+
/// Panics if range `start > end`.
720+
/// Panics if range `start == end` and both bounds are `Excluded`.
721+
///
717722
/// # Examples
718723
///
719724
/// Basic usage:
@@ -739,64 +744,11 @@ impl<K: Ord, V> BTreeMap<K, V> {
739744
pub fn range<T: ?Sized, R>(&self, range: R) -> Range<K, V>
740745
where T: Ord, K: Borrow<T>, R: RangeArgument<T>
741746
{
742-
let min = range.start();
743-
let max = range.end();
744-
let front = match min {
745-
Included(key) => {
746-
match search::search_tree(self.root.as_ref(), key) {
747-
Found(kv_handle) => {
748-
match kv_handle.left_edge().force() {
749-
Leaf(bottom) => bottom,
750-
Internal(internal) => last_leaf_edge(internal.descend()),
751-
}
752-
}
753-
GoDown(bottom) => bottom,
754-
}
755-
}
756-
Excluded(key) => {
757-
match search::search_tree(self.root.as_ref(), key) {
758-
Found(kv_handle) => {
759-
match kv_handle.right_edge().force() {
760-
Leaf(bottom) => bottom,
761-
Internal(internal) => first_leaf_edge(internal.descend()),
762-
}
763-
}
764-
GoDown(bottom) => bottom,
765-
}
766-
}
767-
Unbounded => first_leaf_edge(self.root.as_ref()),
768-
};
747+
let root1 = self.root.as_ref();
748+
let root2 = self.root.as_ref();
749+
let (f, b) = range_search(root1, root2, range);
769750

770-
let back = match max {
771-
Included(key) => {
772-
match search::search_tree(self.root.as_ref(), key) {
773-
Found(kv_handle) => {
774-
match kv_handle.right_edge().force() {
775-
Leaf(bottom) => bottom,
776-
Internal(internal) => first_leaf_edge(internal.descend()),
777-
}
778-
}
779-
GoDown(bottom) => bottom,
780-
}
781-
}
782-
Excluded(key) => {
783-
match search::search_tree(self.root.as_ref(), key) {
784-
Found(kv_handle) => {
785-
match kv_handle.left_edge().force() {
786-
Leaf(bottom) => bottom,
787-
Internal(internal) => last_leaf_edge(internal.descend()),
788-
}
789-
}
790-
GoDown(bottom) => bottom,
791-
}
792-
}
793-
Unbounded => last_leaf_edge(self.root.as_ref()),
794-
};
795-
796-
Range {
797-
front: front,
798-
back: back,
799-
}
751+
Range { front: f, back: b}
800752
}
801753

802754
/// Constructs a mutable double-ended iterator over a sub-range of elements in the map.
@@ -806,6 +758,11 @@ impl<K: Ord, V> BTreeMap<K, V> {
806758
/// `range((Excluded(4), Included(10)))` will yield a left-exclusive, right-inclusive
807759
/// range from 4 to 10.
808760
///
761+
/// # Panics
762+
///
763+
/// Panics if range `start > end`.
764+
/// Panics if range `start == end` and both bounds are `Excluded`.
765+
///
809766
/// # Examples
810767
///
811768
/// Basic usage:
@@ -831,66 +788,13 @@ impl<K: Ord, V> BTreeMap<K, V> {
831788
pub fn range_mut<T: ?Sized, R>(&mut self, range: R) -> RangeMut<K, V>
832789
where T: Ord, K: Borrow<T>, R: RangeArgument<T>
833790
{
834-
let min = range.start();
835-
let max = range.end();
836791
let root1 = self.root.as_mut();
837792
let root2 = unsafe { ptr::read(&root1) };
838-
839-
let front = match min {
840-
Included(key) => {
841-
match search::search_tree(root1, key) {
842-
Found(kv_handle) => {
843-
match kv_handle.left_edge().force() {
844-
Leaf(bottom) => bottom,
845-
Internal(internal) => last_leaf_edge(internal.descend()),
846-
}
847-
}
848-
GoDown(bottom) => bottom,
849-
}
850-
}
851-
Excluded(key) => {
852-
match search::search_tree(root1, key) {
853-
Found(kv_handle) => {
854-
match kv_handle.right_edge().force() {
855-
Leaf(bottom) => bottom,
856-
Internal(internal) => first_leaf_edge(internal.descend()),
857-
}
858-
}
859-
GoDown(bottom) => bottom,
860-
}
861-
}
862-
Unbounded => first_leaf_edge(root1),
863-
};
864-
865-
let back = match max {
866-
Included(key) => {
867-
match search::search_tree(root2, key) {
868-
Found(kv_handle) => {
869-
match kv_handle.right_edge().force() {
870-
Leaf(bottom) => bottom,
871-
Internal(internal) => first_leaf_edge(internal.descend()),
872-
}
873-
}
874-
GoDown(bottom) => bottom,
875-
}
876-
}
877-
Excluded(key) => {
878-
match search::search_tree(root2, key) {
879-
Found(kv_handle) => {
880-
match kv_handle.left_edge().force() {
881-
Leaf(bottom) => bottom,
882-
Internal(internal) => last_leaf_edge(internal.descend()),
883-
}
884-
}
885-
GoDown(bottom) => bottom,
886-
}
887-
}
888-
Unbounded => last_leaf_edge(root2),
889-
};
793+
let (f, b) = range_search(root1, root2, range);
890794

891795
RangeMut {
892-
front: front,
893-
back: back,
796+
front: f,
797+
back: b,
894798
_marker: PhantomData,
895799
}
896800
}
@@ -1827,6 +1731,80 @@ fn last_leaf_edge<BorrowType, K, V>
18271731
}
18281732
}
18291733

1734+
fn range_search<BorrowType, K, V, Q: ?Sized, R: RangeArgument<Q>>(
1735+
root1: NodeRef<BorrowType, K, V, marker::LeafOrInternal>,
1736+
root2: NodeRef<BorrowType, K, V, marker::LeafOrInternal>,
1737+
range: R
1738+
)-> (Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge>,
1739+
Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge>)
1740+
where Q: Ord, K: Borrow<Q>
1741+
{
1742+
match (range.start(), range.end()) {
1743+
(Excluded(s), Excluded(e)) if s==e =>
1744+
panic!("range start and end are equal and excluded in BTreeMap"),
1745+
(Included(s), Included(e)) |
1746+
(Included(s), Excluded(e)) |
1747+
(Excluded(s), Included(e)) |
1748+
(Excluded(s), Excluded(e)) if s>e =>
1749+
panic!("range start is greater than range end in BTreeMap"),
1750+
_ => {},
1751+
};
1752+
1753+
let mut min_node = root1;
1754+
let mut max_node = root2;
1755+
let mut min_found = false;
1756+
let mut max_found = false;
1757+
let mut diverged = false;
1758+
1759+
loop {
1760+
let min_edge = match (min_found, range.start()) {
1761+
(false, Included(key)) => match search::search_linear(&min_node, key) {
1762+
(i, true) => { min_found = true; i },
1763+
(i, false) => i,
1764+
},
1765+
(false, Excluded(key)) => match search::search_linear(&min_node, key) {
1766+
(i, true) => { min_found = true; i+1 },
1767+
(i, false) => i,
1768+
},
1769+
(_, Unbounded) => 0,
1770+
(true, Included(_)) => min_node.keys().len(),
1771+
(true, Excluded(_)) => 0,
1772+
};
1773+
1774+
let max_edge = match (max_found, range.end()) {
1775+
(false, Included(key)) => match search::search_linear(&max_node, key) {
1776+
(i, true) => { max_found = true; i+1 },
1777+
(i, false) => i,
1778+
},
1779+
(false, Excluded(key)) => match search::search_linear(&max_node, key) {
1780+
(i, true) => { max_found = true; i },
1781+
(i, false) => i,
1782+
},
1783+
(_, Unbounded) => max_node.keys().len(),
1784+
(true, Included(_)) => 0,
1785+
(true, Excluded(_)) => max_node.keys().len(),
1786+
};
1787+
1788+
if !diverged {
1789+
if max_edge < min_edge { panic!("Ord is ill-defined in BTreeMap range") }
1790+
if min_edge != max_edge { diverged = true; }
1791+
}
1792+
1793+
let front = Handle::new_edge(min_node, min_edge);
1794+
let back = Handle::new_edge(max_node, max_edge);
1795+
match (front.force(), back.force()) {
1796+
(Leaf(f), Leaf(b)) => {
1797+
return (f, b);
1798+
},
1799+
(Internal(min_int), Internal(max_int)) => {
1800+
min_node = min_int.descend();
1801+
max_node = max_int.descend();
1802+
},
1803+
_ => unreachable!("BTreeMap has different depths"),
1804+
};
1805+
}
1806+
}
1807+
18301808
#[inline(always)]
18311809
unsafe fn unwrap_unchecked<T>(val: Option<T>) -> T {
18321810
val.unwrap_or_else(|| {

‎src/libcollections/btree/search.rs‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -58,7 +58,7 @@ pub fn search_node<BorrowType, K, V, Type, Q: ?Sized>(
5858
}
5959
}
6060

61-
fn search_linear<BorrowType, K, V, Type, Q: ?Sized>(
61+
pub fn search_linear<BorrowType, K, V, Type, Q: ?Sized>(
6262
node: &NodeRef<BorrowType, K, V, Type>,
6363
key: &Q
6464
) -> (usize, bool)
@@ -73,4 +73,3 @@ fn search_linear<BorrowType, K, V, Type, Q: ?Sized>(
7373
}
7474
(node.keys().len(), false)
7575
}
76-

‎src/libcollectionstest/btree/map.rs‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -178,6 +178,48 @@ fn test_range_small() {
178178
assert_eq!(j, size - 2);
179179
}
180180

181+
#[test]
182+
fn test_range_equal_empty_cases() {
183+
let map: BTreeMap<_, _> = (0..5).map(|i| (i, i)).collect();
184+
assert_eq!(map.range((Included(2), Excluded(2))).next(), None);
185+
assert_eq!(map.range((Excluded(2), Included(2))).next(), None);
186+
}
187+
188+
#[test]
189+
#[should_panic]
190+
fn test_range_equal_excluded() {
191+
let map: BTreeMap<_, _> = (0..5).map(|i| (i, i)).collect();
192+
map.range((Excluded(2), Excluded(2)));
193+
}
194+
195+
#[test]
196+
#[should_panic]
197+
fn test_range_backwards_1() {
198+
let map: BTreeMap<_, _> = (0..5).map(|i| (i, i)).collect();
199+
map.range((Included(3), Included(2)));
200+
}
201+
202+
#[test]
203+
#[should_panic]
204+
fn test_range_backwards_2() {
205+
let map: BTreeMap<_, _> = (0..5).map(|i| (i, i)).collect();
206+
map.range((Included(3), Excluded(2)));
207+
}
208+
209+
#[test]
210+
#[should_panic]
211+
fn test_range_backwards_3() {
212+
let map: BTreeMap<_, _> = (0..5).map(|i| (i, i)).collect();
213+
map.range((Excluded(3), Included(2)));
214+
}
215+
216+
#[test]
217+
#[should_panic]
218+
fn test_range_backwards_4() {
219+
let map: BTreeMap<_, _> = (0..5).map(|i| (i, i)).collect();
220+
map.range((Excluded(3), Excluded(2)));
221+
}
222+
181223
#[test]
182224
fn test_range_1000() {
183225
let size = 1000;

0 commit comments

Comments
 (0)