Auto merge of #39457 - bvinc:master, r=alexcrichton

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?

3 files changed, 134 insertions, 115 deletions

diff --git a/src/libcollections/btree/map.rs b/src/libcollections/btree/map.rs
index e1fabe2cc49..7218d15ded5 100644
--- a/src/libcollections/btree/map.rs
+++ b/src/libcollections/btree/map.rs
@@ -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:
@@ -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.
@@ -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:
@@ -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,
         }
     }
@@ -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(|| {
diff --git a/src/libcollections/btree/search.rs b/src/libcollections/btree/search.rs
index c94b570bfed..bc1272fbc78 100644
--- a/src/libcollections/btree/search.rs
+++ b/src/libcollections/btree/search.rs
@@ -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)
@@ -73,4 +73,3 @@ fn search_linear<BorrowType, K, V, Type, Q: ?Sized>(
     }
     (node.keys().len(), false)
 }
-
diff --git a/src/libcollectionstest/btree/map.rs b/src/libcollectionstest/btree/map.rs
index 11be13426e4..f33923f9963 100644
--- a/src/libcollectionstest/btree/map.rs
+++ b/src/libcollectionstest/btree/map.rs
@@ -179,6 +179,48 @@ fn test_range_small() {
 }
 
 #[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;
     let map: BTreeMap<_, _> = (0..size).map(|i| (i, i)).collect();