Add test helper RcTestObject to test that reference counting works properly

objc2/src/rc/id.rs

@@ -606,49 +606,93 @@ impl<T: UnwindSafe + ?Sized> UnwindSafe for Id<T, Owned> {}
 
 #[cfg(test)]
 mod tests {
-    use super::{Id, Owned, Shared};
-    use crate::rc::autoreleasepool;
+    use super::*;
+    use crate::msg_send;
+    use crate::rc::{autoreleasepool, RcTestObject, ThreadTestData};
     use crate::runtime::Object;
-    use crate::{class, msg_send};
 
-    fn retain_count(obj: &Object) -> usize {
-        unsafe { msg_send![obj, retainCount] }
+    #[track_caller]
+    fn assert_retain_count(obj: &Object, expected: usize) {
+        let retain_count: usize = unsafe { msg_send![obj, retainCount] };
+        assert_eq!(retain_count, expected);
     }
 
     #[test]
-    fn test_autorelease() {
-        let obj: Id<Object, Shared> = unsafe { Id::new(msg_send![class!(NSObject), new]).unwrap() };
+    fn test_drop() {
+        let mut expected = ThreadTestData::current();
+
+        let obj = RcTestObject::new();
+        expected.alloc += 1;
+        expected.init += 1;
+        expected.assert_current();
+
+        drop(obj);
+        expected.release += 1;
+        expected.dealloc += 1;
+        expected.assert_current();
+    }
 
+    #[test]
+    fn test_autorelease() {
+        let obj: Id<_, Shared> = RcTestObject::new().into();
         let cloned = obj.clone();
+        let mut expected = ThreadTestData::current();
 
         autoreleasepool(|pool| {
             let _ref = obj.autorelease(pool);
-            assert_eq!(retain_count(&*cloned), 2);
+            expected.autorelease += 1;
+            expected.assert_current();
+            assert_retain_count(&cloned, 2);
         });
+        expected.release += 1;
+        expected.assert_current();
+        assert_retain_count(&cloned, 1);
 
-        // make sure that the autoreleased value has been released
-        // TODO: Investigate if this is flaky on GNUStep
-        assert_eq!(retain_count(&*cloned), 1);
+        autoreleasepool(|pool| {
+            let _ref = cloned.autorelease(pool);
+            expected.autorelease += 1;
+            expected.assert_current();
+        });
+        expected.release += 1;
+        expected.dealloc += 1;
+        expected.assert_current();
     }
 
     #[test]
     fn test_clone() {
-        let cls = class!(NSObject);
-        let obj: Id<Object, Owned> = unsafe {
-            let obj: *mut Object = msg_send![cls, alloc];
-            let obj: *mut Object = msg_send![obj, init];
-            Id::new(obj).unwrap()
-        };
-        assert_eq!(retain_count(&obj), 1);
+        let obj: Id<_, Owned> = RcTestObject::new();
+        assert_retain_count(&obj, 1);
+        let mut expected = ThreadTestData::current();
 
         let obj: Id<_, Shared> = obj.into();
-        assert_eq!(retain_count(&obj), 1);
+        expected.assert_current();
+        assert_retain_count(&obj, 1);
 
         let cloned = obj.clone();
-        assert_eq!(retain_count(&cloned), 2);
-        assert_eq!(retain_count(&obj), 2);
+        expected.retain += 1;
+        expected.assert_current();
+        assert_retain_count(&cloned, 2);
+        assert_retain_count(&obj, 2);
 
         drop(obj);
-        assert_eq!(retain_count(&cloned), 1);
+        expected.release += 1;
+        expected.assert_current();
+        assert_retain_count(&cloned, 1);
+
+        drop(cloned);
+        expected.release += 1;
+        expected.dealloc += 1;
+        expected.assert_current();
+    }
+
+    #[test]
+    fn test_retain_autoreleased_works_as_retain() {
+        let obj: Id<_, Shared> = RcTestObject::new().into();
+        let mut expected = ThreadTestData::current();
+
+        let ptr = Id::as_ptr(&obj) as *mut RcTestObject;
+        let _obj2: Id<_, Shared> = unsafe { Id::retain_autoreleased(ptr) }.unwrap();
+        expected.retain += 1;
+        expected.assert_current();
     }
 }

objc2/src/rc/mod.rs

@@ -62,12 +62,18 @@ mod id_traits;
 mod ownership;
 mod weak_id;
 
+#[cfg(test)]
+mod test_object;
+
 pub use self::autorelease::{autoreleasepool, AutoreleasePool, AutoreleaseSafe};
 pub use self::id::Id;
 pub use self::id_traits::{DefaultId, SliceId, SliceIdMut};
 pub use self::ownership::{Owned, Ownership, Shared};
 pub use self::weak_id::WeakId;
 
+#[cfg(test)]
+pub(crate) use self::test_object::{RcTestObject, ThreadTestData};
+
 #[cfg(test)]
 mod tests {
     use core::marker::PhantomData;

objc2/src/rc/test_object.rs

@@ -0,0 +1,157 @@
+use core::cell::RefCell;
+use core::ops::{Deref, DerefMut};
+use std::sync::Once;
+
+use super::{Id, Owned};
+use crate::declare::ClassBuilder;
+use crate::runtime::{Bool, Class, Object, Sel};
+use crate::{msg_send, msg_send_bool};
+use crate::{Encoding, Message, RefEncode};
+
+#[derive(Debug, Clone, Default, PartialEq)]
+pub(crate) struct ThreadTestData {
+    pub(crate) alloc: usize,
+    pub(crate) dealloc: usize,
+    pub(crate) init: usize,
+    pub(crate) retain: usize,
+    pub(crate) release: usize,
+    pub(crate) autorelease: usize,
+    pub(crate) try_retain: usize,
+    pub(crate) try_retain_fail: usize,
+}
+
+impl ThreadTestData {
+    /// Get the amount of method calls performed on the current thread.
+    pub(crate) fn current() -> ThreadTestData {
+        TEST_DATA.with(|data| data.borrow().clone())
+    }
+
+    #[track_caller]
+    pub(crate) fn assert_current(&self) {
+        let current = Self::current();
+        let mut expected = self.clone();
+        if cfg!(feature = "gnustep-1-7") {
+            // GNUStep doesn't have `tryRetain`, it uses `retain` directly
+            let retain_diff = expected.try_retain - current.try_retain;
+            expected.retain += retain_diff;
+            expected.try_retain -= retain_diff;
+
+            // GNUStep doesn't call `autorelease` if it's overridden
+            expected.autorelease = 0;
+        }
+        assert_eq!(current, expected);
+    }
+}
+
+std::thread_local! {
+    pub(crate) static TEST_DATA: RefCell<ThreadTestData> = RefCell::new(Default::default());
+}
+
+/// A helper object that counts how many times various reference-counting
+/// primitives are called.
+#[repr(C)]
+pub(crate) struct RcTestObject {
+    inner: Object,
+}
+
+unsafe impl RefEncode for RcTestObject {
+    const ENCODING_REF: Encoding<'static> = Object::ENCODING_REF;
+}
+
+unsafe impl Message for RcTestObject {}
+
+unsafe impl Send for RcTestObject {}
+unsafe impl Sync for RcTestObject {}
+
+impl Deref for RcTestObject {
+    type Target = Object;
+    fn deref(&self) -> &Self::Target {
+        &self.inner
+    }
+}
+
+impl DerefMut for RcTestObject {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.inner
+    }
+}
+
+impl RcTestObject {
+    fn class() -> &'static Class {
+        static REGISTER_CLASS: Once = Once::new();
+
+        REGISTER_CLASS.call_once(|| {
+            extern "C" fn alloc(cls: &Class, _cmd: Sel) -> *mut RcTestObject {
+                TEST_DATA.with(|data| data.borrow_mut().alloc += 1);
+                let superclass = class!(NSObject).metaclass();
+                unsafe { msg_send![super(cls, superclass), alloc] }
+            }
+            extern "C" fn init(this: &mut RcTestObject, _cmd: Sel) -> *mut RcTestObject {
+                TEST_DATA.with(|data| data.borrow_mut().init += 1);
+                unsafe { msg_send![super(this, class!(NSObject)), init] }
+            }
+            extern "C" fn retain(this: &RcTestObject, _cmd: Sel) -> *mut RcTestObject {
+                TEST_DATA.with(|data| data.borrow_mut().retain += 1);
+                unsafe { msg_send![super(this, class!(NSObject)), retain] }
+            }
+            extern "C" fn release(this: &RcTestObject, _cmd: Sel) {
+                TEST_DATA.with(|data| data.borrow_mut().release += 1);
+                unsafe { msg_send![super(this, class!(NSObject)), release] }
+            }
+            extern "C" fn autorelease(this: &RcTestObject, _cmd: Sel) -> *mut RcTestObject {
+                TEST_DATA.with(|data| data.borrow_mut().autorelease += 1);
+                unsafe { msg_send![super(this, class!(NSObject)), autorelease] }
+            }
+            unsafe extern "C" fn dealloc(_this: *mut RcTestObject, _cmd: Sel) {
+                TEST_DATA.with(|data| data.borrow_mut().dealloc += 1);
+                // Don't call superclass
+            }
+            unsafe extern "C" fn try_retain(this: &RcTestObject, _cmd: Sel) -> Bool {
+                TEST_DATA.with(|data| data.borrow_mut().try_retain += 1);
+                let res = unsafe { msg_send_bool![super(this, class!(NSObject)), _tryRetain] };
+                if !res {
+                    TEST_DATA.with(|data| data.borrow_mut().try_retain -= 1);
+                    TEST_DATA.with(|data| data.borrow_mut().try_retain_fail += 1);
+                }
+                Bool::from(res)
+            }
+
+            let mut builder = ClassBuilder::new("RcTestObject", class!(NSObject)).unwrap();
+            unsafe {
+                builder.add_class_method(
+                    sel!(alloc),
+                    alloc as extern "C" fn(&Class, Sel) -> *mut RcTestObject,
+                );
+                builder.add_method(
+                    sel!(init),
+                    init as extern "C" fn(&mut RcTestObject, Sel) -> _,
+                );
+                builder.add_method(
+                    sel!(retain),
+                    retain as extern "C" fn(&RcTestObject, Sel) -> _,
+                );
+                builder.add_method(
+                    sel!(_tryRetain),
+                    try_retain as unsafe extern "C" fn(&RcTestObject, Sel) -> Bool,
+                );
+                builder.add_method(sel!(release), release as extern "C" fn(&RcTestObject, Sel));
+                builder.add_method(
+                    sel!(autorelease),
+                    autorelease as extern "C" fn(&RcTestObject, Sel) -> _,
+                );
+                builder.add_method(
+                    sel!(dealloc),
+                    dealloc as unsafe extern "C" fn(*mut RcTestObject, Sel),
+                );
+            }
+
+            builder.register();
+        });
+
+        class!(RcTestObject)
+    }
+
+    pub(crate) fn new() -> Id<Self, Owned> {
+        unsafe { Id::new(msg_send![Self::class(), new]) }.unwrap()
+    }
+}

objc2/src/rc/weak_id.rs

@@ -150,45 +150,68 @@ impl<T: Message> TryFrom<WeakId<T>> for Id<T, Shared> {
 
 #[cfg(test)]
 mod tests {
-    use super::WeakId;
-    use super::{Id, Shared};
+    use super::*;
+    use crate::rc::{RcTestObject, ThreadTestData};
     use crate::runtime::Object;
-    use crate::{class, msg_send};
 
     #[test]
     fn test_weak() {
-        let cls = class!(NSObject);
-        let obj: Id<Object, Shared> = unsafe {
-            let obj: *mut Object = msg_send![cls, alloc];
-            let obj: *mut Object = msg_send![obj, init];
-            Id::new(obj).unwrap()
-        };
+        let obj: Id<_, Shared> = RcTestObject::new().into();
+        let mut expected = ThreadTestData::current();
 
         let weak = WeakId::new(&obj);
+        expected.assert_current();
+
         let strong = weak.load().unwrap();
-        let strong_ptr: *const Object = &*strong;
-        let obj_ptr: *const Object = &*obj;
-        assert_eq!(strong_ptr, obj_ptr);
-        drop(strong);
+        expected.try_retain += 1;
+        expected.assert_current();
+        assert!(ptr::eq(&*strong, &*obj));
 
         drop(obj);
-        assert!(weak.load().is_none());
+        drop(strong);
+        expected.release += 2;
+        expected.dealloc += 1;
+        expected.assert_current();
+
+        if cfg!(not(feature = "gnustep-1-7")) {
+            // This loads the object on GNUStep for some reason??
+            assert!(weak.load().is_none());
+            expected.try_retain_fail += 1;
+            expected.assert_current();
+        }
+
+        drop(weak);
+        expected.assert_current();
     }
 
     #[test]
     fn test_weak_clone() {
-        let obj: Id<Object, Shared> = unsafe { Id::new(msg_send![class!(NSObject), new]).unwrap() };
+        let obj: Id<_, Shared> = RcTestObject::new().into();
+        let mut expected = ThreadTestData::current();
+
         let weak = WeakId::new(&obj);
+        expected.assert_current();
 
         let weak2 = weak.clone();
+        if cfg!(feature = "apple") {
+            expected.try_retain += 1;
+            expected.release += 1;
+        }
+        expected.assert_current();
 
         let strong = weak.load().unwrap();
+        expected.try_retain += 1;
+        expected.assert_current();
+        assert!(ptr::eq(&*strong, &*obj));
+
         let strong2 = weak2.load().unwrap();
-        let strong_ptr: *const Object = &*strong;
-        let strong2_ptr: *const Object = &*strong2;
-        let obj_ptr: *const Object = &*obj;
-        assert_eq!(strong_ptr, obj_ptr);
-        assert_eq!(strong2_ptr, obj_ptr);
+        expected.try_retain += 1;
+        expected.assert_current();
+        assert!(ptr::eq(&*strong, &*strong2));
+
+        drop(weak);
+        drop(weak2);
+        expected.assert_current();
     }
 
     #[test]