Changed with_next() to use remove_next() + insert_after_unchecked() instead of drop handling, refactored code a bit, and updated comments

Author: asder8215

library/alloc/src/collections/btree/map.rs

@@ -1318,7 +1318,7 @@ impl<K, V, A: Allocator + Clone> BTreeMap<K, V, A> {
                     });
                 }
                 Ordering::Less =>
-                // Safety: we know our other_key's ordering is less than self_key,
+                // SAFETY: we know our other_key's ordering is less than self_key,
                 // so inserting before will guarantee sorted order
                 unsafe {
                     self_cursor.insert_before_unchecked(first_other_key, first_other_val);
@@ -1328,61 +1328,45 @@ impl<K, V, A: Allocator + Clone> BTreeMap<K, V, A> {
                 }
             }
         } else {
-            // Safety: if we reach here, that means our cursor has reached
-            // the end of self BTreeMap, (other_key is greater than all the
-            // previous self BTreeMap keys) so we just insert other_key here
-            // at the end of the CursorMut
+            // SAFETY: reaching here means our cursor is at the end
+            // self BTreeMap so we just insert other_key here
             unsafe {
                 self_cursor.insert_after_unchecked(first_other_key, first_other_val);
             }
         }
 
         for (other_key, other_val) in other_iter {
-            if self_cursor.peek_next().is_some() {
-                loop {
-                    let self_entry = self_cursor.peek_next();
-                    if let Some((self_key, _)) = self_entry {
-                        match K::cmp(&other_key, self_key) {
-                            Ordering::Equal => {
-                                self_cursor.with_next(|self_key, self_val| {
-                                    conflict(self_key, self_val, other_val)
-                                });
-                                break;
-                            }
-                            Ordering::Less => {
-                                // Safety: we know our other_key's ordering is less than self_key,
-                                // so inserting before will guarantee sorted order
-                                unsafe {
-                                    self_cursor.insert_before_unchecked(other_key, other_val);
-                                }
-                                break;
-                            }
-                            Ordering::Greater => {
-                                // advance the cursor forward
-                                self_cursor.next();
+            loop {
+                if let Some((self_key, _)) = self_cursor.peek_next() {
+                    match K::cmp(&other_key, self_key) {
+                        Ordering::Equal => {
+                            self_cursor.with_next(|self_key, self_val| {
+                                conflict(self_key, self_val, other_val)
+                            });
+                            break;
+                        }
+                        Ordering::Less => {
+                            // SAFETY: we know our other_key's ordering is less than self_key,
+                            // so inserting before will guarantee sorted order
+                            unsafe {
+                                self_cursor.insert_before_unchecked(other_key, other_val);
                             }
+                            break;
                         }
-                    } else {
-                        // Safety: if we reach here, that means our cursor has reached
-                        // the end of self BTreeMap, (other_key is greater than all the
-                        // previous self BTreeMap keys) so we just insert other_key here
-                        // at the end of the Cursor
-                        unsafe {
-                            self_cursor.insert_after_unchecked(other_key, other_val);
+                        Ordering::Greater => {
+                            // advance the cursor forward
+                            self_cursor.next();
                         }
-                        self_cursor.next();
-                        break;
                     }
+                } else {
+                    // SAFETY: reaching here means our cursor is at the end
+                    // self BTreeMap so we just insert other_key here
+                    unsafe {
+                        self_cursor.insert_after_unchecked(other_key, other_val);
+                    }
+                    self_cursor.next();
+                    break;
                 }
-            } else {
-                // Safety: if we reach here, that means our cursor has reached
-                // the end of self BTreeMap, (other_key is greater than all the
-                // previous self BTreeMap keys) so we just insert the rest of
-                // other_keys here at the end of CursorMut
-                unsafe {
-                    self_cursor.insert_after_unchecked(other_key, other_val);
-                }
-                self_cursor.next();
             }
         }
     }
@@ -3416,91 +3400,27 @@ impl<'a, K: Ord, V, A: Allocator + Clone> CursorMutKey<'a, K, V, A> {
     ///
     /// If the cursor is at the end of the map then the function is not called
     /// and this essentially does not do anything.
+    ///
+    /// # Safety
+    ///
+    /// You must ensure that the `BTreeMap` invariants are maintained.
+    /// Specifically:
+    ///
+    /// * The next element's key must be unique in the tree.
+    /// * All keys in the tree must remain in sorted order.
     #[allow(dead_code)] /* This function exists for consistency with CursorMut */
     pub(super) fn with_next(&mut self, f: impl FnOnce(K, V) -> (K, V)) {
-        /// Call `f` with the next entry's key and value, replacing the
-        /// next entry's key and value with the returned key and value.
-        /// if `f` unwinds, the next entry is removed.
-        /// Equivalent to a more efficient version of:
-        /// ```ignore (this is an example equivalent)
-        /// if let Some((k, v)) = self.remove_next() {
-        ///     let (k, v) = f(k, v);
-        ///     // Safety: key is unmodified
-        ///     unsafe { self.insert_after_unchecked(k, v) };
-        /// }
-        /// ```
-        struct RemoveNextOnDrop<'a, 'b, K, V, A = Global>
-        where
-            K: Ord,
-            A: Allocator + Clone,
-        {
-            cursor: &'a mut CursorMutKey<'b, K, V, A>,
-            forget_next: bool,
-        }
-
-        impl<K: Ord, V, A: Allocator + Clone> Drop for RemoveNextOnDrop<'_, '_, K, V, A> {
-            fn drop(&mut self) {
-                if self.forget_next {
-                    // Removes the next entry of the Cursor without running
-                    // running the destructor on the key and value (prevents double
-                    // free)
-                    self.cursor.forget_next_key_and_value();
-                }
-            }
-        }
-
-        let mut remove_next_on_drop = RemoveNextOnDrop {
-            cursor: self,
-            forget_next: false, // next value is not known yet
-        };
-
-        if let Some((k_mut, v_mut)) = remove_next_on_drop.cursor.peek_next() {
-            remove_next_on_drop.forget_next = true;
-            // Safety: we move the K, V out of the next entry,
-            // we marked the entry's value to be forgotten
-            // when remove_next_on_drop is dropped that
-            // way we avoid returning to the caller leaving
-            // a moved-out invalid value if `f` unwinds.
-            let k = unsafe { ptr::read(k_mut) };
-            let v = unsafe { ptr::read(v_mut) };
+        // if `f` unwinds, the next entry is already removed leaving
+        // the tree in valid state.
+        // FIXME: Once `MaybeDangling` is implemented, we can optimize
+        // this through using a drop handler and transmutating CursorMutKey<K, V>
+        // to CursorMutKey<ManuallyDrop<K>, ManuallyDrop<V>> (see PR #152418)
+        if let Some((k, v)) = self.remove_next() {
+            // SAFETY: we remove the K, V out of the next entry,
+            // apply 'f' to get a new (K, V), and insert it back
+            // into the next entry that the cursor is pointing at
             let (k, v) = f(k, v);
-            unsafe { ptr::write(k_mut, k) };
-            unsafe { ptr::write(v_mut, v) };
-
-            remove_next_on_drop.forget_next = false;
-        }
-    }
-
-    /// Forgets the next element's key and value from the `BTreeMap`.
-    ///
-    /// The element's key and value does not run its destructor. The cursor position is
-    /// unchanged (before the forgotten element).
-    fn forget_next_key_and_value(&mut self) {
-        let current = match self.current.take() {
-            Some(current) => current,
-            None => return,
-        };
-
-        if current.reborrow().next_kv().is_err() {
-            self.current = Some(current);
-            return;
-        }
-
-        let mut emptied_internal_root = false;
-        if let Ok(next_kv) = current.next_kv() {
-            let ((key, val), pos) =
-                next_kv.remove_kv_tracking(|| emptied_internal_root = true, self.alloc.clone());
-            // Don't run destructor on next element's key and value
-            mem::forget(key);
-            mem::forget(val);
-            self.current = Some(pos);
-            *self.length -= 1;
-            if emptied_internal_root {
-                // SAFETY: This is safe since current does not point within the now
-                // empty root node.
-                let root = unsafe { self.root.reborrow().as_mut().unwrap() };
-                root.pop_internal_level(self.alloc.clone());
-            }
+            unsafe { self.insert_after_unchecked(k, v) };
         }
     }
 
@@ -3717,85 +3637,17 @@ impl<'a, K: Ord, V, A: Allocator + Clone> CursorMut<'a, K, V, A> {
     /// If the cursor is at the end of the map then the function is not called
     /// and this essentially does not do anything.
     pub(super) fn with_next(&mut self, f: impl FnOnce(&K, V) -> V) {
-        /// Call `f` with the next entry's key and value, replacing the
-        /// next entry's value with the returned value. if `f` unwinds,
-        /// the next entry is removed.
-        /// Equivalent to a more efficient version of:
-        /// ```ignore (this is an example equivalent)
-        /// if let Some((k, v)) = self.remove_next() {
-        ///     let v = f(&k, v);
-        ///     // Safety: key is unmodified
-        ///     unsafe { self.insert_after_unchecked(k, v) };
-        /// }
-        /// ```
-        struct RemoveNextOnDrop<'a, 'b, K, V, A = Global>
-        where
-            K: Ord,
-            A: Allocator + Clone,
-        {
-            cursor: &'a mut CursorMut<'b, K, V, A>,
-            forget_next: bool,
-        }
-
-        impl<K: Ord, V, A: Allocator + Clone> Drop for RemoveNextOnDrop<'_, '_, K, V, A> {
-            fn drop(&mut self) {
-                if self.forget_next {
-                    // Removes the next entry of the Cursor without running
-                    // running the destructor on the value (prevents double
-                    // free)
-                    self.cursor.forget_next_value();
-                }
-            }
-        }
-
-        let mut remove_next_on_drop = RemoveNextOnDrop {
-            cursor: self,
-            forget_next: false, // next value is not known yet
-        };
-
-        if let Some((k, v_mut)) = remove_next_on_drop.cursor.peek_next() {
-            remove_next_on_drop.forget_next = true;
-            // Safety: we move the V out of the next entry,
-            // we marked the entry's value to be forgotten
-            // when remove_next_on_drop is dropped that
-            // way we avoid returning to the caller leaving
-            // a moved-out invalid value if `f` unwinds.
-            let v = unsafe { ptr::read(v_mut) };
-            let v = f(k, v);
-            unsafe { ptr::write(v_mut, v) };
-            remove_next_on_drop.forget_next = false;
-        }
-    }
-
-    /// Forgets the next element's value from the `BTreeMap`.
-    ///
-    /// The element's value does not run its destructor. The cursor position is
-    /// unchanged (before the forgotten element).
-    fn forget_next_value(&mut self) {
-        let current = match self.inner.current.take() {
-            Some(current) => current,
-            None => return,
-        };
-
-        if current.reborrow().next_kv().is_err() {
-            self.inner.current = Some(current);
-            return;
-        }
-
-        let mut emptied_internal_root = false;
-        if let Ok(next_kv) = current.next_kv() {
-            let ((_, val), pos) = next_kv
-                .remove_kv_tracking(|| emptied_internal_root = true, self.inner.alloc.clone());
-            // Don't run destructor on next element's value
-            mem::forget(val);
-            self.inner.current = Some(pos);
-            *self.inner.length -= 1;
-            if emptied_internal_root {
-                // SAFETY: This is safe since current does not point within the now
-                // empty root node.
-                let root = unsafe { self.inner.root.reborrow().as_mut().unwrap() };
-                root.pop_internal_level(self.inner.alloc.clone());
-            }
+        // if `f` unwinds, the next entry is already removed leaving
+        // the tree in valid state.
+        // FIXME: Once `MaybeDangling` is implemented, we can optimize
+        // this through using a drop handler and transmutating CursorMut<K, V>
+        // to CursorMut<K, ManuallyDrop<V>> (see PR #152418)
+        if let Some((k, v)) = self.remove_next() {
+            // SAFETY: we remove the K, V out of the next entry,
+            // apply 'f' to get a new (K, V), and insert it back
+            // into the next entry that the cursor is pointing at
+            let v = f(&k, v);
+            unsafe { self.insert_after_unchecked(k, v) };
         }
     }