Skip to content

Commit

Permalink
btree: simplify the backdoor between set and map
Browse files Browse the repository at this point in the history
The internal `btree::Recover` trait acted as a private API between
`BTreeSet` and `BTreeMap`, but we can use `pub(_)` restrictions these
days, and some of the methods don't need special handling anymore.

* `BTreeSet::get` can use `BTreeMap::get_key_value`
* `BTreeSet::take` can use `BTreeMap::remove_entry`
* `BTreeSet::replace` does need help, but this now uses a `pub(super)`
  method on `BTreeMap` instead of the trait.
* `btree::Recover` is now removed.
  • Loading branch information
cuviper committed Nov 13, 2024
1 parent 75609d6 commit d4e21f5
Show file tree
Hide file tree
Showing 4 changed files with 11 additions and 48 deletions.
40 changes: 6 additions & 34 deletions alloc/src/collections/btree/map.rs
Original file line number Diff line number Diff line change
Expand Up @@ -289,40 +289,12 @@ impl<K: Clone, V: Clone, A: Allocator + Clone> Clone for BTreeMap<K, V, A> {
}
}

impl<K, Q: ?Sized, A: Allocator + Clone> super::Recover<Q> for BTreeMap<K, SetValZST, A>
where
K: Borrow<Q> + Ord,
Q: Ord,
{
type Key = K;

fn get(&self, key: &Q) -> Option<&K> {
let root_node = self.root.as_ref()?.reborrow();
match root_node.search_tree(key) {
Found(handle) => Some(handle.into_kv().0),
GoDown(_) => None,
}
}

fn take(&mut self, key: &Q) -> Option<K> {
let (map, dormant_map) = DormantMutRef::new(self);
let root_node = map.root.as_mut()?.borrow_mut();
match root_node.search_tree(key) {
Found(handle) => Some(
OccupiedEntry {
handle,
dormant_map,
alloc: (*map.alloc).clone(),
_marker: PhantomData,
}
.remove_kv()
.0,
),
GoDown(_) => None,
}
}

fn replace(&mut self, key: K) -> Option<K> {
/// Internal functionality for `BTreeSet`.
impl<K, A: Allocator + Clone> BTreeMap<K, SetValZST, A> {
pub(super) fn replace(&mut self, key: K) -> Option<K>
where
K: Ord,
{
let (map, dormant_map) = DormantMutRef::new(self);
let root_node =
map.root.get_or_insert_with(|| Root::new((*map.alloc).clone())).borrow_mut();
Expand Down
8 changes: 0 additions & 8 deletions alloc/src/collections/btree/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,11 +12,3 @@ mod search;
pub mod set;
mod set_val;
mod split;

trait Recover<Q: ?Sized> {
type Key;

fn get(&self, key: &Q) -> Option<&Self::Key>;
fn take(&mut self, key: &Q) -> Option<Self::Key>;
fn replace(&mut self, key: Self::Key) -> Option<Self::Key>;
}
7 changes: 3 additions & 4 deletions alloc/src/collections/btree/set.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,6 @@ use core::iter::{FusedIterator, Peekable};
use core::mem::ManuallyDrop;
use core::ops::{BitAnd, BitOr, BitXor, Bound, RangeBounds, Sub};

use super::Recover;
use super::map::{BTreeMap, Keys};
use super::merge_iter::MergeIterInner;
use super::set_val::SetValZST;
Expand Down Expand Up @@ -635,7 +634,7 @@ impl<T, A: Allocator + Clone> BTreeSet<T, A> {
T: Borrow<Q> + Ord,
Q: Ord,
{
Recover::get(&self.map, value)
self.map.get_key_value(value).map(|(k, _)| k)
}

/// Returns `true` if `self` has no elements in common with `other`.
Expand Down Expand Up @@ -926,7 +925,7 @@ impl<T, A: Allocator + Clone> BTreeSet<T, A> {
where
T: Ord,
{
Recover::replace(&mut self.map, value)
self.map.replace(value)
}

/// If the set contains an element equal to the value, removes it from the
Expand Down Expand Up @@ -978,7 +977,7 @@ impl<T, A: Allocator + Clone> BTreeSet<T, A> {
T: Borrow<Q> + Ord,
Q: Ord,
{
Recover::take(&mut self.map, value)
self.map.remove_entry(value).map(|(k, _)| k)
}

/// Retains only the elements specified by the predicate.
Expand Down
4 changes: 2 additions & 2 deletions alloc/src/collections/btree/set_val.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,14 +3,14 @@
/// * `BTreeMap<T, ()>` (possible user-defined map)
/// * `BTreeMap<T, SetValZST>` (internal set representation)
#[derive(Debug, Eq, PartialEq, Ord, PartialOrd, Hash, Clone, Default)]
pub struct SetValZST;
pub(super) struct SetValZST;

/// A trait to differentiate between `BTreeMap` and `BTreeSet` values.
/// Returns `true` only for type `SetValZST`, `false` for all other types (blanket implementation).
/// `TypeId` requires a `'static` lifetime, use of this trait avoids that restriction.
///
/// [`TypeId`]: std::any::TypeId
pub trait IsSetVal {
pub(super) trait IsSetVal {
fn is_set_val() -> bool;
}

Expand Down

0 comments on commit d4e21f5

Please sign in to comment.