with_next() and forget_next_*() created for Cursor* to handle panic behavior for conflict func

Author: asder8215

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

@@ -1306,29 +1306,25 @@ impl<K, V, A: Allocator + Clone> BTreeMap<K, V, A> {
         let mut other_iter = other.into_iter();
         let (first_other_key, first_other_val) = other_iter.next().unwrap();
 
-        // find the first gap that has the smallest key greater than the first key from other
+        // find the first gap that has the smallest key greater than or equal to
+        // the first key from other
         let mut self_cursor = self.lower_bound_mut(Bound::Included(&first_other_key));
 
-        if let Some((self_key, self_val)) = self_cursor.peek_next() {
+        if let Some((self_key, _)) = self_cursor.peek_next() {
             match K::cmp(&first_other_key, self_key) {
                 Ordering::Equal => {
-                    // SAFETY: We read in self_val's and hand it over to our conflict function
-                    // which will always return a value that we can use to overwrite what's
-                    // in self_val
-                    unsafe {
-                        let val = ptr::read(self_val);
-                        let next_val = (conflict)(self_key, val, first_other_val);
-                        ptr::write(self_val, next_val);
-                    }
+                    self_cursor.with_next(|self_key, self_val| {
+                        conflict(self_key, self_val, first_other_val)
+                    });
                 }
                 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);
                 },
                 Ordering::Greater => {
-                    unreachable!("cursor's peek_next would return None in this case");
+                    unreachable!("Cursor's peek_next should return None.");
                 }
             }
         } else {
@@ -1345,33 +1341,29 @@ impl<K, V, A: Allocator + Clone> BTreeMap<K, V, A> {
             if self_cursor.peek_next().is_some() {
                 loop {
                     let self_entry = self_cursor.peek_next();
-                    if let Some((self_key, self_val)) = self_entry {
+                    if let Some((self_key, _)) = self_entry {
                         match K::cmp(&other_key, self_key) {
                             Ordering::Equal => {
-                                // SAFETY: We read in self_val's and hand it over to our conflict function
-                                // which will always return a value that we can use to overwrite what's
-                                // in self_val
-                                unsafe {
-                                    let val = ptr::read(self_val);
-                                    let next_val = (conflict)(self_key, val, other_val);
-                                    ptr::write(self_val, next_val);
-                                }
+                                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,
+                                // 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();
                             }
                         }
                     } else {
-                        // SAFETY: if we reach here, that means our cursor has reached
+                        // 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
@@ -1383,7 +1375,7 @@ impl<K, V, A: Allocator + Clone> BTreeMap<K, V, A> {
                     }
                 }
             } else {
-                // SAFETY: if we reach here, that means our cursor has reached
+                // 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
@@ -3417,6 +3409,100 @@ impl<'a, K, V, A> CursorMutKey<'a, K, V, A> {
 
 // Now the tree editing operations
 impl<'a, K: Ord, V, A: Allocator + Clone> CursorMutKey<'a, K, V, A> {
+    /// Calls a function with ownership of the next element's key and
+    /// and value and expects it to return a value to write
+    /// back to the next element's key and value. The cursor is not
+    /// advanced forward.
+    ///
+    /// If the cursor is at the end of the map then the function is not called
+    /// and this essentially does not do anything.
+    #[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) };
+            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 ((_, val), pos) =
+                next_kv.remove_kv_tracking(|| emptied_internal_root = true, self.alloc.clone());
+            // Don't run destructor on next element's value
+            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());
+            }
+        }
+    }
+
     /// Inserts a new key-value pair into the map in the gap that the
     /// cursor is currently pointing to.
     ///
@@ -3622,6 +3708,96 @@ impl<'a, K: Ord, V, A: Allocator + Clone> CursorMutKey<'a, K, V, A> {
 }
 
 impl<'a, K: Ord, V, A: Allocator + Clone> CursorMut<'a, K, V, A> {
+    /// Calls a function with a reference to the next element's key and
+    /// ownership of its value. The function is expected to return a value
+    /// to write back to the next element's value. The cursor is not
+    /// advanced forward.
+    ///
+    /// 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());
+            }
+        }
+    }
+
     /// Inserts a new key-value pair into the map in the gap that the
     /// cursor is currently pointing to.
     ///

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

@@ -2239,6 +2239,45 @@ fn test_append_ord_chaos() {
     map2.check();
 }
 
+#[test]
+#[cfg_attr(not(panic = "unwind"), ignore = "test requires unwinding support")]
+fn test_merge_drop_leak() {
+    let a = CrashTestDummy::new(0);
+    let b = CrashTestDummy::new(1);
+    let c = CrashTestDummy::new(2);
+    let mut left = BTreeMap::new();
+    let mut right = BTreeMap::new();
+    left.insert(a.spawn(Panic::Never), ());
+    left.insert(b.spawn(Panic::Never), ());
+    left.insert(c.spawn(Panic::Never), ());
+    right.insert(b.spawn(Panic::InDrop), ()); // first duplicate key, dropped during append
+    right.insert(c.spawn(Panic::Never), ());
+
+    catch_unwind(move || left.merge(right, |_, _, _| ())).unwrap_err();
+    assert_eq!(a.dropped(), 1);
+    assert_eq!(b.dropped(), 2);
+    assert_eq!(c.dropped(), 2);
+}
+
+#[test]
+fn test_merge_ord_chaos() {
+    let mut map1 = BTreeMap::new();
+    map1.insert(Cyclic3::A, ());
+    map1.insert(Cyclic3::B, ());
+    let mut map2 = BTreeMap::new();
+    map2.insert(Cyclic3::A, ());
+    map2.insert(Cyclic3::B, ());
+    map2.insert(Cyclic3::C, ()); // lands first, before A
+    map2.insert(Cyclic3::B, ()); // lands first, before C
+    map1.check();
+    map2.check(); // keys are not unique but still strictly ascending
+    assert_eq!(map1.len(), 2);
+    assert_eq!(map2.len(), 4);
+    map1.merge(map2, |_, _, _| ());
+    assert_eq!(map1.len(), 5);
+    map1.check();
+}
+
 fn rand_data(len: usize) -> Vec<(u32, u32)> {
     let mut rng = DeterministicRng::new();
     Vec::from_iter((0..len).map(|_| (rng.next(), rng.next())))