diff options
Diffstat (limited to 'library/alloc/src/collections/btree/map')
-rw-r--r-- | library/alloc/src/collections/btree/map/entry.rs | 40 | ||||
-rw-r--r-- | library/alloc/src/collections/btree/map/tests.rs | 49 |
2 files changed, 70 insertions, 19 deletions
diff --git a/library/alloc/src/collections/btree/map/entry.rs b/library/alloc/src/collections/btree/map/entry.rs index 370b58864..e9366eec9 100644 --- a/library/alloc/src/collections/btree/map/entry.rs +++ b/library/alloc/src/collections/btree/map/entry.rs @@ -347,7 +347,7 @@ impl<'a, K: Ord, V, A: Allocator + Clone> VacantEntry<'a, K, V, A> { /// assert_eq!(map["poneyland"], 37); /// ``` #[stable(feature = "rust1", since = "1.0.0")] - pub fn insert(self, value: V) -> &'a mut V { + pub fn insert(mut self, value: V) -> &'a mut V { let out_ptr = match self.handle { None => { // SAFETY: There is no tree yet so no reference to it exists. @@ -358,25 +358,27 @@ impl<'a, K: Ord, V, A: Allocator + Clone> VacantEntry<'a, K, V, A> { map.length = 1; val_ptr } - Some(handle) => match handle.insert_recursing(self.key, value, self.alloc.clone()) { - (None, val_ptr) => { - // SAFETY: We have consumed self.handle. - let map = unsafe { self.dormant_map.awaken() }; - map.length += 1; - val_ptr - } - (Some(ins), val_ptr) => { - drop(ins.left); - // SAFETY: We have consumed self.handle and dropped the - // remaining reference to the tree, ins.left. - let map = unsafe { self.dormant_map.awaken() }; - let root = map.root.as_mut().unwrap(); // same as ins.left - root.push_internal_level(self.alloc).push(ins.kv.0, ins.kv.1, ins.right); - map.length += 1; - val_ptr - } - }, + Some(handle) => { + let new_handle = + handle.insert_recursing(self.key, value, self.alloc.clone(), |ins| { + drop(ins.left); + // SAFETY: Pushing a new root node doesn't invalidate + // handles to existing nodes. + let map = unsafe { self.dormant_map.reborrow() }; + let root = map.root.as_mut().unwrap(); // same as ins.left + root.push_internal_level(self.alloc).push(ins.kv.0, ins.kv.1, ins.right) + }); + + // Get the pointer to the value + let val_ptr = new_handle.into_val_mut(); + + // SAFETY: We have consumed self.handle. + let map = unsafe { self.dormant_map.awaken() }; + map.length += 1; + val_ptr + } }; + // Now that we have finished growing the tree using borrowed references, // dereference the pointer to a part of it, that we picked up along the way. unsafe { &mut *out_ptr } diff --git a/library/alloc/src/collections/btree/map/tests.rs b/library/alloc/src/collections/btree/map/tests.rs index 700b1463b..76c2f27b4 100644 --- a/library/alloc/src/collections/btree/map/tests.rs +++ b/library/alloc/src/collections/btree/map/tests.rs @@ -2336,3 +2336,52 @@ fn from_array() { let unordered_duplicates = BTreeMap::from([(3, 4), (1, 2), (1, 2)]); assert_eq!(map, unordered_duplicates); } + +#[test] +fn test_cursor() { + let map = BTreeMap::from([(1, 'a'), (2, 'b'), (3, 'c')]); + + let mut cur = map.lower_bound(Bound::Unbounded); + assert_eq!(cur.key(), Some(&1)); + cur.move_next(); + assert_eq!(cur.key(), Some(&2)); + assert_eq!(cur.peek_next(), Some((&3, &'c'))); + cur.move_prev(); + assert_eq!(cur.key(), Some(&1)); + assert_eq!(cur.peek_prev(), None); + + let mut cur = map.upper_bound(Bound::Excluded(&1)); + assert_eq!(cur.key(), None); + cur.move_next(); + assert_eq!(cur.key(), Some(&1)); + cur.move_prev(); + assert_eq!(cur.key(), None); + assert_eq!(cur.peek_prev(), Some((&3, &'c'))); +} + +#[test] +fn test_cursor_mut() { + let mut map = BTreeMap::from([(1, 'a'), (3, 'c'), (5, 'e')]); + let mut cur = map.lower_bound_mut(Bound::Excluded(&3)); + assert_eq!(cur.key(), Some(&5)); + cur.insert_before(4, 'd'); + assert_eq!(cur.key(), Some(&5)); + assert_eq!(cur.peek_prev(), Some((&4, &mut 'd'))); + cur.move_next(); + assert_eq!(cur.key(), None); + cur.insert_before(6, 'f'); + assert_eq!(cur.key(), None); + assert_eq!(cur.remove_current(), None); + assert_eq!(cur.key(), None); + cur.insert_after(0, '?'); + assert_eq!(cur.key(), None); + assert_eq!(map, BTreeMap::from([(0, '?'), (1, 'a'), (3, 'c'), (4, 'd'), (5, 'e'), (6, 'f')])); + + let mut cur = map.upper_bound_mut(Bound::Included(&5)); + assert_eq!(cur.key(), Some(&5)); + assert_eq!(cur.remove_current(), Some((5, 'e'))); + assert_eq!(cur.key(), Some(&6)); + assert_eq!(cur.remove_current_and_move_back(), Some((6, 'f'))); + assert_eq!(cur.key(), Some(&4)); + assert_eq!(map, BTreeMap::from([(0, '?'), (1, 'a'), (3, 'c'), (4, 'd')])); +} |