Auto merge of #51466 - joshlf:ref-split, r=dtolnay

Add Ref/RefMut map_split method

As proposed [here](https://internals.rust-lang.org/t/make-refcell-support-slice-splitting/7707).

TLDR: Add a `map_split` method that allows multiple `RefMut`s to exist simultaneously so long as they refer to non-overlapping regions of the original `RefCell`. This is useful for things like the slice `split_at_mut` method.

src/libcore/cell.rs

@@ -570,11 +570,20 @@ fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
     }
 }
 
-// Values [1, MAX-1] represent the number of `Ref` active
-// (will not outgrow its range since `usize` is the size of the address space)
+// Values [1, MIN_WRITING-1] represent the number of `Ref` active. Values in
+// [MIN_WRITING, MAX-1] represent the number of `RefMut` active. Multiple
+// `RefMut`s can only be active at a time if they refer to distinct,
+// nonoverlapping components of a `RefCell` (e.g., different ranges of a slice).
+//
+// `Ref` and `RefMut` are both two words in size, and so there will likely never
+// be enough `Ref`s or `RefMut`s in existence to overflow half of the `usize`
+// range. Thus, a `BorrowFlag` will probably never overflow. However, this is
+// not a guarantee, as a pathological program could repeatedly create and then
+// mem::forget `Ref`s or `RefMut`s. Thus, all code must explicitly check for
+// overflow in order to avoid unsafety.
 type BorrowFlag = usize;
 const UNUSED: BorrowFlag = 0;
