Cleanup hook blocks (#14)

Co-authored-by: Victor Pavlychko <[email protected]>

Author: steipete

…ose.xcworkspace/contents.xcworkspacedata …Kit.xcworkspace/contents.xcworkspacedataInterpose.xcworkspace/contents.xcworkspacedata renamed to InterposeKit.xcworkspace/contents.xcworkspacedata Interpose.xcworkspace/contents.xcworkspacedata renamed to InterposeKit.xcworkspace/contents.xcworkspacedata

@@ -29,6 +29,10 @@ final public class Interpose {
         }
     }
 
+    deinit {
+        tasks.forEach({ $0.cleanup() })
+    }
+
     /// Hook an `@objc dynamic` instance method via selector name on the current class.
     @discardableResult public func hook(_ selName: String,
                                         _ implementation: (Task) -> Any) throws -> Task {
@@ -146,6 +150,19 @@ extension Interpose {
             try execute(newState: .prepared) { try resetImplementation() }
         }
 
+        /// Release the hook block if possible.
+        fileprivate func cleanup() {
+            switch state {
+            case .prepared:
+                Interpose.log("Releasing -[\(`class`).\(selector)] IMP: \(replacementIMP!)")
+                imp_removeBlock(replacementIMP)
+            case .interposed:
+                Interpose.log("Keeping -[\(`class`).\(selector)] IMP: \(replacementIMP!)")
+            case let .error(error):
+                Interpose.log("Leaking -[\(`class`).\(selector)] IMP: \(replacementIMP!) due to error: \(error)")
+            }
+        }
+
         private func execute(newState: State, task: () throws -> Void) throws {
             do {
                 try task()
@@ -324,4 +341,5 @@ func method_getTypeEncoding(_ m: Method) -> UnsafePointer<Int8>? { return nil }
 // swiftlint:disable:next identifier_name
 func _dyld_register_func_for_add_image(_ func: (@convention(c) (UnsafePointer<Int8>?, Int) -> Void)!) {}
 func imp_implementationWithBlock(_ block: Any) -> IMP { IMP() }
+func imp_removeBlock(_ anImp: IMP) -> Bool { false }
 #endif

…ce/xcshareddata/IDEWorkspaceChecks.plist …ce/xcshareddata/IDEWorkspaceChecks.plistInterpose.xcworkspace/xcshareddata/IDEWorkspaceChecks.plist renamed to InterposeKit.xcworkspace/xcshareddata/IDEWorkspaceChecks.plist Interpose.xcworkspace/xcshareddata/IDEWorkspaceChecks.plist renamed to InterposeKit.xcworkspace/xcshareddata/IDEWorkspaceChecks.plist

@@ -10,6 +10,8 @@ class TestClass: NSObject {
         print(testClassHi)
         return testClassHi
     }
+
+    @objc dynamic func doNothing() { }
 }
 
 class TestSubclass: TestClass {
@@ -93,8 +95,100 @@ final class InterposeKitTests: XCTestCase {
         XCTAssertEqual(testObj.sayHi(), testClassHi + testSubclass)
     }
 
+    func testInterposedCleanup() throws {
+        var deallocated = false
+
+        try autoreleasepool {
+            let tracker = LifetimeTracker {
+                deallocated = true
+            }
+
+            // Swizzle test class
+            let interposer = try Interpose(TestClass.self) {
+                try $0.hook(#selector(TestClass.doNothing), { store in { `self` in
+                    tracker.keep()
+                    let origCall = store((@convention(c) (AnyObject, Selector) -> Void).self)
+                    return origCall(`self`, store.selector)
+                } as @convention(block) (AnyObject) -> Void })
+            }
+
+            // Dealloc interposer without removing hooks
+            _ = interposer
+        }
+
+        // Unreverted block should not be deallocated
+        XCTAssertFalse(deallocated)
+    }
+
+    func testRevertedCleanup() throws {
+        var deallocated = false
+
+        try autoreleasepool {
+            let tracker = LifetimeTracker {
+                deallocated = true
+            }
+
+            // Swizzle test class
+            let interposer = try Interpose(TestClass.self) {
+                try $0.hook(#selector(TestClass.doNothing), { store in { `self` in
+                    tracker.keep()
+                    let origCall = store((@convention(c) (AnyObject, Selector) -> Void).self)
+                    return origCall(`self`, store.selector)
+                } as @convention(block) (AnyObject) -> Void })
+            }
+
+            try interposer.revert()
+        }
+
+        // Verify that the block was deallocated
+        XCTAssertTrue(deallocated)
+    }
+
+    func testImpRemoveBlockWorks() {
+        var deallocated = false
+
+        let imp: IMP = autoreleasepool {
+            let tracker = LifetimeTracker {
+                deallocated = true
+            }
+
+            let block: @convention(block) (AnyObject) -> Void = { _ in
+                // retain `tracker` inside a block
+                tracker.keep()
+            }
+
+            return imp_implementationWithBlock(block)
+        }
+
+        // `imp` retains `block` which retains `tracker`
+        XCTAssertFalse(deallocated)
+
+        // Detach `block` from `imp`
+        imp_removeBlock(imp)
+
+        // `block` and `tracker` should be deallocated now
+        XCTAssertTrue(deallocated)
+    }
+
+    class LifetimeTracker {
+        let deinitCalled: () -> Void
+
+        init(deinitCalled: @escaping () -> Void) {
+            self.deinitCalled = deinitCalled
+        }
+
+        deinit {
+            deinitCalled()
+        }
+
+        func keep() { }
+    }
+
     static var allTests = [
         ("testClassOverrideAndRevert", testClassOverrideAndRevert),
-        ("testSubclassOverride", testSubclassOverride)
+        ("testSubclassOverride", testSubclassOverride),
+        ("testInterposedCleanup", testInterposedCleanup),
+        ("testRevertedCleanup", testRevertedCleanup),
+        ("testImpRemoveBlockWorks", testImpRemoveBlockWorks)
     ]
 }