Auto merge of #56696 - jonas-schievink:weak-counts, r=alexcrichton

Implement Weak::{strong_count, weak_count}

The counters are also useful on `Weak`, not just on strong references (`Rc` or `Arc`).

In situations where there are still strong references around, you can also get these counts by temporarily upgrading and adjusting the values accordingly. Using the methods introduced here is simpler to do, less error-prone (since you can't forget to adjust the counts), can also be used when no strong references are around anymore, and might be more efficient due to not having to temporarily create an `Rc`.

This is mainly useful in assertions or tests of complex data structures. Data structures might have internal invariants that make them the sole owner of a `Weak` pointer, and an assertion on the weak count could be used to ensure that this indeed happens as expected. Due to the presence of `Weak::upgrade`, the `strong_count` becomes less useful, but it still seems worthwhile to mirror the API of `Rc`.

TODO:
* [X] Tracking issue - #57977

Closes #50158

src/liballoc/rc.rs

@@ -1305,6 +1305,38 @@ impl<T: ?Sized> Weak<T> {
         }
     }
 
+    /// Gets the number of strong (`Rc`) pointers pointing to this value.
+    ///
+    /// If `self` was created using [`Weak::new`], this will return 0.
+    ///
+    /// [`Weak::new`]: #method.new
+    #[unstable(feature = "weak_counts", issue = "57977")]
+    pub fn strong_count(&self) -> usize {
+        if let Some(inner) = self.inner() {
+            inner.strong()
+        } else {
+            0
+        }
+    }
+
+    /// Gets the number of `Weak` pointers pointing to this value.
+    ///
+    /// If `self` was created using [`Weak::new`], this will return `None`. If
+    /// not, the returned value is at least 1, since `self` still points to the
+    /// value.
+    ///
+    /// [`Weak::new`]: #method.new
+    #[unstable(feature = "weak_counts", issue = "57977")]
+    pub fn weak_count(&self) -> Option<usize> {
+        self.inner().map(|inner| {
+            if inner.strong() > 0 {
+                inner.weak() - 1  // subtract the implicit weak ptr
+            } else {
+                inner.weak()
+            }
+        })
+    }
+
     /// Return `None` when the pointer is dangling and there is no allocated `RcBox`,
     /// i.e., this `Weak` was created by `Weak::new`
     #[inline]
@@ -1643,6 +1675,33 @@ mod tests {
         drop(c);
     }
 
+    #[test]
+    fn weak_counts() {
+        assert_eq!(Weak::weak_count(&Weak::<u64>::new()), None);
+        assert_eq!(Weak::strong_count(&Weak::<u64>::new()), 0);
+
+        let a = Rc::new(0);
+        let w = Rc::downgrade(&a);
+        assert_eq!(Weak::strong_count(&w), 1);
+        assert_eq!(Weak::weak_count(&w), Some(1));
+        let w2 = w.clone();
+        assert_eq!(Weak::strong_count(&w), 1);
+        assert_eq!(Weak::weak_count(&w), Some(2));
+        assert_eq!(Weak::strong_count(&w2), 1);
+        assert_eq!(Weak::weak_count(&w2), Some(2));
+        drop(w);
+        assert_eq!(Weak::strong_count(&w2), 1);
+        assert_eq!(Weak::weak_count(&w2), Some(1));
+        let a2 = a.clone();
+        assert_eq!(Weak::strong_count(&w2), 2);
+        assert_eq!(Weak::weak_count(&w2), Some(1));
+        drop(a2);
+        drop(a);
+        assert_eq!(Weak::strong_count(&w2), 0);
+        assert_eq!(Weak::weak_count(&w2), Some(1));
+        drop(w2);
+    }
+
     #[test]
     fn try_unwrap() {
         let x = Rc::new(3);

src/liballoc/sync.rs

@@ -11,7 +11,7 @@ use core::sync::atomic;
 use core::sync::atomic::Ordering::{Acquire, Relaxed, Release, SeqCst};
 use core::borrow;
 use core::fmt;
-use core::cmp::Ordering;
+use core::cmp::{self, Ordering};
 use core::intrinsics::abort;
 use core::mem::{self, align_of_val, size_of_val};
 use core::ops::{Deref, Receiver};
@@ -1138,6 +1138,61 @@ impl<T: ?Sized> Weak<T> {
         }
     }
 
