From deebb63cc8f706b7e92d5e67073a6a94f9a05b2d Mon Sep 17 00:00:00 2001 From: Stein Somers Date: Wed, 18 Nov 2020 18:19:38 +0100 Subject: [PATCH] BTree: split off reusable components from range_search --- .../alloc/src/collections/btree/navigate.rs | 142 +++++--------- library/alloc/src/collections/btree/node.rs | 10 - .../alloc/src/collections/btree/node/tests.rs | 10 +- library/alloc/src/collections/btree/search.rs | 183 +++++++++++++++++- library/alloc/src/lib.rs | 1 + 5 files changed, 228 insertions(+), 118 deletions(-) diff --git a/library/alloc/src/collections/btree/navigate.rs b/library/alloc/src/collections/btree/navigate.rs index e0b9a01d11056..c2c99a9360cf1 100644 --- a/library/alloc/src/collections/btree/navigate.rs +++ b/library/alloc/src/collections/btree/navigate.rs @@ -1,11 +1,8 @@ use core::borrow::Borrow; -use core::cmp::Ordering; -use core::ops::Bound::{Excluded, Included, Unbounded}; use core::ops::RangeBounds; use core::ptr; use super::node::{marker, ForceResult::*, Handle, NodeRef}; -use super::search::SearchResult; pub struct LeafRange { pub front: Option, marker::Edge>>, @@ -30,100 +27,50 @@ impl LeafRange { } } -/// Finds the leaf edges delimiting a specified range in or underneath a node. -/// -/// The result is meaningful only if the tree is ordered by key, like the tree -/// in a `BTreeMap` is. -fn range_search( - root1: NodeRef, - root2: NodeRef, - range: R, -) -> LeafRange -where - Q: ?Sized + Ord, - K: Borrow, - R: RangeBounds, -{ - // WARNING: Inlining these variables would be unsound (#81138) - // We assume the bounds reported by `range` remain the same, but - // an adversarial implementation could change between calls - let start = range.start_bound(); - let end = range.end_bound(); - match (start, end) { - (Excluded(s), Excluded(e)) if s == e => { - panic!("range start and end are equal and excluded in BTreeMap") - } - (Included(s) | Excluded(s), Included(e) | 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; - - loop { - // Using `range` again would be unsound (#81138) - let front = match (min_found, start) { - (false, Included(key)) => match min_node.search_node(key) { - SearchResult::Found(kv) => { - min_found = true; - kv.left_edge() - } - SearchResult::GoDown(edge) => edge, - }, - (false, Excluded(key)) => match min_node.search_node(key) { - SearchResult::Found(kv) => { - min_found = true; - kv.right_edge() - } - SearchResult::GoDown(edge) => edge, - }, - (true, Included(_)) => min_node.last_edge(), - (true, Excluded(_)) => min_node.first_edge(), - (_, Unbounded) => min_node.first_edge(), - }; - - // Using `range` again would be unsound (#81138) - let back = match (max_found, end) { - (false, Included(key)) => match max_node.search_node(key) { - SearchResult::Found(kv) => { - max_found = true; - kv.right_edge() - } - SearchResult::GoDown(edge) => edge, - }, - (false, Excluded(key)) => match max_node.search_node(key) { - SearchResult::Found(kv) => { - max_found = true; - kv.left_edge() +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. + /// # Safety + /// Unless `BorrowType` is `Immut`, do not use the duplicate handles to + /// visit the same KV twice. + unsafe fn find_leaf_edges_spanning_range( + self, + range: R, + ) -> LeafRange + where + Q: Ord, + K: Borrow, + R: RangeBounds, + { + match self.search_tree_for_bifurcation(&range) { + Err(_) => LeafRange::none(), + Ok(( + node, + lower_edge_idx, + upper_edge_idx, + mut lower_child_bound, + mut upper_child_bound, + )) => { + let mut lower_edge = unsafe { Handle::new_edge(ptr::read(&node), lower_edge_idx) }; + let mut upper_edge = unsafe { Handle::new_edge(node, upper_edge_idx) }; + loop { + match (lower_edge.force(), upper_edge.force()) { + (Leaf(f), Leaf(b)) => return LeafRange { front: Some(f), back: Some(b) }, + (Internal(f), Internal(b)) => { + (lower_edge, lower_child_bound) = + f.descend().find_lower_bound_edge(lower_child_bound); + (upper_edge, upper_child_bound) = + b.descend().find_upper_bound_edge(upper_child_bound); + } + _ => unreachable!("BTreeMap has different depths"), + } } - SearchResult::GoDown(edge) => edge, - }, - (true, Included(_)) => max_node.first_edge(), - (true, Excluded(_)) => max_node.last_edge(), - (_, Unbounded) => max_node.last_edge(), - }; - - if front.partial_cmp(&back) == Some(Ordering::Greater) { - panic!("Ord is ill-defined in BTreeMap range"); - } - match (front.force(), back.force()) { - (Leaf(f), Leaf(b)) => { - return LeafRange { front: Some(f), back: Some(b) }; } - (Internal(min_int), Internal(max_int)) => { - min_node = min_int.descend(); - max_node = max_int.descend(); - } - _ => unreachable!("BTreeMap has different depths"), - }; + } } } -/// Equivalent to `range_search(root1, root2, ..)` but without the `Ord` bound. /// Equivalent to `(root1.first_leaf_edge(), root2.last_leaf_edge())` but more efficient. fn full_range( root1: NodeRef, @@ -158,7 +105,8 @@ impl<'a, K: 'a, V: 'a> NodeRef, K, V, marker::LeafOrInternal> K: Borrow, R: RangeBounds, { - range_search(self, self, range) + // SAFETY: our borrow type is immutable. + unsafe { self.find_leaf_edges_spanning_range(range) } } /// Finds the pair of leaf edges delimiting an entire tree. @@ -174,16 +122,16 @@ impl<'a, K: 'a, V: 'a> NodeRef, K, V, marker::LeafOrInternal> /// /// The result is meaningful only if the tree is ordered by key, like the tree /// in a `BTreeMap` is. + /// + /// # Safety + /// Do not use the duplicate handles to visit the same KV twice. pub fn range_search(self, range: R) -> LeafRange, K, V> where Q: ?Sized + Ord, K: Borrow, R: RangeBounds, { - // We duplicate the root NodeRef here -- we will never visit the same KV - // twice, and never end up with overlapping value references. - let self2 = unsafe { ptr::read(&self) }; - range_search(self, self2, range) + unsafe { self.find_leaf_edges_spanning_range(range) } } /// Splits a unique reference into a pair of leaf edges delimiting the full range of the tree. diff --git a/library/alloc/src/collections/btree/node.rs b/library/alloc/src/collections/btree/node.rs index 4fc32305f1e30..4622fd6cc1c8f 100644 --- a/library/alloc/src/collections/btree/node.rs +++ b/library/alloc/src/collections/btree/node.rs @@ -31,7 +31,6 @@ // since leaf edges are empty and need no data representation. In an internal node, // an edge both identifies a position and contains a pointer to a child node. -use core::cmp::Ordering; use core::marker::PhantomData; use core::mem::{self, MaybeUninit}; use core::ptr::{self, NonNull}; @@ -742,15 +741,6 @@ impl PartialEq } } -impl PartialOrd - for Handle, HandleType> -{ - fn partial_cmp(&self, other: &Self) -> Option { - let Self { node, idx, _marker } = self; - if node.eq(&other.node) { Some(idx.cmp(&other.idx)) } else { None } - } -} - impl Handle, HandleType> { diff --git a/library/alloc/src/collections/btree/node/tests.rs b/library/alloc/src/collections/btree/node/tests.rs index acb7210ca7c8b..58e78d99e0a9e 100644 --- a/library/alloc/src/collections/btree/node/tests.rs +++ b/library/alloc/src/collections/btree/node/tests.rs @@ -2,7 +2,6 @@ use super::super::navigate; use super::*; use crate::fmt::Debug; use crate::string::String; -use core::cmp::Ordering::*; impl<'a, K: 'a, V: 'a> NodeRef, K, V, marker::LeafOrInternal> { // Asserts that the back pointer in each reachable node points to its parent. @@ -67,7 +66,7 @@ fn test_splitpoint() { } #[test] -fn test_partial_cmp_eq() { +fn test_partial_eq() { let mut root1 = NodeRef::new_leaf(); root1.borrow_mut().push(1, ()); let mut root1 = NodeRef::new_internal(root1.forget_type()).forget_type(); @@ -87,13 +86,6 @@ fn test_partial_cmp_eq() { assert!(top_edge_1 == top_edge_1); assert!(top_edge_1 != top_edge_2); - assert_eq!(leaf_edge_1a.partial_cmp(&leaf_edge_1a), Some(Equal)); - assert_eq!(leaf_edge_1a.partial_cmp(&leaf_edge_1b), Some(Less)); - assert_eq!(leaf_edge_1a.partial_cmp(&top_edge_1), None); - assert_eq!(leaf_edge_1a.partial_cmp(&top_edge_2), None); - assert_eq!(top_edge_1.partial_cmp(&top_edge_1), Some(Equal)); - assert_eq!(top_edge_1.partial_cmp(&top_edge_2), None); - root1.pop_internal_level(); unsafe { root1.into_dying().deallocate_and_ascend() }; unsafe { root2.into_dying().deallocate_and_ascend() }; diff --git a/library/alloc/src/collections/btree/search.rs b/library/alloc/src/collections/btree/search.rs index f87444b7cd3e5..f376b3cde02e5 100644 --- a/library/alloc/src/collections/btree/search.rs +++ b/library/alloc/src/collections/btree/search.rs @@ -1,10 +1,33 @@ use core::borrow::Borrow; use core::cmp::Ordering; +use core::ops::{Bound, RangeBounds}; use super::node::{marker, ForceResult::*, Handle, NodeRef}; +use SearchBound::*; use SearchResult::*; +pub enum SearchBound { + /// An inclusive bound to look for, just like `Bound::Included(T)`. + Included(T), + /// An exclusive bound to look for, just like `Bound::Excluded(T)`. + Excluded(T), + /// An unconditional inclusive bound, just like `Bound::Unbounded`. + AllIncluded, + /// An unconditional exclusive bound. + AllExcluded, +} + +impl SearchBound { + pub fn from_range(range_bound: Bound) -> Self { + match range_bound { + Bound::Included(t) => Included(t), + Bound::Excluded(t) => Excluded(t), + Bound::Unbounded => AllIncluded, + } + } +} + pub enum SearchResult { Found(Handle, marker::KV>), GoDown(Handle, marker::Edge>), @@ -40,6 +63,112 @@ impl NodeRef( + mut self, + range: &'r R, + ) -> Result< + ( + NodeRef, + usize, + usize, + SearchBound<&'r Q>, + SearchBound<&'r Q>, + ), + Handle, marker::Edge>, + > + where + Q: Ord, + K: Borrow, + R: RangeBounds, + { + // WARNING: Inlining these variables would be unsound (#81138) + // We assume the bounds reported by `range` remain the same, but + // an adversarial implementation could change between calls + let (start, end) = (range.start_bound(), range.end_bound()); + match (start, end) { + (Bound::Excluded(s), Bound::Excluded(e)) if s == e => { + panic!("range start and end are equal and excluded in BTreeMap") + } + (Bound::Included(s) | Bound::Excluded(s), Bound::Included(e) | Bound::Excluded(e)) + if s > e => + { + panic!("range start is greater than range end in BTreeMap") + } + _ => {} + } + let mut lower_bound = SearchBound::from_range(start); + let mut upper_bound = SearchBound::from_range(end); + loop { + let (lower_edge_idx, lower_child_bound) = self.find_lower_bound_index(lower_bound); + let (upper_edge_idx, upper_child_bound) = self.find_upper_bound_index(upper_bound); + if lower_edge_idx > upper_edge_idx { + panic!("Ord is ill-defined in BTreeMap range") + } + if lower_edge_idx < upper_edge_idx { + return Ok(( + self, + lower_edge_idx, + upper_edge_idx, + lower_child_bound, + upper_child_bound, + )); + } + let common_edge = unsafe { Handle::new_edge(self, lower_edge_idx) }; + match common_edge.force() { + Leaf(common_edge) => return Err(common_edge), + Internal(common_edge) => { + self = common_edge.descend(); + lower_bound = lower_child_bound; + upper_bound = upper_child_bound; + } + } + } + } + + /// Finds an edge in the node delimiting the lower bound of a range. + /// Also returns the lower bound to be used for continuing the search in + /// the matching child node, if `self` is an internal node. + /// + /// The result is meaningful only if the tree is ordered by key. + pub fn find_lower_bound_edge<'r, Q>( + self, + bound: SearchBound<&'r Q>, + ) -> (Handle, SearchBound<&'r Q>) + where + Q: ?Sized + Ord, + K: Borrow, + { + let (edge_idx, bound) = self.find_lower_bound_index(bound); + let edge = unsafe { Handle::new_edge(self, edge_idx) }; + (edge, bound) + } + + /// Clone of `find_lower_bound_edge` for the upper bound. + pub fn find_upper_bound_edge<'r, Q>( + self, + bound: SearchBound<&'r Q>, + ) -> (Handle, SearchBound<&'r Q>) + where + Q: ?Sized + Ord, + K: Borrow, + { + let (edge_idx, bound) = self.find_upper_bound_index(bound); + let edge = unsafe { Handle::new_edge(self, edge_idx) }; + (edge, bound) + } } impl NodeRef { @@ -55,7 +184,7 @@ impl NodeRef { Q: Ord, K: Borrow, { - match self.find_index(key) { + match self.find_key_index(key) { IndexResult::KV(idx) => Found(unsafe { Handle::new_kv(self, idx) }), IndexResult::Edge(idx) => GoDown(unsafe { Handle::new_edge(self, idx) }), } @@ -66,7 +195,7 @@ impl NodeRef { /// /// The result is meaningful only if the tree is ordered by key, like the tree /// in a `BTreeMap` is. - fn find_index(&self, key: &Q) -> IndexResult + fn find_key_index(&self, key: &Q) -> IndexResult where Q: Ord, K: Borrow, @@ -82,4 +211,54 @@ impl NodeRef { } IndexResult::Edge(keys.len()) } + + /// Finds an edge index in the node delimiting the lower bound of a range. + /// Also returns the lower bound to be used for continuing the search in + /// the matching child node, if `self` is an internal node. + /// + /// The result is meaningful only if the tree is ordered by key. + fn find_lower_bound_index<'r, Q>( + &self, + bound: SearchBound<&'r Q>, + ) -> (usize, SearchBound<&'r Q>) + where + Q: ?Sized + Ord, + K: Borrow, + { + match bound { + Included(key) => match self.find_key_index(key) { + IndexResult::KV(idx) => (idx, AllExcluded), + IndexResult::Edge(idx) => (idx, bound), + }, + Excluded(key) => match self.find_key_index(key) { + IndexResult::KV(idx) => (idx + 1, AllIncluded), + IndexResult::Edge(idx) => (idx, bound), + }, + AllIncluded => (0, AllIncluded), + AllExcluded => (self.len(), AllExcluded), + } + } + + /// Clone of `find_lower_bound_index` for the upper bound. + fn find_upper_bound_index<'r, Q>( + &self, + bound: SearchBound<&'r Q>, + ) -> (usize, SearchBound<&'r Q>) + where + Q: ?Sized + Ord, + K: Borrow, + { + match bound { + Included(key) => match self.find_key_index(key) { + IndexResult::KV(idx) => (idx + 1, AllExcluded), + IndexResult::Edge(idx) => (idx, bound), + }, + Excluded(key) => match self.find_key_index(key) { + IndexResult::KV(idx) => (idx, AllIncluded), + IndexResult::Edge(idx) => (idx, bound), + }, + AllIncluded => (self.len(), AllIncluded), + AllExcluded => (0, AllExcluded), + } + } } diff --git a/library/alloc/src/lib.rs b/library/alloc/src/lib.rs index c020a969f1fb9..c4f05a4d09c51 100644 --- a/library/alloc/src/lib.rs +++ b/library/alloc/src/lib.rs @@ -92,6 +92,7 @@ #![feature(const_fn)] #![feature(cow_is_borrowed)] #![feature(const_cow_is_borrowed)] +#![feature(destructuring_assignment)] #![feature(dispatch_from_dyn)] #![feature(core_intrinsics)] #![feature(dropck_eyepatch)]