diff --git a/library/alloc/src/rc.rs b/library/alloc/src/rc.rs index 2f0a374015af6..a9b293856e57b 100644 --- a/library/alloc/src/rc.rs +++ b/library/alloc/src/rc.rs @@ -325,6 +325,50 @@ impl Rc { ) } + /// Constructs a new `Rc` 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) -> Rc { + // 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::::uninit(), + }) + .into(); + + let init_ptr: NonNull> = 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 e88385faf4fd4..fed48a59f809e 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, _> = r.clone().try_into(); assert!(a.is_err()); } + +#[test] +fn test_rc_cyclic_with_zero_refs() { + struct ZeroRefs { + inner: Weak, + } + + 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, + } + + 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, + inner1: Weak, + } + + 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); +}