Make Arc::make_unique check weak refs; add make_unique to Rc

src/liballoc/arc.rs

@@ -152,7 +152,11 @@ impl<T: Send + Share + Clone> Arc<T> {
     #[inline]
     #[experimental]
     pub fn make_unique<'a>(&'a mut self) -> &'a mut T {
-        if self.inner().strong.load(atomics::SeqCst) != 1 {
+        // Note that we hold a strong reference, which also counts as
+        // a weak reference, so we only clone if there is an
+        // additional reference of either kind.
+        if self.inner().strong.load(atomics::SeqCst) != 1 ||
+           self.inner().weak.load(atomics::SeqCst) != 1 {
             *self = Arc::new(self.deref().clone())
         }
         // This unsafety is ok because we're guaranteed that the pointer
@@ -356,6 +360,20 @@ mod tests {
         assert!(*cow1 == *cow2);
     }
 
+    #[test]
+    fn test_cowarc_clone_weak() {
+        let mut cow0 = Arc::new(75u);
+        let cow1_weak = cow0.downgrade();
+
+        assert!(75 == *cow0);
+        assert!(75 == *cow1_weak.upgrade().unwrap());
+
+        *cow0.make_unique() += 1;
+
+        assert!(76 == *cow0);
+        assert!(cow1_weak.upgrade().is_none());
+    }
+
     #[test]
     fn test_live() {
         let x = Arc::new(5);

src/liballoc/rc.rs

@@ -86,6 +86,31 @@ impl<T> Rc<T> {
     }
 }
 
+impl<T: Clone> Rc<T> {
+    /// Acquires a mutable pointer to the inner contents by guaranteeing that
+    /// the reference count is one (no sharing is possible).
+    ///
+    /// This is also referred to as a copy-on-write operation because the inner
+    /// data is cloned if the reference count is greater than one.
+    #[inline]
+    #[experimental]
+    pub fn make_unique<'a>(&'a mut self) -> &'a mut T {
+        // Note that we hold a strong reference, which also counts as
+        // a weak reference, so we only clone if there is an
+        // additional reference of either kind.
+        if self.strong() != 1 || self.weak() != 1 {
+            *self = Rc::new(self.deref().clone())
+        }
+        // This unsafety is ok because we're guaranteed that the pointer
+        // returned is the *only* pointer that will ever be returned to T. Our
+        // reference count is guaranteed to be 1 at this point, and we required
+        // the Rc itself to be `mut`, so we're returning the only possible
+        // reference to the inner data.
+        let inner = unsafe { &mut *self._ptr };
+        &mut inner.value
+    }
+}
+
 impl<T> Deref<T> for Rc<T> {
     /// Borrow the value contained in the reference-counted box
     #[inline(always)]
@@ -234,6 +259,7 @@ impl<T> RcBoxPtr<T> for Weak<T> {
 }
 
 #[cfg(test)]
+#[allow(experimental)]
 mod tests {
     use super::{Rc, Weak};
     use std::cell::RefCell;
@@ -304,4 +330,66 @@ mod tests {
 
         // hopefully we don't double-free (or leak)...
     }
+
+    #[test]
+    fn test_cowrc_clone_make_unique() {
+        let mut cow0 = Rc::new(75u);
+        let mut cow1 = cow0.clone();
+        let mut cow2 = cow1.clone();
+
+        assert!(75 == *cow0.make_unique());
+        assert!(75 == *cow1.make_unique());
+        assert!(75 == *cow2.make_unique());
+
+        *cow0.make_unique() += 1;
+        *cow1.make_unique() += 2;
+        *cow2.make_unique() += 3;
+
+        assert!(76 == *cow0);
+        assert!(77 == *cow1);
+        assert!(78 == *cow2);
+
+        // none should point to the same backing memory
+        assert!(*cow0 != *cow1);
+        assert!(*cow0 != *cow2);
+        assert!(*cow1 != *cow2);
+    }
+
+    #[test]
+    fn test_cowrc_clone_unique2() {
+        let mut cow0 = Rc::new(75u);
+        let cow1 = cow0.clone();
+        let cow2 = cow1.clone();
+
+        assert!(75 == *cow0);
+        assert!(75 == *cow1);
+        assert!(75 == *cow2);
+
+        *cow0.make_unique() += 1;
+
+        assert!(76 == *cow0);
+        assert!(75 == *cow1);
+        assert!(75 == *cow2);
+
+        // cow1 and cow2 should share the same contents
+        // cow0 should have a unique reference
+        assert!(*cow0 != *cow1);
+        assert!(*cow0 != *cow2);
+        assert!(*cow1 == *cow2);
+    }
+
+    #[test]
+    fn test_cowrc_clone_weak() {
+        let mut cow0 = Rc::new(75u);
+        let cow1_weak = cow0.downgrade();
+
+        assert!(75 == *cow0);
+        assert!(75 == *cow1_weak.upgrade().unwrap());
+
+        *cow0.make_unique() += 1;
+
+        assert!(76 == *cow0);
+        assert!(cow1_weak.upgrade().is_none());
+    }
+
 }