@@ -269,6 +269,31 @@ impl<'a, K: Ord, V, A: Allocator + Clone> Entry<'a, K, V, A> {
269
269
Vacant ( entry) => Vacant ( entry) ,
270
270
}
271
271
}
272
+
273
+ /// Sets the value of the entry, and returns an `OccupiedEntry`.
274
+ ///
275
+ /// # Examples
276
+ ///
277
+ /// ```
278
+ /// #![feature(btree_entry_insert)]
279
+ /// use std::collections::BTreeMap;
280
+ ///
281
+ /// let mut map: BTreeMap<&str, String> = BTreeMap::new();
282
+ /// let entry = map.entry("poneyland").insert_entry("hoho".to_string());
283
+ ///
284
+ /// assert_eq!(entry.key(), &"poneyland");
285
+ /// ```
286
+ #[ inline]
287
+ #[ unstable( feature = "btree_entry_insert" , issue = "65225" ) ]
288
+ pub fn insert_entry ( self , value : V ) -> OccupiedEntry < ' a , K , V , A > {
289
+ match self {
290
+ Occupied ( mut entry) => {
291
+ entry. insert ( value) ;
292
+ entry
293
+ }
294
+ Vacant ( entry) => entry. insert_entry ( value) ,
295
+ }
296
+ }
272
297
}
273
298
274
299
impl < ' a , K : Ord , V : Default , A : Allocator + Clone > Entry < ' a , K , V , A > {
@@ -348,41 +373,61 @@ impl<'a, K: Ord, V, A: Allocator + Clone> VacantEntry<'a, K, V, A> {
348
373
/// ```
349
374
#[ stable( feature = "rust1" , since = "1.0.0" ) ]
350
375
#[ rustc_confusables( "push" , "put" ) ]
351
- pub fn insert ( mut self , value : V ) -> & ' a mut V {
352
- let out_ptr = match self . handle {
376
+ pub fn insert ( self , value : V ) -> & ' a mut V {
377
+ self . insert_entry ( value) . into_mut ( )
378
+ }
379
+
380
+ /// Sets the value of the entry with the `VacantEntry`'s key,
381
+ /// and returns an `OccupiedEntry`.
382
+ ///
383
+ /// # Examples
384
+ ///
385
+ /// ```
386
+ /// #![feature(btree_entry_insert)]
387
+ /// use std::collections::BTreeMap;
388
+ /// use std::collections::btree_map::Entry;
389
+ ///
390
+ /// let mut map: BTreeMap<&str, u32> = BTreeMap::new();
391
+ ///
392
+ /// if let Entry::Vacant(o) = map.entry("poneyland") {
393
+ /// let entry = o.insert_entry(37);
394
+ /// assert_eq!(entry.get(), &37);
395
+ /// }
396
+ /// assert_eq!(map["poneyland"], 37);
397
+ /// ```
398
+ #[ unstable( feature = "btree_entry_insert" , issue = "65225" ) ]
399
+ pub fn insert_entry ( mut self , value : V ) -> OccupiedEntry < ' a , K , V , A > {
400
+ let handle = match self . handle {
353
401
None => {
354
402
// SAFETY: There is no tree yet so no reference to it exists.
355
- let map = unsafe { self . dormant_map . awaken ( ) } ;
356
- let mut root = NodeRef :: new_leaf ( self . alloc . clone ( ) ) ;
357
- let val_ptr = root. borrow_mut ( ) . push ( self . key , value) ;
358
- map. root = Some ( root. forget_type ( ) ) ;
359
- map. length = 1 ;
360
- val_ptr
361
- }
362
- Some ( handle) => {
363
- let new_handle =
364
- handle. insert_recursing ( self . key , value, self . alloc . clone ( ) , |ins| {
365
- drop ( ins. left ) ;
366
- // SAFETY: Pushing a new root node doesn't invalidate
367
- // handles to existing nodes.
368
- let map = unsafe { self . dormant_map . reborrow ( ) } ;
369
- let root = map. root . as_mut ( ) . unwrap ( ) ; // same as ins.left
370
- root. push_internal_level ( self . alloc ) . push ( ins. kv . 0 , ins. kv . 1 , ins. right )
371
- } ) ;
372
-
373
- // Get the pointer to the value
374
- let val_ptr = new_handle. into_val_mut ( ) ;
375
-
376
- // SAFETY: We have consumed self.handle.
377
- let map = unsafe { self . dormant_map . awaken ( ) } ;
378
- map. length += 1 ;
379
- val_ptr
403
+ let map = unsafe { self . dormant_map . reborrow ( ) } ;
404
+ let root = map. root . insert ( NodeRef :: new_leaf ( self . alloc . clone ( ) ) . forget_type ( ) ) ;
405
+ // SAFETY: We *just* created the root as a leaf, and we're
406
+ // stacking the new handle on the original borrow lifetime.
407
+ unsafe {
408
+ let mut leaf = root. borrow_mut ( ) . cast_to_leaf_unchecked ( ) ;
409
+ leaf. push_with_handle ( self . key , value)
410
+ }
380
411
}
412
+ Some ( handle) => handle. insert_recursing ( self . key , value, self . alloc . clone ( ) , |ins| {
413
+ drop ( ins. left ) ;
414
+ // SAFETY: Pushing a new root node doesn't invalidate
415
+ // handles to existing nodes.
416
+ let map = unsafe { self . dormant_map . reborrow ( ) } ;
417
+ let root = map. root . as_mut ( ) . unwrap ( ) ; // same as ins.left
418
+ root. push_internal_level ( self . alloc . clone ( ) ) . push ( ins. kv . 0 , ins. kv . 1 , ins. right )
419
+ } ) ,
381
420
} ;
382
421
383
- // Now that we have finished growing the tree using borrowed references,
384
- // dereference the pointer to a part of it, that we picked up along the way.
385
- unsafe { & mut * out_ptr }
422
+ // SAFETY: modifying the length doesn't invalidate handles to existing nodes.
423
+ unsafe { self . dormant_map . reborrow ( ) . length += 1 } ;
424
+
425
+ OccupiedEntry {
426
+ handle : handle. forget_node_type ( ) ,
427
+ dormant_map : self . dormant_map ,
428
+ alloc : self . alloc ,
429
+ _marker : PhantomData ,
430
+ }
386
431
}
387
432
}
388
433
0 commit comments