Re-stabilize Weak::as_ptr &friends for unsized T

As per T-lang consensus, this uses a branch to handle the dangling case. The discussed optimization of only doing the branch in the T: ?Sized case is left for a followup patch, as doing so is not trivial (as it requires specialization for correctness, not just optimization).
rust-lang · m-ou-se · Dec 18, 2020 · Dec 23, 2020 · Jan 6, 2021 · Jan 7, 2021
commit 6bc772cdc0a9021a6286ebd550c2f6221179f21d
diff --git a/library/alloc/src/lib.rs b/library/alloc/src/lib.rs
@@ -120,6 +120,7 @@
 #![feature(receiver_trait)]
 #![cfg_attr(bootstrap, feature(min_const_generics))]
 #![feature(min_specialization)]
+#![feature(set_ptr_value)]
 #![feature(slice_ptr_get)]
 #![feature(slice_ptr_len)]
 #![feature(staged_api)]

diff --git a/library/alloc/src/rc.rs b/library/alloc/src/rc.rs
@@ -1862,7 +1862,7 @@ struct WeakInner<'a> {
     strong: &'a Cell<usize>,
 }
 
-impl<T> Weak<T> {
+impl<T: ?Sized> Weak<T> {
     /// Returns a raw pointer to the object `T` pointed to by this `Weak<T>`.
     ///
     /// The pointer is valid only if there are some strong references. The pointer may be dangling,
@@ -1892,15 +1892,17 @@ impl<T> Weak<T> {
     pub fn as_ptr(&self) -> *const T {
         let ptr: *mut RcBox<T> = NonNull::as_ptr(self.ptr);
 
-        // SAFETY: we must offset the pointer manually, and said pointer may be
-        // a dangling weak (usize::MAX) if T is sized. data_offset is safe to call,
-        // because we know that a pointer to unsized T was derived from a real
-        // unsized T, as dangling weaks are only created for sized T. wrapping_offset
-        // is used so that we can use the same code path for the non-dangling
-        // unsized case and the potentially dangling sized case.
-        unsafe {
-            let offset = data_offset(ptr as *mut T);
-            set_data_ptr(ptr as *mut T, (ptr as *mut u8).wrapping_offset(offset))
+        if is_dangling(self.ptr) {
+            // If the pointer is dangling, we return a null pointer as the dangling sentinel.
+            // We can't return the usize::MAX sentinel, as that could valid if T is ZST.
+            // SAFETY: we have to return a known sentinel here that cannot be produced for
+            // a valid pointer, so that `from_raw` can reverse this transformation.
+            (ptr as *mut T).set_ptr_value(ptr::null_mut())
+        } else {
+            // SAFETY: If the pointer is not dangling, it describes to a valid allocation.
+            // The payload may be dropped at this point, and we have to maintain provenance,
+            // so use raw pointer manipulation.
+            unsafe { &raw mut (*ptr).value }
         }
     }
 
@@ -1982,22 +1984,25 @@ impl<T> Weak<T> {
     /// [`new`]: Weak::new
     #[stable(feature = "weak_into_raw", since = "1.45.0")]
     pub unsafe fn from_raw(ptr: *const T) -> Self {
-        // SAFETY: data_offset is safe to call, because this pointer originates from a Weak.
         // See Weak::as_ptr for context on how the input pointer is derived.
-        let offset = unsafe { data_offset(ptr) };
 
-        // Reverse the offset to find the original RcBox.
-        // SAFETY: we use wrapping_offset here because the pointer may be dangling (but only if T: Sized).
-        let ptr = unsafe {
-            set_data_ptr(ptr as *mut RcBox<T>, (ptr as *mut u8).wrapping_offset(-offset))
+        let ptr = if ptr.is_null() {
+            // If we get a null pointer, this is a dangling weak.
+            // SAFETY: this is the same sentinel as used in Weak::new and is_dangling
+            (ptr as *mut RcBox<T>).set_ptr_value(usize::MAX as *mut _)
+        } else {
+            // Otherwise, this describes a real allocation.
+            // SAFETY: data_offset is safe to call, as ptr describes a real allocation.
+            let offset = unsafe { data_offset(ptr) };
+            // Thus, we reverse the offset to get the whole RcBox.
+            // SAFETY: the pointer originated from a Weak, so this offset is safe.
+            unsafe { (ptr as *mut RcBox<T>).set_ptr_value((ptr as *mut u8).offset(-offset)) }
         };
 
         // SAFETY: we now have recovered the original Weak pointer, so can create the Weak.
         Weak { ptr: unsafe { NonNull::new_unchecked(ptr) } }
     }
-}
 
-impl<T: ?Sized> Weak<T> {
     /// Attempts to upgrade the `Weak` pointer to an [`Rc`], delaying
     /// dropping of the inner value if successful.
     ///

diff --git a/library/alloc/src/rc/tests.rs b/library/alloc/src/rc/tests.rs
@@ -208,6 +208,30 @@ fn into_from_weak_raw() {
     }
 }
 
+#[test]
+fn test_into_from_weak_raw_unsized() {
+    use std::fmt::Display;
+    use std::string::ToString;
+
+    let arc: Rc<str> = Rc::from("foo");
+    let weak: Weak<str> = Rc::downgrade(&arc);
+
+    let ptr = Weak::into_raw(weak.clone());
+    let weak2 = unsafe { Weak::from_raw(ptr) };
+
+    assert_eq!(unsafe { &*ptr }, "foo");
+    assert!(weak.ptr_eq(&weak2));
+
+    let arc: Rc<dyn Display> = Rc::new(123);
+    let weak: Weak<dyn Display> = Rc::downgrade(&arc);
+
+    let ptr = Weak::into_raw(weak.clone());
+    let weak2 = unsafe { Weak::from_raw(ptr) };
+
+    assert_eq!(unsafe { &*ptr }.to_string(), "123");
+    assert!(weak.ptr_eq(&weak2));
+}
+
 #[test]
 fn get_mut() {
     let mut x = Rc::new(3);
@@ -294,6 +318,23 @@ fn test_unsized() {
     assert_eq!(foo, foo.clone());
 }
 
+#[test]
+fn test_maybe_thin_unsized() {
+    // If/when custom thin DSTs exist, this test should be updated to use one
+    use std::ffi::{CStr, CString};
+
+    let x: Rc<CStr> = Rc::from(CString::new("swordfish").unwrap().into_boxed_c_str());
+    assert_eq!(format!("{:?}", x), "\"swordfish\"");
+    let y: Weak<CStr> = Rc::downgrade(&x);
+    drop(x);
+
+    // At this point, the weak points to a dropped DST
+    assert!(y.upgrade().is_none());
+    // But we still need to be able to get the alloc layout to drop.
+    // CStr has no drop glue, but custom DSTs might, and need to work.
+    drop(y);
+}
+
 #[test]
 fn test_from_owned() {
     let foo = 123;

diff --git a/library/alloc/src/sync.rs b/library/alloc/src/sync.rs
@@ -1648,7 +1648,7 @@ struct WeakInner<'a> {
     strong: &'a atomic::AtomicUsize,
 }
 
-impl<T> Weak<T> {
+impl<T: ?Sized> Weak<T> {
     /// Returns a raw pointer to the object `T` pointed to by this `Weak<T>`.
     ///
     /// The pointer is valid only if there are some strong references. The pointer may be dangling,
@@ -1678,15 +1678,17 @@ impl<T> Weak<T> {
     pub fn as_ptr(&self) -> *const T {
         let ptr: *mut ArcInner<T> = NonNull::as_ptr(self.ptr);
 
-        // SAFETY: we must offset the pointer manually, and said pointer may be
-        // a dangling weak (usize::MAX) if T is sized. data_offset is safe to call,
-        // because we know that a pointer to unsized T was derived from a real
-        // unsized T, as dangling weaks are only created for sized T. wrapping_offset
-        // is used so that we can use the same code path for the non-dangling
-        // unsized case and the potentially dangling sized case.
-        unsafe {
-            let offset = data_offset(ptr as *mut T);
-            set_data_ptr(ptr as *mut T, (ptr as *mut u8).wrapping_offset(offset))
+        if is_dangling(self.ptr) {
+            // If the pointer is dangling, we return a null pointer as the dangling sentinel.
+            // We can't return the usize::MAX sentinel, as that could valid if T is ZST.
+            // SAFETY: we have to return a known sentinel here that cannot be produced for
+            // a valid pointer, so that `from_raw` can reverse this transformation.
+            (ptr as *mut T).set_ptr_value(ptr::null_mut())
+        } else {
+            // SAFETY: If the pointer is not dangling, it describes to a valid allocation.
+            // The payload may be dropped at this point, and we have to maintain provenance,
+            // so use raw pointer manipulation.
+            unsafe { &raw mut (*ptr).data }
         }
     }
 
@@ -1768,18 +1770,23 @@ impl<T> Weak<T> {
     /// [`forget`]: std::mem::forget
     #[stable(feature = "weak_into_raw", since = "1.45.0")]
     pub unsafe fn from_raw(ptr: *const T) -> Self {
-        // SAFETY: data_offset is safe to call, because this pointer originates from a Weak.
         // See Weak::as_ptr for context on how the input pointer is derived.
-        let offset = unsafe { data_offset(ptr) };
 
-        // Reverse the offset to find the original ArcInner.
-        // SAFETY: we use wrapping_offset here because the pointer may be dangling (but only if T: Sized)
-        let ptr = unsafe {
-            set_data_ptr(ptr as *mut ArcInner<T>, (ptr as *mut u8).wrapping_offset(-offset))
+        let ptr = if ptr.is_null() {
+            // If we get a null pointer, this is a dangling weak.
+            // SAFETY: this is the same sentinel as used in Weak::new and is_dangling
+            (ptr as *mut ArcInner<T>).set_ptr_value(usize::MAX as *mut _)
+        } else {
+            // Otherwise, this describes a real allocation.
+            // SAFETY: data_offset is safe to call, as ptr describes a real allocation.
+            let offset = unsafe { data_offset(ptr) };
+            // Thus, we reverse the offset to get the whole RcBox.
+            // SAFETY: the pointer originated from a Weak, so this offset is safe.
+            unsafe { (ptr as *mut ArcInner<T>).set_ptr_value((ptr as *mut u8).offset(-offset)) }
         };
 
         // SAFETY: we now have recovered the original Weak pointer, so can create the Weak.
-        unsafe { Weak { ptr: NonNull::new_unchecked(ptr) } }
+        Weak { ptr: unsafe { NonNull::new_unchecked(ptr) } }
     }
 }
 

diff --git a/library/alloc/src/sync/tests.rs b/library/alloc/src/sync/tests.rs
@@ -158,6 +158,30 @@ fn into_from_weak_raw() {
     }
 }
 
+#[test]
+fn test_into_from_weak_raw_unsized() {
+    use std::fmt::Display;
+    use std::string::ToString;
+
+    let arc: Arc<str> = Arc::from("foo");
+    let weak: Weak<str> = Arc::downgrade(&arc);
+
+    let ptr = Weak::into_raw(weak.clone());
+    let weak2 = unsafe { Weak::from_raw(ptr) };
+
+    assert_eq!(unsafe { &*ptr }, "foo");
+    assert!(weak.ptr_eq(&weak2));
+
+    let arc: Arc<dyn Display> = Arc::new(123);
+    let weak: Weak<dyn Display> = Arc::downgrade(&arc);
+
+    let ptr = Weak::into_raw(weak.clone());
+    let weak2 = unsafe { Weak::from_raw(ptr) };
+
+    assert_eq!(unsafe { &*ptr }.to_string(), "123");
+    assert!(weak.ptr_eq(&weak2));
+}
+
 #[test]
 fn test_cowarc_clone_make_mut() {
     let mut cow0 = Arc::new(75);
@@ -329,6 +353,23 @@ fn test_unsized() {
     assert!(y.upgrade().is_none());
 }
 
+#[test]
+fn test_maybe_thin_unsized() {
+    // If/when custom thin DSTs exist, this test should be updated to use one
+    use std::ffi::{CStr, CString};
+
+    let x: Arc<CStr> = Arc::from(CString::new("swordfish").unwrap().into_boxed_c_str());
+    assert_eq!(format!("{:?}", x), "\"swordfish\"");
+    let y: Weak<CStr> = Arc::downgrade(&x);
+    drop(x);
+
+    // At this point, the weak points to a dropped DST
+    assert!(y.upgrade().is_none());
+    // But we still need to be able to get the alloc layout to drop.
+    // CStr has no drop glue, but custom DSTs might, and need to work.
+    drop(y);
+}
+
 #[test]
 fn test_from_owned() {
     let foo = 123;