Rollup merge of #75994 - mental32:impl-rc-new-cyclic, r=KodrAus

`impl Rc::new_cyclic`

References #75861

r? @Dylan-DPC

2 files changed, 110 insertions, 0 deletions

diff --git a/library/alloc/src/rc.rs b/library/alloc/src/rc.rs
index 2f0a374015a..a9b293856e5 100644
--- a/library/alloc/src/rc.rs
+++ b/library/alloc/src/rc.rs
@@ -325,6 +325,50 @@ impl<T> Rc<T> {
         )
     }
 
+    /// Constructs a new `Rc<T>` using a weak reference to itself. Attempting
+    /// to upgrade the weak reference before this function returns will result
+    /// in a `None` value. However, the weak reference may be cloned freely and
+    /// stored for use at a later time.
+    #[unstable(feature = "arc_new_cyclic", issue = "75861")]
+    pub fn new_cyclic(data_fn: impl FnOnce(&Weak<T>) -> T) -> Rc<T> {
+        // Construct the inner in the "uninitialized" state with a single
+        // weak reference.
+        let uninit_ptr: NonNull<_> = Box::leak(box RcBox {
+            strong: Cell::new(0),
+            weak: Cell::new(1),
+            value: mem::MaybeUninit::<T>::uninit(),
+        })
+        .into();
+
+        let init_ptr: NonNull<RcBox<T>> = uninit_ptr.cast();
+
+        let weak = Weak { ptr: init_ptr };
+
+        // It's important we don't give up ownership of the weak pointer, or
+        // else the memory might be freed by the time `data_fn` returns. If
+        // we really wanted to pass ownership, we could create an additional
+        // weak pointer for ourselves, but this would result in additional
+        // updates to the weak reference count which might not be necessary
+        // otherwise.
+        let data = data_fn(&weak);
+
+        unsafe {
+            let inner = init_ptr.as_ptr();
+            ptr::write(&raw mut (*inner).value, data);
+
+            let prev_value = (*inner).strong.get();
+            debug_assert_eq!(prev_value, 0, "No prior strong references should exist");
+            (*inner).strong.set(1);
+        }
+
+        let strong = Rc::from_inner(init_ptr);
+
+        // Strong references should collectively own a shared weak reference,
+        // so don't run the destructor for our old weak reference.
+        mem::forget(weak);
+        strong
+    }
+
     /// Constructs a new `Rc` with uninitialized contents.
     ///
     /// # Examples
diff --git a/library/alloc/src/rc/tests.rs b/library/alloc/src/rc/tests.rs
index e88385faf4f..fed48a59f80 100644
--- a/library/alloc/src/rc/tests.rs
+++ b/library/alloc/src/rc/tests.rs
@@ -434,3 +434,69 @@ fn test_array_from_slice() {
     let a: Result<Rc<[u32; 2]>, _> = r.clone().try_into();
     assert!(a.is_err());
 }
+
+#[test]
+fn test_rc_cyclic_with_zero_refs() {
+    struct ZeroRefs {
+        inner: Weak<ZeroRefs>,
+    }
+
+    let zero_refs = Rc::new_cyclic(|inner| {
+        assert_eq!(inner.strong_count(), 0);
+        assert!(inner.upgrade().is_none());
+        ZeroRefs { inner: Weak::new() }
+    });
+
+    assert_eq!(Rc::strong_count(&zero_refs), 1);
+    assert_eq!(Rc::weak_count(&zero_refs), 0);
+    assert_eq!(zero_refs.inner.strong_count(), 0);
+    assert_eq!(zero_refs.inner.weak_count(), 0);
+}
+
+#[test]
+fn test_rc_cyclic_with_one_ref() {
+    struct OneRef {
+        inner: Weak<OneRef>,
+    }
+
+    let one_ref = Rc::new_cyclic(|inner| {
+        assert_eq!(inner.strong_count(), 0);
+        assert!(inner.upgrade().is_none());
+        OneRef { inner: inner.clone() }
+    });
+
+    assert_eq!(Rc::strong_count(&one_ref), 1);
+    assert_eq!(Rc::weak_count(&one_ref), 1);
+
+    let one_ref2 = Weak::upgrade(&one_ref.inner).unwrap();
+    assert!(Rc::ptr_eq(&one_ref, &one_ref2));
+
+    assert_eq!(one_ref.inner.strong_count(), 2);
+    assert_eq!(one_ref.inner.weak_count(), 1);
+}
+
+#[test]
+fn test_rc_cyclic_with_two_ref() {
+    struct TwoRefs {
+        inner: Weak<TwoRefs>,
+        inner1: Weak<TwoRefs>,
+    }
+
+    let two_refs = Rc::new_cyclic(|inner| {
+        assert_eq!(inner.strong_count(), 0);
+        assert!(inner.upgrade().is_none());
+        TwoRefs { inner: inner.clone(), inner1: inner.clone() }
+    });
+
+    assert_eq!(Rc::strong_count(&two_refs), 1);
+    assert_eq!(Rc::weak_count(&two_refs), 2);
+
+    let two_ref3 = Weak::upgrade(&two_refs.inner).unwrap();
+    assert!(Rc::ptr_eq(&two_refs, &two_ref3));
+
+    let two_ref2 = Weak::upgrade(&two_refs.inner1).unwrap();
+    assert!(Rc::ptr_eq(&two_refs, &two_ref2));
+
+    assert_eq!(Rc::strong_count(&two_refs), 3);
+    assert_eq!(Rc::weak_count(&two_refs), 2);
+}