diff --git a/library/alloc/src/collections/btree/map/tests.rs b/library/alloc/src/collections/btree/map/tests.rs index e636e490e1bf4..3a74b6a6fa85c 100644 --- a/library/alloc/src/collections/btree/map/tests.rs +++ b/library/alloc/src/collections/btree/map/tests.rs @@ -776,7 +776,6 @@ fn test_range_backwards_4() { } #[test] -#[should_panic] fn test_range_finding_ill_order_in_map() { let mut map = BTreeMap::new(); map.insert(Cyclic3::B, ()); @@ -789,7 +788,6 @@ fn test_range_finding_ill_order_in_map() { } #[test] -#[should_panic] fn test_range_finding_ill_order_in_range_ord() { // Has proper order the first time asked, then flips around. struct EvilTwin(i32); diff --git a/library/alloc/src/collections/btree/navigate.rs b/library/alloc/src/collections/btree/navigate.rs index c2c99a9360cf1..4399feaccc9b7 100644 --- a/library/alloc/src/collections/btree/navigate.rs +++ b/library/alloc/src/collections/btree/navigate.rs @@ -29,11 +29,18 @@ impl LeafRange { impl NodeRef { /// Finds the distinct leaf edges delimiting a specified range in a tree. - /// Returns either a pair of different handles into the same tree or a pair - /// of empty options. + /// + /// If such distinct edges exist, returns them in ascending order, meaning + /// that a non-zero number of calls to `next_unchecked` on the `front` of + /// the result and/or calls to `next_back_unchecked` on the `back` of the + /// result will eventually reach the same edge. + /// + /// If there are no such edges, i.e., if the tree contains no key within + /// the range, returns a pair of empty options. + /// /// # Safety - /// Unless `BorrowType` is `Immut`, do not use the duplicate handles to - /// visit the same KV twice. + /// Unless `BorrowType` is `Immut`, do not use the handles to visit the same + /// KV twice. unsafe fn find_leaf_edges_spanning_range( self, range: R, diff --git a/library/alloc/src/collections/btree/search.rs b/library/alloc/src/collections/btree/search.rs index f376b3cde02e5..62a048e61c6be 100644 --- a/library/alloc/src/collections/btree/search.rs +++ b/library/alloc/src/collections/btree/search.rs @@ -68,13 +68,16 @@ impl NodeRef( mut self, @@ -113,10 +116,8 @@ impl NodeRef upper_edge_idx { - panic!("Ord is ill-defined in BTreeMap range") - } + let (upper_edge_idx, upper_child_bound) = + unsafe { self.find_upper_bound_index(upper_bound, lower_edge_idx) }; if lower_edge_idx < upper_edge_idx { return Ok(( self, @@ -126,6 +127,7 @@ impl NodeRef return Err(common_edge), @@ -165,7 +167,7 @@ impl NodeRef, { - let (edge_idx, bound) = self.find_upper_bound_index(bound); + let (edge_idx, bound) = unsafe { self.find_upper_bound_index(bound, 0) }; let edge = unsafe { Handle::new_edge(self, edge_idx) }; (edge, bound) } @@ -184,29 +186,33 @@ impl NodeRef { Q: Ord, K: Borrow, { - match self.find_key_index(key) { + match unsafe { self.find_key_index(key, 0) } { IndexResult::KV(idx) => Found(unsafe { Handle::new_kv(self, idx) }), IndexResult::Edge(idx) => GoDown(unsafe { Handle::new_edge(self, idx) }), } } /// Returns either the KV index in the node at which the key (or an equivalent) - /// exists, or the edge index where the key belongs. + /// exists, or the edge index where the key belongs, starting from a particular index. /// /// The result is meaningful only if the tree is ordered by key, like the tree /// in a `BTreeMap` is. - fn find_key_index(&self, key: &Q) -> IndexResult + /// + /// # Safety + /// `start_index` must be a valid edge index for the node. + unsafe fn find_key_index(&self, key: &Q, start_index: usize) -> IndexResult where Q: Ord, K: Borrow, { let node = self.reborrow(); let keys = node.keys(); - for (i, k) in keys.iter().enumerate() { + debug_assert!(start_index <= keys.len()); + for (offset, k) in unsafe { keys.get_unchecked(start_index..) }.iter().enumerate() { match key.cmp(k.borrow()) { Ordering::Greater => {} - Ordering::Equal => return IndexResult::KV(i), - Ordering::Less => return IndexResult::Edge(i), + Ordering::Equal => return IndexResult::KV(start_index + offset), + Ordering::Less => return IndexResult::Edge(start_index + offset), } } IndexResult::Edge(keys.len()) @@ -226,11 +232,11 @@ impl NodeRef { K: Borrow, { match bound { - Included(key) => match self.find_key_index(key) { + Included(key) => match unsafe { self.find_key_index(key, 0) } { IndexResult::KV(idx) => (idx, AllExcluded), IndexResult::Edge(idx) => (idx, bound), }, - Excluded(key) => match self.find_key_index(key) { + Excluded(key) => match unsafe { self.find_key_index(key, 0) } { IndexResult::KV(idx) => (idx + 1, AllIncluded), IndexResult::Edge(idx) => (idx, bound), }, @@ -239,26 +245,31 @@ impl NodeRef { } } - /// Clone of `find_lower_bound_index` for the upper bound. - fn find_upper_bound_index<'r, Q>( + /// Mirror image of `find_lower_bound_index` for the upper bound, + /// with an additional parameter to skip part of the key array. + /// + /// # Safety + /// `start_index` must be a valid edge index for the node. + unsafe fn find_upper_bound_index<'r, Q>( &self, bound: SearchBound<&'r Q>, + start_index: usize, ) -> (usize, SearchBound<&'r Q>) where Q: ?Sized + Ord, K: Borrow, { match bound { - Included(key) => match self.find_key_index(key) { + Included(key) => match unsafe { self.find_key_index(key, start_index) } { IndexResult::KV(idx) => (idx + 1, AllExcluded), IndexResult::Edge(idx) => (idx, bound), }, - Excluded(key) => match self.find_key_index(key) { + Excluded(key) => match unsafe { self.find_key_index(key, start_index) } { IndexResult::KV(idx) => (idx, AllIncluded), IndexResult::Edge(idx) => (idx, bound), }, AllIncluded => (self.len(), AllIncluded), - AllExcluded => (0, AllExcluded), + AllExcluded => (start_index, AllExcluded), } } }