rollup merge of #19193: scialex/rc-counts

These functions allow you to see how many weak and strong references
there are to an `Arc`, `Rc`, or an `rc::Weak`. Due to the design of
`Arc` it is not possible to get the number of weak references of an
arbitrary `arc::Weak`. Look in `arc.rs` for a more in-depth explanation.

On `arc::Arc` and `arc::Weak` these operations are wait-free and atomic.

This sort of information is useful for creating dynamically cleared caches for use in OS development, for example holding pages of files in memory until the address space is needed for something else.

2 files changed, 100 insertions, 4 deletions

diff --git a/src/liballoc/arc.rs b/src/liballoc/arc.rs
index c4f53d74467..4f744b0b2de 100644
--- a/src/liballoc/arc.rs
+++ b/src/liballoc/arc.rs
@@ -119,6 +119,16 @@ impl<T> Arc<T> {
     }
 }
 
+/// Get the number of weak references to this value.
+#[inline]
+#[experimental]
+pub fn weak_count<T>(this: &Arc<T>) -> uint { this.inner().weak.load(atomic::SeqCst) - 1 }
+
+/// Get the number of strong references to this value.
+#[inline]
+#[experimental]
+pub fn strong_count<T>(this: &Arc<T>) -> uint { this.inner().strong.load(atomic::SeqCst) }
+
 #[unstable = "waiting on stability of Clone"]
 impl<T> Clone for Arc<T> {
     /// Duplicate an atomically reference counted wrapper.
@@ -321,7 +331,7 @@ mod tests {
     use std::sync::atomic;
     use std::task;
     use std::vec::Vec;
-    use super::{Arc, Weak};
+    use super::{Arc, Weak, weak_count, strong_count};
     use std::sync::Mutex;
 
     struct Canary(*mut atomic::AtomicUint);
@@ -466,6 +476,49 @@ mod tests {
     }
 
     #[test]
+    fn test_strong_count() {
+        let a = Arc::new(0u32);
+        assert!(strong_count(&a) == 1);
+        let w = a.downgrade();
+        assert!(strong_count(&a) == 1);
+        let b = w.upgrade().expect("");
+        assert!(strong_count(&b) == 2);
+        assert!(strong_count(&a) == 2);
+        drop(w);
+        drop(a);
+        assert!(strong_count(&b) == 1);
+        let c = b.clone();
+        assert!(strong_count(&b) == 2);
+        assert!(strong_count(&c) == 2);
+    }
+
+    #[test]
+    fn test_weak_count() {
+        let a = Arc::new(0u32);
+        assert!(strong_count(&a) == 1);
+        assert!(weak_count(&a) == 0);
+        let w = a.downgrade();
+        assert!(strong_count(&a) == 1);
+        assert!(weak_count(&a) == 1);
+        let x = w.clone();
+        assert!(weak_count(&a) == 2);
+        drop(w);
+        drop(x);
+        assert!(strong_count(&a) == 1);
+        assert!(weak_count(&a) == 0);
+        let c = a.clone();
+        assert!(strong_count(&a) == 2);
+        assert!(weak_count(&a) == 0);
+        let d = c.downgrade();
+        assert!(weak_count(&c) == 1);
+        assert!(strong_count(&c) == 2);
+
+        drop(a);
+        drop(c);
+        drop(d);
+    }
+
+    #[test]
     fn show_arc() {
         let a = Arc::new(5u32);
         assert!(format!("{}", a).as_slice() == "5")
diff --git a/src/liballoc/rc.rs b/src/liballoc/rc.rs
index 501f915461a..df84ac9aec9 100644
--- a/src/liballoc/rc.rs
+++ b/src/liballoc/rc.rs
@@ -213,6 +213,16 @@ impl<T> Rc<T> {
     }
 }
 
+/// Get the number of weak references to this value.
+#[inline]
+#[experimental]
+pub fn weak_count<T>(this: &Rc<T>) -> uint { this.weak() - 1 }
+
+/// Get the number of strong references to this value.
+#[inline]
+#[experimental]
+pub fn strong_count<T>(this: &Rc<T>) -> uint { this.strong() }
+
 /// Returns true if the `Rc` currently has unique ownership.
 ///
 /// Unique ownership means that there are no other `Rc` or `Weak` values
@@ -220,8 +230,7 @@ impl<T> Rc<T> {
 #[inline]
 #[experimental]
 pub fn is_unique<T>(rc: &Rc<T>) -> bool {
-    // note that we hold both a strong and a weak reference
-    rc.strong() == 1 && rc.weak() == 1
+    weak_count(rc) == 0 && strong_count(rc) == 1
 }
 
 /// Unwraps the contained value if the `Rc` has unique ownership.
@@ -489,7 +498,7 @@ impl<T> RcBoxPtr<T> for Weak<T> {
 #[cfg(test)]
 #[allow(experimental)]
 mod tests {
-    use super::{Rc, Weak};
+    use super::{Rc, Weak, weak_count, strong_count};
     use std::cell::RefCell;
     use std::option::{Option, Some, None};
     use std::result::{Err, Ok};
@@ -567,6 +576,40 @@ mod tests {
     }
 
     #[test]
+    fn test_strong_count() {
+        let a = Rc::new(0u32);
+        assert!(strong_count(&a) == 1);
+        let w = a.downgrade();
+        assert!(strong_count(&a) == 1);
+        let b = w.upgrade().expect("upgrade of live rc failed");
+        assert!(strong_count(&b) == 2);
+        assert!(strong_count(&a) == 2);
+        drop(w);
+        drop(a);
+        assert!(strong_count(&b) == 1);
+        let c = b.clone();
+        assert!(strong_count(&b) == 2);
+        assert!(strong_count(&c) == 2);
+    }
+
+    #[test]
+    fn test_weak_count() {
+        let a = Rc::new(0u32);
+        assert!(strong_count(&a) == 1);
+        assert!(weak_count(&a) == 0);
+        let w = a.downgrade();
+        assert!(strong_count(&a) == 1);
+        assert!(weak_count(&a) == 1);
+        drop(w);
+        assert!(strong_count(&a) == 1);
+        assert!(weak_count(&a) == 0);
+        let c = a.clone();
+        assert!(strong_count(&a) == 2);
+        assert!(weak_count(&a) == 0);
+        drop(c);
+    }
+
+    #[test]
     fn try_unwrap() {
         let x = Rc::new(3u);
         assert_eq!(super::try_unwrap(x), Ok(3u));