-const WRITING: BorrowFlag = !0;
+const MIN_WRITING: BorrowFlag = (!0)/2 + 1; // 0b1000...
 
 impl<T> RefCell<T> {
     /// Creates a new `RefCell` containing `value`.
@@ -775,8 +784,9 @@ pub fn try_borrow(&self) -> Result<Ref<T>, BorrowError> {
 
     /// Mutably borrows the wrapped value.
     ///
-    /// The borrow lasts until the returned `RefMut` exits scope. The value
-    /// cannot be borrowed while this borrow is active.
+    /// The borrow lasts until the returned `RefMut` or all `RefMut`s derived
+    /// from it exit scope. The value cannot be borrowed while this borrow is
+    /// active.
     ///
     /// # Panics
     ///
@@ -818,8 +828,9 @@ pub fn borrow_mut(&self) -> RefMut<T> {
 
     /// Mutably borrows the wrapped value, returning an error if the value is currently borrowed.
     ///
-    /// The borrow lasts until the returned `RefMut` exits scope. The value cannot be borrowed
-    /// while this borrow is active.
+    /// The borrow lasts until the returned `RefMut` or all `RefMut`s derived
+    /// from it exit scope. The value cannot be borrowed while this borrow is
+    /// active.
     ///
     /// This is the non-panicking variant of [`borrow_mut`](#method.borrow_mut).
     ///
@@ -1010,12 +1021,15 @@ struct BorrowRef<'b> {
 impl<'b> BorrowRef<'b> {
     #[inline]
     fn new(borrow: &'b Cell<BorrowFlag>) -> Option<BorrowRef<'b>> {
-        match borrow.get() {
-            WRITING => None,
-            b => {
-                borrow.set(b + 1);
-                Some(BorrowRef { borrow: borrow })
-            },
+        let b = borrow.get();
+        if b >= MIN_WRITING {
+            None
+        } else {
+            // Prevent the borrow counter from overflowing into
+            // a writing borrow.
+            assert!(b < MIN_WRITING - 1);
+            borrow.set(b + 1);
+            Some(BorrowRef { borrow })
         }
     }
 }
@@ -1024,7 +1038,7 @@ impl<'b> Drop for BorrowRef<'b> {
     #[inline]
     fn drop(&mut self) {
         let borrow = self.borrow.get();
-        debug_assert!(borrow != WRITING && borrow != UNUSED);
+        debug_assert!(borrow < MIN_WRITING && borrow != UNUSED);
         self.borrow.set(borrow - 1);
     }
 }
@@ -1036,8 +1050,9 @@ fn clone(&self) -> BorrowRef<'b> {
         // is not set to WRITING.
         let borrow = self.borrow.get();
         debug_assert!(borrow != UNUSED);
-        // Prevent the borrow counter from overflowing.
-        assert!(borrow != WRITING);
+        // Prevent the borrow counter from overflowing into
+        // a writing borrow.
+        assert!(borrow < MIN_WRITING - 1);
         self.borrow.set(borrow + 1);
         BorrowRef { borrow: self.borrow }
     }
@@ -1109,6 +1124,37 @@ pub fn map<U: ?Sized, F>(orig: Ref<'b, T>, f: F) -> Ref<'b, U>
             borrow: orig.borrow,
         }
     }
+
+    /// Split a `Ref` into multiple `Ref`s for different components of the
+    /// borrowed data.
+    ///
+    /// The `RefCell` is already immutably borrowed, so this cannot fail.
+    ///
+    /// This is an associated function that needs to be used as
+    /// `Ref::map_split(...)`. A method would interfere with methods of the same
+    /// name on the contents of a `RefCell` used through `Deref`.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(refcell_map_split)]
+    /// use std::cell::{Ref, RefCell};
+    ///
+    /// let cell = RefCell::new([1, 2, 3, 4]);
+    /// let borrow = cell.borrow();
+    /// let (begin, end) = Ref::map_split(borrow, |slice| slice.split_at(2));
+    /// assert_eq!(*begin, [1, 2]);
+    /// assert_eq!(*end, [3, 4]);
+    /// ```
+    #[unstable(feature = "refcell_map_split", issue = "51476")]
+    #[inline]
+    pub fn map_split<U: ?Sized, V: ?Sized, F>(orig: Ref<'b, T>, f: F) -> (Ref<'b, U>, Ref<'b, V>)
+        where F: FnOnce(&T) -> (&U, &V)
+    {
+        let (a, b) = f(orig.value);
+        let borrow = orig.borrow.clone();
+        (Ref { value: a, borrow }, Ref { value: b, borrow: orig.borrow })
+    }
 }
 
 #[unstable(feature = "coerce_unsized", issue = "27732")]
@@ -1157,6 +1203,44 @@ pub fn map<U: ?Sized, F>(orig: RefMut<'b, T>, f: F) -> RefMut<'b, U>
             borrow: borrow,
         }
     }
+
+    /// Split a `RefMut` into multiple `RefMut`s for different components of the
+    /// borrowed data.
+    ///
+    /// The underlying `RefCell` will remain mutably borrowed until both
+    /// returned `RefMut`s go out of scope.
+    ///
+    /// The `RefCell` is already mutably borrowed, so this cannot fail.
+    ///
+    /// This is an associated function that needs to be used as
+    /// `RefMut::map_split(...)`. A method would interfere with methods of the
+    /// same name on the contents of a `RefCell` used through `Deref`.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(refcell_map_split)]
+    /// use std::cell::{RefCell, RefMut};
+    ///
+    /// let cell = RefCell::new([1, 2, 3, 4]);
+    /// let borrow = cell.borrow_mut();
+    /// let (mut begin, mut end) = RefMut::map_split(borrow, |slice| slice.split_at_mut(2));
+    /// assert_eq!(*begin, [1, 2]);
+    /// assert_eq!(*end, [3, 4]);
+    /// begin.copy_from_slice(&[4, 3]);
+    /// end.copy_from_slice(&[2, 1]);
+    /// ```
+    #[unstable(feature = "refcell_map_split", issue = "51476")]
+    #[inline]
+    pub fn map_split<U: ?Sized, V: ?Sized, F>(
+        orig: RefMut<'b, T>, f: F
+    ) -> (RefMut<'b, U>, RefMut<'b, V>)
+        where F: FnOnce(&mut T) -> (&mut U, &mut V)
+    {
+        let (a, b) = f(orig.value);
+        let borrow = orig.borrow.clone();
+        (RefMut { value: a, borrow }, RefMut { value: b, borrow: orig.borrow })
+    }
 }
 
 struct BorrowRefMut<'b> {
@@ -1167,22 +1251,45 @@ impl<'b> Drop for BorrowRefMut<'b> {
     #[inline]
     fn drop(&mut self) {
         let borrow = self.borrow.get();
-        debug_assert!(borrow == WRITING);
-        self.borrow.set(UNUSED);
+        debug_assert!(borrow >= MIN_WRITING);
+        self.borrow.set(if borrow == MIN_WRITING {
+            UNUSED
+        } else {
+            borrow - 1
+        });
     }
 }
 
 impl<'b> BorrowRefMut<'b> {
     #[inline]
     fn new(borrow: &'b Cell<BorrowFlag>) -> Option<BorrowRefMut<'b>> {
+        // NOTE: Unlike BorrowRefMut::clone, new is called to create the initial
+        // mutable reference, and so there must currently be no existing
+        // references. Thus, while clone increments the mutable refcount, here
+        // we simply go directly from UNUSED to MIN_WRITING.
         match borrow.get() {
             UNUSED => {
-                borrow.set(WRITING);
+                borrow.set(MIN_WRITING);
                 Some(BorrowRefMut { borrow: borrow })
             },
             _ => None,
         }
     }
+
+    // Clone a `BorrowRefMut`.
+    //
+    // This is only valid if each `BorrowRefMut` is used to track a mutable
+    // reference to a distinct, nonoverlapping range of the original object.
+    // This isn't in a Clone impl so that code doesn't call this implicitly.
+    #[inline]
+    fn clone(&self) -> BorrowRefMut<'b> {
+        let borrow = self.borrow.get();
+        debug_assert!(borrow >= MIN_WRITING);
+        // Prevent the borrow counter from overflowing.
+        assert!(borrow != !0);
+        self.borrow.set(borrow + 1);
+        BorrowRefMut { borrow: self.borrow }
+    }
 }
 
 /// A wrapper type for a mutably borrowed value from a `RefCell<T>`.

src/libcore/tests/cell.rs

@@ -165,6 +165,64 @@ fn ref_map_does_not_update_flag() {
     assert!(x.try_borrow_mut().is_ok());
 }
 
+#[test]
+fn ref_map_split_updates_flag() {
+    let x = RefCell::new([1, 2]);
+    {
+        let b1 = x.borrow();
+        assert!(x.try_borrow().is_ok());
+        assert!(x.try_borrow_mut().is_err());
+        {
+            let (_b2, _b3) = Ref::map_split(b1, |slc| slc.split_at(1));
+            assert!(x.try_borrow().is_ok());
+            assert!(x.try_borrow_mut().is_err());
+        }
+        assert!(x.try_borrow().is_ok());
+        assert!(x.try_borrow_mut().is_ok());
+    }
+    assert!(x.try_borrow().is_ok());
+    assert!(x.try_borrow_mut().is_ok());
+
+    {
+        let b1 = x.borrow_mut();
+        assert!(x.try_borrow().is_err());
+        assert!(x.try_borrow_mut().is_err());
+        {
+            let (_b2, _b3) = RefMut::map_split(b1, |slc| slc.split_at_mut(1));
+            assert!(x.try_borrow().is_err());
+            assert!(x.try_borrow_mut().is_err());
+            drop(_b2);
+            assert!(x.try_borrow().is_err());
+            assert!(x.try_borrow_mut().is_err());
+        }
+        assert!(x.try_borrow().is_ok());
+        assert!(x.try_borrow_mut().is_ok());
+    }
+    assert!(x.try_borrow().is_ok());
+    assert!(x.try_borrow_mut().is_ok());
+}
+
+#[test]
+fn ref_map_split() {
+    let x = RefCell::new([1, 2]);
+    let (b1, b2) = Ref::map_split(x.borrow(), |slc| slc.split_at(1));
+    assert_eq!(*b1, [1]);
+    assert_eq!(*b2, [2]);
+}
+
+#[test]
+fn ref_mut_map_split() {
+    let x = RefCell::new([1, 2]);
+    {
+        let (mut b1, mut b2) = RefMut::map_split(x.borrow_mut(), |slc| slc.split_at_mut(1));
+        assert_eq!(*b1, [1]);
+        assert_eq!(*b2, [2]);
+        b1[0] = 2;
+        b2[0] = 1;
+    }
+    assert_eq!(*x.borrow(), [2, 1]);
+}
+
 #[test]
 fn ref_map_accessor() {
     struct X(RefCell<(u32, char)>);

src/libcore/tests/lib.rs

@@ -27,6 +27,7 @@
 #![feature(pattern)]
 #![feature(range_is_empty)]
 #![feature(raw)]
+#![feature(refcell_map_split)]
 #![feature(refcell_replace_swap)]
 #![feature(slice_patterns)]
 #![feature(slice_rotate)]