From de6e53a3274acf4a15e12be4c3a7b0108bd567a8 Mon Sep 17 00:00:00 2001 From: Stein Somers Date: Mon, 30 Nov 2020 12:15:25 +0100 Subject: [PATCH] BTreeMap: convert search functions to methods --- library/alloc/src/collections/btree/map.rs | 18 +-- .../alloc/src/collections/btree/navigate.rs | 10 +- library/alloc/src/collections/btree/search.rs | 134 +++++++++--------- library/alloc/src/collections/btree/split.rs | 4 +- 4 files changed, 83 insertions(+), 83 deletions(-) diff --git a/library/alloc/src/collections/btree/map.rs b/library/alloc/src/collections/btree/map.rs index 5e63a303d22c3..ecc2873187467 100644 --- a/library/alloc/src/collections/btree/map.rs +++ b/library/alloc/src/collections/btree/map.rs @@ -10,7 +10,7 @@ use core::ptr; use super::borrow::DormantMutRef; use super::node::{self, marker, ForceResult::*, Handle, NodeRef, Root}; -use super::search::{self, SearchResult::*}; +use super::search::SearchResult::*; use super::unwrap_unchecked; mod entry; @@ -230,7 +230,7 @@ where fn get(&self, key: &Q) -> Option<&K> { let root_node = self.root.as_ref()?.reborrow(); - match search::search_tree(root_node, key) { + match root_node.search_tree(key) { Found(handle) => Some(handle.into_kv().0), GoDown(_) => None, } @@ -239,7 +239,7 @@ where fn take(&mut self, key: &Q) -> Option { let (map, dormant_map) = DormantMutRef::new(self); let root_node = map.root.as_mut()?.borrow_mut(); - match search::search_tree(root_node, key) { + match root_node.search_tree(key) { Found(handle) => { Some(OccupiedEntry { handle, dormant_map, _marker: PhantomData }.remove_kv().0) } @@ -250,7 +250,7 @@ where fn replace(&mut self, key: K) -> Option { let (map, dormant_map) = DormantMutRef::new(self); let root_node = Self::ensure_is_owned(&mut map.root).borrow_mut(); - match search::search_tree::, K, (), K>(root_node, &key) { + match root_node.search_tree::(&key) { Found(mut kv) => Some(mem::replace(kv.key_mut(), key)), GoDown(handle) => { VacantEntry { key, handle, dormant_map, _marker: PhantomData }.insert(()); @@ -526,7 +526,7 @@ impl BTreeMap { Q: Ord, { let root_node = self.root.as_ref()?.reborrow(); - match search::search_tree(root_node, key) { + match root_node.search_tree(key) { Found(handle) => Some(handle.into_kv().1), GoDown(_) => None, } @@ -554,7 +554,7 @@ impl BTreeMap { Q: Ord, { let root_node = self.root.as_ref()?.reborrow(); - match search::search_tree(root_node, k) { + match root_node.search_tree(k) { Found(handle) => Some(handle.into_kv()), GoDown(_) => None, } @@ -762,7 +762,7 @@ impl BTreeMap { Q: Ord, { let root_node = self.root.as_mut()?.borrow_mut(); - match search::search_tree(root_node, key) { + match root_node.search_tree(key) { Found(handle) => Some(handle.into_val_mut()), GoDown(_) => None, } @@ -858,7 +858,7 @@ impl BTreeMap { { let (map, dormant_map) = DormantMutRef::new(self); let root_node = map.root.as_mut()?.borrow_mut(); - match search::search_tree(root_node, key) { + match root_node.search_tree(key) { Found(handle) => { Some(OccupiedEntry { handle, dormant_map, _marker: PhantomData }.remove_entry()) } @@ -1051,7 +1051,7 @@ impl BTreeMap { // FIXME(@porglezomp) Avoid allocating if we don't insert let (map, dormant_map) = DormantMutRef::new(self); let root_node = Self::ensure_is_owned(&mut map.root).borrow_mut(); - match search::search_tree(root_node, &key) { + match root_node.search_tree(&key) { Found(handle) => Occupied(OccupiedEntry { handle, dormant_map, _marker: PhantomData }), GoDown(handle) => { Vacant(VacantEntry { key, handle, dormant_map, _marker: PhantomData }) diff --git a/library/alloc/src/collections/btree/navigate.rs b/library/alloc/src/collections/btree/navigate.rs index cce8b21a2bc41..45a8e0103a496 100644 --- a/library/alloc/src/collections/btree/navigate.rs +++ b/library/alloc/src/collections/btree/navigate.rs @@ -5,7 +5,7 @@ use core::ops::RangeBounds; use core::ptr; use super::node::{marker, ForceResult::*, Handle, NodeRef}; -use super::search::{self, SearchResult}; +use super::search::SearchResult; use super::unwrap_unchecked; /// Finds the leaf edges delimiting a specified range in or underneath a node. @@ -42,14 +42,14 @@ where loop { let front = match (min_found, range.start_bound()) { - (false, Included(key)) => match search::search_node(min_node, key) { + (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 search::search_node(min_node, key) { + (false, Excluded(key)) => match min_node.search_node(key) { SearchResult::Found(kv) => { min_found = true; kv.right_edge() @@ -62,14 +62,14 @@ where }; let back = match (max_found, range.end_bound()) { - (false, Included(key)) => match search::search_node(max_node, key) { + (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 search::search_node(max_node, key) { + (false, Excluded(key)) => match max_node.search_node(key) { SearchResult::Found(kv) => { max_found = true; kv.left_edge() diff --git a/library/alloc/src/collections/btree/search.rs b/library/alloc/src/collections/btree/search.rs index efe94ef175ce7..f62eae3f4d5a2 100644 --- a/library/alloc/src/collections/btree/search.rs +++ b/library/alloc/src/collections/btree/search.rs @@ -10,79 +10,79 @@ pub enum SearchResult { GoDown(Handle, marker::Edge>), } -/// Looks up a given key in a (sub)tree headed by the given node, recursively. -/// Returns a `Found` with the handle of the matching KV, if any. Otherwise, -/// returns a `GoDown` with the handle of the leaf edge where the key belongs. -/// -/// The result is meaningful only if the tree is ordered by key, like the tree -/// in a `BTreeMap` is. -pub fn search_tree( - mut node: NodeRef, - key: &Q, -) -> SearchResult -where - Q: Ord, - K: Borrow, -{ - loop { - match search_node(node, key) { - Found(handle) => return Found(handle), - GoDown(handle) => match handle.force() { - Leaf(leaf) => return GoDown(leaf), - Internal(internal) => { - node = internal.descend(); - continue; - } - }, +pub enum IndexResult { + KV(usize), + Edge(usize), +} + +impl NodeRef { + /// Looks up a given key in a (sub)tree headed by the node, recursively. + /// Returns a `Found` with the handle of the matching KV, if any. Otherwise, + /// returns a `GoDown` with the handle of the leaf edge where the key belongs. + /// + /// The result is meaningful only if the tree is ordered by key, like the tree + /// in a `BTreeMap` is. + pub fn search_tree( + mut self, + key: &Q, + ) -> SearchResult + where + Q: Ord, + K: Borrow, + { + loop { + self = match self.search_node(key) { + Found(handle) => return Found(handle), + GoDown(handle) => match handle.force() { + Leaf(leaf) => return GoDown(leaf), + Internal(internal) => internal.descend(), + }, + } } } } -/// Looks up a given key in a given node, without recursion. -/// Returns a `Found` with the handle of the matching KV, if any. Otherwise, -/// returns a `GoDown` with the handle of the edge where the key might be found -/// (if the node is internal) or where the key can be inserted. -/// -/// The result is meaningful only if the tree is ordered by key, like the tree -/// in a `BTreeMap` is. -pub fn search_node( - node: NodeRef, - key: &Q, -) -> SearchResult -where - Q: Ord, - K: Borrow, -{ - match search_linear(&node, key) { - (idx, true) => Found(unsafe { Handle::new_kv(node, idx) }), - (idx, false) => GoDown(unsafe { Handle::new_edge(node, idx) }), +impl NodeRef { + /// Looks up a given key in the node, without recursion. + /// Returns a `Found` with the handle of the matching KV, if any. Otherwise, + /// returns a `GoDown` with the handle of the edge where the key might be found + /// (if the node is internal) or where the key can be inserted. + /// + /// The result is meaningful only if the tree is ordered by key, like the tree + /// in a `BTreeMap` is. + pub fn search_node(self, key: &Q) -> SearchResult + where + Q: Ord, + K: Borrow, + { + match self.find_index(key) { + 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 and `true`, or the edge index where the key belongs and `false`. -/// -/// The result is meaningful only if the tree is ordered by key, like the tree -/// in a `BTreeMap` is. -fn search_linear( - node: &NodeRef, - key: &Q, -) -> (usize, bool) -where - Q: Ord, - K: Borrow, -{ - // This function is defined over all borrow types (immutable, mutable, owned). - // Using `keys_at()` is fine here even if BorrowType is mutable, as all we return - // is an index -- not a reference. - let len = node.len(); - for i in 0..len { - let k = unsafe { node.reborrow().key_at(i) }; - match key.cmp(k.borrow()) { - Ordering::Greater => {} - Ordering::Equal => return (i, true), - Ordering::Less => return (i, false), + /// Returns either the KV index in the node at which the key (or an equivalent) + /// exists, or the edge index where the key belongs. + /// + /// 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 + where + Q: Ord, + K: Borrow, + { + // This function is defined over all borrow types (immutable, mutable, owned). + // Using `keys_at()` is fine here even if BorrowType is mutable, as all we return + // is an index -- not a reference. + let len = self.len(); + for i in 0..len { + let k = unsafe { self.reborrow().key_at(i) }; + match key.cmp(k.borrow()) { + Ordering::Greater => {} + Ordering::Equal => return IndexResult::KV(i), + Ordering::Less => return IndexResult::Edge(i), + } } + IndexResult::Edge(len) } - (len, false) } diff --git a/library/alloc/src/collections/btree/split.rs b/library/alloc/src/collections/btree/split.rs index 375c6173a093e..62c5e3a46d74d 100644 --- a/library/alloc/src/collections/btree/split.rs +++ b/library/alloc/src/collections/btree/split.rs @@ -1,6 +1,6 @@ use super::map::MIN_LEN; use super::node::{ForceResult::*, Root}; -use super::search::{search_node, SearchResult::*}; +use super::search::SearchResult::*; use core::borrow::Borrow; impl Root { @@ -21,7 +21,7 @@ impl Root { let mut right_node = right_root.borrow_mut(); loop { - let mut split_edge = match search_node(left_node, key) { + let mut split_edge = match left_node.search_node(key) { // key is going to the right tree Found(kv) => kv.left_edge(), GoDown(edge) => edge,