Make Weak::weak_count() return zero when no strong refs remain

4 files changed, 31 insertions, 45 deletions

diff --git a/src/liballoc/rc.rs b/src/liballoc/rc.rs
index 12f86bf880f..0dab68b155d 100644
--- a/src/liballoc/rc.rs
+++ b/src/liballoc/rc.rs
@@ -1814,20 +1814,16 @@ impl<T: ?Sized> Weak<T> {
 
     /// Gets the number of `Weak` pointers pointing to this allocation.
     ///
-    /// 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
-    /// allocation.
-    ///
-    /// [`Weak::new`]: #method.new
+    /// If no strong pointers remain, this will return zero.
     #[stable(feature = "weak_counts", since = "1.40.0")]
-    pub fn weak_count(&self) -> Option<usize> {
+    pub fn weak_count(&self) -> usize {
         self.inner().map(|inner| {
             if inner.strong() > 0 {
                 inner.weak() - 1  // subtract the implicit weak ptr
             } else {
-                inner.weak()
+                0
             }
-        })
+        }).unwrap_or(0)
     }
 
     /// Returns `None` when the pointer is dangling and there is no allocated `RcBox`
diff --git a/src/liballoc/rc/tests.rs b/src/liballoc/rc/tests.rs
index 6fd3f909357..bf5c85a5c59 100644
--- a/src/liballoc/rc/tests.rs
+++ b/src/liballoc/rc/tests.rs
@@ -114,28 +114,28 @@ fn test_weak_count() {
 
 #[test]
 fn weak_counts() {
-    assert_eq!(Weak::weak_count(&Weak::<u64>::new()), None);
+    assert_eq!(Weak::weak_count(&Weak::<u64>::new()), 0);
     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));
+    assert_eq!(Weak::weak_count(&w), 1);
     let w2 = w.clone();
     assert_eq!(Weak::strong_count(&w), 1);
-    assert_eq!(Weak::weak_count(&w), Some(2));
+    assert_eq!(Weak::weak_count(&w), 2);
     assert_eq!(Weak::strong_count(&w2), 1);
-    assert_eq!(Weak::weak_count(&w2), Some(2));
+    assert_eq!(Weak::weak_count(&w2), 2);
     drop(w);
     assert_eq!(Weak::strong_count(&w2), 1);
-    assert_eq!(Weak::weak_count(&w2), Some(1));
+    assert_eq!(Weak::weak_count(&w2), 1);
     let a2 = a.clone();
     assert_eq!(Weak::strong_count(&w2), 2);
-    assert_eq!(Weak::weak_count(&w2), Some(1));
+    assert_eq!(Weak::weak_count(&w2), 1);
     drop(a2);
     drop(a);
     assert_eq!(Weak::strong_count(&w2), 0);
-    assert_eq!(Weak::weak_count(&w2), Some(1));
+    assert_eq!(Weak::weak_count(&w2), 0);
     drop(w2);
 }
 
diff --git a/src/liballoc/sync.rs b/src/liballoc/sync.rs
index 3f86dfb469e..1bfe3b80249 100644
--- a/src/liballoc/sync.rs
+++ b/src/liballoc/sync.rs
@@ -12,7 +12,7 @@ use core::sync::atomic;
 use core::sync::atomic::Ordering::{Acquire, Relaxed, Release, SeqCst};
 use core::borrow;
 use core::fmt;
-use core::cmp::{self, Ordering};
+use core::cmp::Ordering;
 use core::iter;
 use core::intrinsics::abort;
 use core::mem::{self, align_of, align_of_val, size_of_val};
@@ -1508,9 +1508,8 @@ impl<T: ?Sized> Weak<T> {
     /// Gets an approximation of the number of `Weak` pointers pointing to this
     /// allocation.
     ///
-    /// 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
-    /// allocation.
+    /// If `self` was created using [`Weak::new`], or if there are no remaining
+    /// strong pointers, this will return 0.
     ///
     /// # Accuracy
     ///
@@ -1520,30 +1519,21 @@ impl<T: ?Sized> Weak<T> {
     ///
     /// [`Weak::new`]: #method.new
     #[stable(feature = "weak_counts", since = "1.40.0")]
-    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.
+    pub fn weak_count(&self) -> usize {
         self.inner().map(|inner| {
-            let strong = inner.strong.load(SeqCst);
             let weak = inner.weak.load(SeqCst);
+            let strong = inner.strong.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
+                0
             } 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)
+                // Since we observed that there was at least one strong pointer
+                // after reading the weak count, we know that the implicit weak
+                // reference (present whenever any strong references are alive)
+                // was still around when we observed the weak count, and can
+                // therefore safely subtract it.
+                weak - 1
             }
-        })
+        }).unwrap_or(0)
     }
 
     /// Returns `None` when the pointer is dangling and there is no allocated `ArcInner`,
diff --git a/src/liballoc/sync/tests.rs b/src/liballoc/sync/tests.rs
index 9220f5e0333..be0200c9a46 100644
--- a/src/liballoc/sync/tests.rs
+++ b/src/liballoc/sync/tests.rs
@@ -62,28 +62,28 @@ fn test_arc_get_mut() {
 
 #[test]
 fn weak_counts() {
-    assert_eq!(Weak::weak_count(&Weak::<u64>::new()), None);
+    assert_eq!(Weak::weak_count(&Weak::<u64>::new()), 0);
     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));
+    assert_eq!(Weak::weak_count(&w), 1);
     let w2 = w.clone();
     assert_eq!(Weak::strong_count(&w), 1);
-    assert_eq!(Weak::weak_count(&w), Some(2));
+    assert_eq!(Weak::weak_count(&w), 2);
     assert_eq!(Weak::strong_count(&w2), 1);
-    assert_eq!(Weak::weak_count(&w2), Some(2));
+    assert_eq!(Weak::weak_count(&w2), 2);
     drop(w);
     assert_eq!(Weak::strong_count(&w2), 1);
-    assert_eq!(Weak::weak_count(&w2), Some(1));
+    assert_eq!(Weak::weak_count(&w2), 1);
     let a2 = a.clone();
     assert_eq!(Weak::strong_count(&w2), 2);
-    assert_eq!(Weak::weak_count(&w2), Some(1));
+    assert_eq!(Weak::weak_count(&w2), 1);
     drop(a2);
     drop(a);
     assert_eq!(Weak::strong_count(&w2), 0);
-    assert_eq!(Weak::weak_count(&w2), Some(1));
+    assert_eq!(Weak::weak_count(&w2), 0);
     drop(w2);
 }