insert_mut, and friends

Author: balt-dev

library/alloc/src/collections/linked_list.rs

@@ -844,11 +844,34 @@ impl<T, A: Allocator> LinkedList<T, A> {
     /// ```
     #[stable(feature = "rust1", since = "1.0.0")]
     pub fn push_front(&mut self, elt: T) {
+        let _ = self.push_front_mut(elt);
+    }
+
+    /// Adds an element first in the list, returning a reference to it.
+    ///
+    /// This operation should compute in *O*(1) time.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(push_mut)]
+    /// use std::collections::LinkedList;
+    ///
+    /// let mut dl = LinkedList::from([1, 2, 3]);
+    ///
+    /// let ptr = dl.push_front_mut(2);
+    /// *ptr += 4;
+    /// assert_eq!(dl.front().unwrap(), &6);
+    /// ```
+    #[unstable(feature = "push_mut", issue = "135974")]
+    #[must_use = "if you don't need a reference to the value, use LinkedList::push_front instead"]
+    pub fn push_front_mut(&mut self, elt: T) -> &mut T {
         let node = Box::new_in(Node::new(elt), &self.alloc);
-        let node_ptr = NonNull::from(Box::leak(node));
+        let mut node_ptr = NonNull::from(Box::leak(node));
         // SAFETY: node_ptr is a unique pointer to a node we boxed with self.alloc and leaked
         unsafe {
             self.push_front_node(node_ptr);
+            &mut node_ptr.as_mut().element
         }
     }
 
@@ -893,11 +916,34 @@ impl<T, A: Allocator> LinkedList<T, A> {
     #[stable(feature = "rust1", since = "1.0.0")]
     #[rustc_confusables("push", "append")]
     pub fn push_back(&mut self, elt: T) {
+        let _ = self.push_back_mut(elt);
+    }
+
+    /// Adds an element last in the list, returning a reference to it.
+    ///
+    /// This operation should compute in *O*(1) time.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(push_mut)]
+    /// use std::collections::LinkedList;
+    ///
+    /// let mut dl = LinkedList::from([1, 2, 3]);
+    ///
+    /// let ptr = dl.push_back_mut(2);
+    /// *ptr += 4;
+    /// assert_eq!(dl.back().unwrap(), &6);
+    /// ```
+    #[unstable(feature = "push_mut", issue = "135974")]
+    #[must_use = "if you don't need a reference to the value, use LinkedList::push_back instead"]
+    pub fn push_back_mut(&mut self, elt: T) -> &mut T {
         let node = Box::new_in(Node::new(elt), &self.alloc);
-        let node_ptr = NonNull::from(Box::leak(node));
+        let mut node_ptr = NonNull::from(Box::leak(node));
         // SAFETY: node_ptr is a unique pointer to a node we boxed with self.alloc and leaked
         unsafe {
             self.push_back_node(node_ptr);
+            &mut node_ptr.as_mut().element
         }
     }
 

library/alloc/src/collections/vec_deque/mod.rs

@@ -182,11 +182,13 @@ impl<T, A: Allocator> VecDeque<T, A> {
         unsafe { ptr::read(self.ptr().add(off)) }
     }
 
-    /// Writes an element into the buffer, moving it.
+    /// Writes an element into the buffer, moving it and returning a pointer to it.
     #[inline]
-    unsafe fn buffer_write(&mut self, off: usize, value: T) {
+    unsafe fn buffer_write(&mut self, off: usize, value: T) -> *mut T {
         unsafe {
-            ptr::write(self.ptr().add(off), value);
+            let ptr = self.ptr().add(off);
+            ptr::write(ptr, value);
+            ptr
         }
     }
 
@@ -1920,10 +1922,66 @@ impl<T, A: Allocator> VecDeque<T, A> {
             self.grow();
         }
 
-        unsafe { self.buffer_write(self.to_physical_idx(self.len), value) }
+        unsafe { self.buffer_write(self.to_physical_idx(self.len), value) };
         self.len += 1;
     }
 
+    /// Prepends an element to the deque, returning a reference to it.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(push_mut)]
+    /// use std::collections::VecDeque;
+    ///
+    /// let mut d = VecDeque::from([1, 2, 3]);
+    /// let x = d.push_front_mut(1);
+    /// *x -= 1;
+    /// assert_eq!(d.front(), Some(&0));
+    /// ```
+    #[unstable(feature = "push_mut", issue = "135974")]
+    #[track_caller]
+    #[must_use = "if you don't need a reference to the value, use VecDeque::push_front instead"]
+    pub fn push_front_mut(&mut self, value: T) -> &mut T {
+        if self.is_full() {
+            self.grow();
+        }
+
+        self.head = self.wrap_sub(self.head, 1);
+        self.len += 1;
+
+        unsafe {
+            let ptr = self.buffer_write(self.head, value);
+            &mut *ptr
+        }
+    }
+
+    /// Appends an element to the back of the deque, returning a reference to it.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(push_mut)]
+    /// use std::collections::VecDeque;
+    ///
+    /// let mut d = VecDeque::from([1, 2, 3]);
+    /// let x = d.push_back_mut(3);
+    /// *x += 1;
+    /// assert_eq!(d.back(), Some(&4));
+    /// ```
+    #[unstable(feature = "push_mut", issue = "135974")]
+    #[track_caller]
+    #[must_use = "if you don't need a reference to the value, use VecDeque::push_back instead"]
+    pub fn push_back_mut(&mut self, value: T) -> &mut T {
+        if self.is_full() {
+            self.grow();
+        }
+
+        let back_ptr = unsafe { self.buffer_write(self.to_physical_idx(self.len), value) };
+        self.len += 1;
+        unsafe { &mut *back_ptr }
+    }
+
     #[inline]
     fn is_contiguous(&self) -> bool {
         // Do the calculation like this to avoid overflowing if len + head > usize::MAX
@@ -2056,6 +2114,64 @@ impl<T, A: Allocator> VecDeque<T, A> {
         }
     }
 
+    /// Inserts an element at `index` within the deque, shifting all elements
+    /// with indices greater than or equal to `index` towards the back,
+    /// and returning a reference to it.
+    ///
+    /// Element at index 0 is the front of the queue.
+    ///
+    /// Returns [`None`] if `index` is strictly greater than deque's length.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(push_mut)]
+    /// use std::collections::VecDeque;
+    ///
+    /// let mut vec_deque = VecDeque::from([1, 2, 3]);
+    ///
+    /// let x = vec_deque.insert_mut(1, 5).unwrap();
+    /// *x += 7;
+    /// assert_eq!(vec_deque, &[1, 12, 2, 3]);
+    ///
+    /// let y = vec_deque.insert_mut(7, 5);
+    /// assert!(y.is_none());
+    /// ```
+    #[unstable(feature = "push_mut", issue = "135974")]
+    #[track_caller]
+    #[must_use = "if you don't need a reference to the value or type-safe bound checking, use VecDeque::insert instead"]
+    pub fn insert_mut(&mut self, index: usize, value: T) -> Option<&mut T> {
+        if index > self.len() {
+            return None;
+        }
+        if self.is_full() {
+            self.grow();
+        }
+
+        let k = self.len - index;
+        if k < index {
+            // `index + 1` can't overflow, because if index was usize::MAX, then either the
+            // assert would've failed, or the deque would've tried to grow past usize::MAX
+            // and panicked.
+            unsafe {
+                // see `remove()` for explanation why this wrap_copy() call is safe.
+                self.wrap_copy(self.to_physical_idx(index), self.to_physical_idx(index + 1), k);
+                let ptr = self.buffer_write(self.to_physical_idx(index), value);
+                self.len += 1;
+                Some(&mut *ptr)
+            }
+        } else {
+            let old_head = self.head;
+            self.head = self.wrap_sub(self.head, 1);
+            unsafe {
+                self.wrap_copy(old_head, self.head, index);
+                let ptr = self.buffer_write(self.to_physical_idx(index), value);
+                self.len += 1;
+                Some(&mut *ptr)
+            }
+        }
+    }
+
     /// Removes and returns the element at `index` from the deque.
     /// Whichever end is closer to the removal point will be moved to make
     /// room, and all the affected elements will be moved to new positions.

library/alloc/src/vec/mod.rs

@@ -1974,6 +1974,8 @@ impl<T, A: Allocator> Vec<T, A> {
         }
     }
 
+   
+
     /// Inserts an element at position `index` within the vector, shifting all
     /// elements after it to the right.
     ///
@@ -2000,7 +2002,7 @@ impl<T, A: Allocator> Vec<T, A> {
     #[stable(feature = "rust1", since = "1.0.0")]
     #[track_caller]
     pub fn insert(&mut self, index: usize, element: T) {
-        #[cold]
+         #[cold]
         #[cfg_attr(not(feature = "panic_immediate_abort"), inline(never))]
         #[track_caller]
         #[optimize(size)]
@@ -2036,6 +2038,64 @@ impl<T, A: Allocator> Vec<T, A> {
         }
     }
 
+    /// Inserts an element at position `index` within the vector, shifting all
+    /// elements after it to the right.
+    /// 
+    /// Returns [`None`] if the specified index is out of bounds.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(push_mut)]
+    /// let mut vec = vec![1, 3, 5, 9];
+    /// let x = vec.insert_mut(3, 6).unwrap();
+    /// *x += 1;
+    /// assert_eq!(vec, [1, 3, 5, 7, 9]);
+    /// 
+    /// let y = vec.insert_mut(7, 6);
+    /// assert!(y.is_none());
+    /// ```
+    ///
+    /// # Time complexity
+    ///
+    /// Takes *O*([`Vec::len`]) time. All items after the insertion index must be
+    /// shifted to the right. In the worst case, all elements are shifted when
+    /// the insertion index is 0.
+    #[cfg(not(no_global_oom_handling))]
+    #[inline]
+    #[unstable(feature = "push_mut", issue = "135974")]
+    #[track_caller]
+    #[must_use = "if you don't need a reference to the value or type-safe bound checking, use Vec::insert instead"]
+    pub fn insert_mut(&mut self, index: usize, element: T) -> Option<&mut T> {
+        let len = self.len();
+        if index > len {
+            return None;
+        }
+
+        // space for the new element
+        if len == self.buf.capacity() {
+            self.buf.grow_one();
+        }
+
+        unsafe {
+            // infallible
+            // The spot to put the new value
+            let p = self.as_mut_ptr().add(index);
+            {
+                if index < len {
+                    // Shift everything over to make space. (Duplicating the
+                    // `index`th element into two consecutive places.)
+                    ptr::copy(p, p.add(1), len - index);
+                }
+                // Write it in, overwriting the first copy of the `index`th
+                // element.
+                ptr::write(p, element);
+            }
+            self.set_len(len + 1);
+            Some(&mut *p)
+        }
+    }
+
     /// Removes and returns the element at position `index` within the vector,
     /// shifting all elements after it to the left.
     ///