From 2f906ee24ffbaead6e18e1d573e2de5865e19b8f Mon Sep 17 00:00:00 2001
From: Jeff Muizelaar <jmuizelaar@mozilla.com>
Date: Wed, 11 Jul 2018 11:43:40 -0400
Subject: [PATCH] Add autoreleasepool functionality.

This replicates the behaviour of @autoreleasepool blocks in
Objective-C and allows higher performance creation of autorelease pools.
---
 src/rc/autorelease.rs | 29 +++++++++++++++++++++++++++++
 src/rc/mod.rs         | 23 +++++++++++++++++++++++
 src/runtime.rs        |  3 +++
 3 files changed, 55 insertions(+)
 create mode 100644 src/rc/autorelease.rs

diff --git a/src/rc/autorelease.rs b/src/rc/autorelease.rs
new file mode 100644
index 000000000..3b68b8d76
--- /dev/null
+++ b/src/rc/autorelease.rs
@@ -0,0 +1,29 @@
+use std::os::raw::c_void;
+use runtime::{objc_autoreleasePoolPush, objc_autoreleasePoolPop};
+
+// we use a struct to ensure that objc_autoreleasePoolPop during unwinding.
+struct AutoReleaseHelper {
+    context: *mut c_void,
+}
+
+impl AutoReleaseHelper {
+    unsafe fn new() -> Self {
+        AutoReleaseHelper { context: objc_autoreleasePoolPush() }
+    }
+}
+
+impl Drop for AutoReleaseHelper {
+    fn drop(&mut self) {
+        unsafe { objc_autoreleasePoolPop(self.context) }
+    }
+}
+
+/**
+Execute `f` in the context of a new autorelease pool. The pool is drained
+after the execution of `f` completes. This corresponds to @autoreleasepool blocks
+in Objective-c and Swift.
+*/
+pub fn autoreleasepool<F: FnOnce()>(f: F) {
+    let _context = unsafe { AutoReleaseHelper::new() };
+    f();
+}
diff --git a/src/rc/mod.rs b/src/rc/mod.rs
index 4eda1b3dd..8f70ece86 100644
--- a/src/rc/mod.rs
+++ b/src/rc/mod.rs
@@ -11,15 +11,18 @@ For more information on Objective-C's reference counting, see Apple's documentat
 
 mod strong;
 mod weak;
+mod autorelease;
 
 pub use self::strong::StrongPtr;
 pub use self::weak::WeakPtr;
+pub use self::autorelease::autoreleasepool;
 
 // These tests use NSObject, which isn't present for GNUstep
 #[cfg(all(test, any(target_os = "macos", target_os = "ios")))]
 mod tests {
     use runtime::Object;
     use super::StrongPtr;
+    use super::autoreleasepool;
 
     #[test]
     fn test_strong_clone() {
@@ -66,4 +69,24 @@ mod tests {
         let strong = weak2.load();
         assert!(*strong == *obj);
     }
+
+    #[test]
+    fn test_autorelease() {
+        let obj = unsafe {
+            StrongPtr::new(msg_send![class!(NSObject), new])
+        };
+
+        fn retain_count(obj: *mut Object) -> usize {
+            unsafe { msg_send![obj, retainCount] }
+        }
+        let cloned = obj.clone();
+
+        autoreleasepool(|| {
+                        obj.autorelease();
+                        assert!(retain_count(*cloned) == 2);
+        });
+
+        // make sure that the autoreleased value has been released
+        assert!(retain_count(*cloned) == 1);
+    }
 }
diff --git a/src/runtime.rs b/src/runtime.rs
index 46ae9396e..3b533efe6 100644
--- a/src/runtime.rs
+++ b/src/runtime.rs
@@ -109,6 +109,9 @@ extern {
     pub fn objc_allocateProtocol(name: *const c_char) -> *mut Protocol;
     pub fn objc_registerProtocol(proto: *mut Protocol);
 
+    pub fn objc_autoreleasePoolPush() -> *mut c_void;
+    pub fn objc_autoreleasePoolPop(context: *mut c_void);
+
     pub fn protocol_addMethodDescription(proto: *mut Protocol, name: Sel, types: *const c_char, isRequiredMethod: BOOL,
                                          isInstanceMethod: BOOL);
     pub fn protocol_addProtocol(proto: *mut Protocol, addition: *const Protocol);