+    /// Gets the number of strong (`Arc`) pointers pointing to this value.
+    ///
+    /// If `self` was created using [`Weak::new`], this will return 0.
+    ///
+    /// [`Weak::new`]: #method.new
+    #[unstable(feature = "weak_counts", issue = "57977")]
+    pub fn strong_count(&self) -> usize {
+        if let Some(inner) = self.inner() {
+            inner.strong.load(SeqCst)
+        } else {
+            0
+        }
+    }
+
+    /// Gets an approximation of the number of `Weak` pointers pointing to this
+    /// value.
+    ///
+    /// If `self` was created using [`Weak::new`], this will return 0. If not,
+    /// the returned value is at least 1, since `self` still points to the
+    /// value.
+    ///
+    /// # Accuracy
+    ///
+    /// Due to implementation details, the returned value can be off by 1 in
+    /// either direction when other threads are manipulating any `Arc`s or
+    /// `Weak`s pointing to the same value.
+    ///
+    /// [`Weak::new`]: #method.new
+    #[unstable(feature = "weak_counts", issue = "57977")]
+    pub fn weak_count(&self) -> Option<usize> {
+        // Due to the implicit weak pointer added when any strong pointers are
+        // around, we cannot implement `weak_count` correctly since it
+        // necessarily requires accessing the strong count and weak count in an
+        // unsynchronized fashion. So this version is a bit racy.
+        self.inner().map(|inner| {
+            let strong = inner.strong.load(SeqCst);
+            let weak = inner.weak.load(SeqCst);
+            if strong == 0 {
+                // If the last `Arc` has *just* been dropped, it might not yet
+                // have removed the implicit weak count, so the value we get
+                // here might be 1 too high.
+                weak
+            } else {
+                // As long as there's still at least 1 `Arc` around, subtract
+                // the implicit weak pointer.
+                // Note that the last `Arc` might get dropped between the 2
+                // loads we do above, removing the implicit weak pointer. This
+                // means that the value might be 1 too low here. In order to not
+                // return 0 here (which would happen if we're the only weak
+                // pointer), we guard against that specifically.
+                cmp::max(1, weak - 1)
+            }
+        })
+    }
+
     /// Return `None` when the pointer is dangling and there is no allocated `ArcInner`,
     /// i.e., this `Weak` was created by `Weak::new`
     #[inline]
@@ -1657,6 +1712,33 @@ mod tests {
         assert!(Arc::get_mut(&mut x).is_none());
     }
 
+    #[test]
+    fn weak_counts() {
+        assert_eq!(Weak::weak_count(&Weak::<u64>::new()), None);
+        assert_eq!(Weak::strong_count(&Weak::<u64>::new()), 0);
+
+        let a = Arc::new(0);
+        let w = Arc::downgrade(&a);
+        assert_eq!(Weak::strong_count(&w), 1);
+        assert_eq!(Weak::weak_count(&w), Some(1));
+        let w2 = w.clone();
+        assert_eq!(Weak::strong_count(&w), 1);
+        assert_eq!(Weak::weak_count(&w), Some(2));
+        assert_eq!(Weak::strong_count(&w2), 1);
+        assert_eq!(Weak::weak_count(&w2), Some(2));
+        drop(w);
+        assert_eq!(Weak::strong_count(&w2), 1);
+        assert_eq!(Weak::weak_count(&w2), Some(1));
+        let a2 = a.clone();
+        assert_eq!(Weak::strong_count(&w2), 2);
+        assert_eq!(Weak::weak_count(&w2), Some(1));
+        drop(a2);
+        drop(a);
+        assert_eq!(Weak::strong_count(&w2), 0);
+        assert_eq!(Weak::weak_count(&w2), Some(1));
+        drop(w2);
+    }
+
     #[test]
     fn try_unwrap() {
         let x = Arc::new(3);