From f03d0b38d6a33a64307d83f8ddd3df8ef57ca537 Mon Sep 17 00:00:00 2001
From: mental32 <m3nta1@yahoo.com>
Date: Thu, 27 Aug 2020 19:19:29 +0100
Subject: [PATCH] `impl Rc::new_cyclic`

---
 library/alloc/src/rc.rs       | 34 ++++++++++++++++++
 library/alloc/src/rc/tests.rs | 66 +++++++++++++++++++++++++++++++++++
 2 files changed, 100 insertions(+)

diff --git a/library/alloc/src/rc.rs b/library/alloc/src/rc.rs
index 1046397f4be..76266d77bb0 100644
--- a/library/alloc/src/rc.rs
+++ b/library/alloc/src/rc.rs
@@ -325,6 +325,40 @@ impl<T> Rc<T> {
         )
     }
 
+    /// Constructs a new `Rc<T>` using a weak reference to itself. Attempting
+    /// to upgrade the weak reference before this function retuns will result
+    /// in a `None` value. However, the weak reference may be cloned freely and
+    /// stored for use at a later time.
+    #[inline]
+    #[unstable(feature = "arc_new_cyclic", issue = "75861")]
+    pub fn new_cyclic(data_fn: impl FnOnce(&Weak<T>) -> T) -> Rc<T> {
+        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 };
+
+        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);
+        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);
+}