From d4e21f55ec3085cf46af8304c33651256737f2fa Mon Sep 17 00:00:00 2001
From: Josh Stone <jistone@redhat.com>
Date: Wed, 13 Nov 2024 11:29:32 -0800
Subject: [PATCH] btree: simplify the backdoor between set and map

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.
---
 alloc/src/collections/btree/map.rs     | 40 ++++----------------------
 alloc/src/collections/btree/mod.rs     |  8 ------
 alloc/src/collections/btree/set.rs     |  7 ++---
 alloc/src/collections/btree/set_val.rs |  4 +--
 4 files changed, 11 insertions(+), 48 deletions(-)

diff --git a/alloc/src/collections/btree/map.rs b/alloc/src/collections/btree/map.rs
index 55649d865fc63..66353ccfa47d3 100644
--- a/alloc/src/collections/btree/map.rs
+++ b/alloc/src/collections/btree/map.rs
@@ -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();
diff --git a/alloc/src/collections/btree/mod.rs b/alloc/src/collections/btree/mod.rs
index c7d0144de30cb..b8667d09c33b3 100644
--- a/alloc/src/collections/btree/mod.rs
+++ b/alloc/src/collections/btree/mod.rs
@@ -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>;
-}
diff --git a/alloc/src/collections/btree/set.rs b/alloc/src/collections/btree/set.rs
index a40209fa2e397..8daee6030c270 100644
--- a/alloc/src/collections/btree/set.rs
+++ b/alloc/src/collections/btree/set.rs
@@ -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;
@@ -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`.
@@ -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
@@ -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.
diff --git a/alloc/src/collections/btree/set_val.rs b/alloc/src/collections/btree/set_val.rs
index 80c459bcf81db..cf30160bfbbc2 100644
--- a/alloc/src/collections/btree/set_val.rs
+++ b/alloc/src/collections/btree/set_val.rs
@@ -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;
 }