Support strong, weak and unowned reference type while creating instance-based Interpose

Author: ezoushen

Sources/InterposeKit/AnyObjectContainer.swift

@@ -0,0 +1,63 @@
+import Foundation
+
+/// Container for AnyObject
+public class AnyObjectContainer {
+    public var object: AnyObject {
+        fatalError("Always override")
+    }
+}
+
+/// The container that hold a strong reference to the object
+internal class StrongObjectContainer: AnyObjectContainer {
+    override var object: AnyObject {
+        _object
+    }
+
+    private let _object: AnyObject
+
+    init(_ object: AnyObject) {
+        self._object = object
+    }
+}
+
+/// The container that hold a weak reference to the object
+internal class WeakObjectContainer: AnyObjectContainer {
+    override var object: AnyObject {
+        guard let object = _object else { fatalError("Bad Access") }
+        return object
+    }
+
+    private weak var _object: AnyObject?
+
+    init(_ object: AnyObject) {
+        self._object = object
+    }
+}
+
+/// The container that hold an unowned reference to the object
+internal class UnownedObjectContainer: AnyObjectContainer {
+    override var object: AnyObject { _object }
+
+    private unowned let _object: AnyObject
+
+    init(_ object: AnyObject) {
+        self._object = object
+    }
+}
+
+extension AnyObjectContainer {
+    /// Create a strong reference container
+    public static func strong(_ object: AnyObject) -> AnyObjectContainer {
+        StrongObjectContainer(object)
+    }
+
+    /// Create a weak reference container
+    public static func weak(_ object: AnyObject) -> AnyObjectContainer {
+        WeakObjectContainer(object)
+    }
+
+    /// Create an unowned reference container
+    public static func unowned(_ object: AnyObject) -> AnyObjectContainer {
+        UnownedObjectContainer(object)
+    }
+}

Sources/InterposeKit/InterposeKit.swift

@@ -30,13 +30,18 @@ extension NSObject {
 /// Methods are hooked via replacing the implementation, instead of the usual exchange.
 /// Supports both swizzling classes and individual objects.
 final public class Interpose {
+
     /// Stores swizzle hooks and executes them at once.
     public let `class`: AnyClass
     /// Lists all hooks for the current interpose class object.
     public private(set) var hooks: [AnyHook] = []
 
     /// If Interposing is object-based, this is set.
-    public let object: AnyObject?
+    public var object: AnyObject? {
+        objectContainer?.object
+    }
+
+    internal let objectContainer: AnyObjectContainer?
 
     // Checks if a object is posing as a different class
     // via implementing 'class' and returning something else.
@@ -58,7 +63,7 @@ final public class Interpose {
     /// If `builder` is present, `apply()` is automatically called.
     public init(_ `class`: AnyClass, builder: ((Interpose) throws -> Void)? = nil) throws {
         self.class = `class`
-        self.object = nil
+        self.objectContainer = nil
 
         // Only apply if a builder is present
         if let builder = builder {
@@ -67,9 +72,16 @@ final public class Interpose {
     }
 
     /// Initialize with a single object to interpose.
-    public init(_ object: NSObject, builder: ((Interpose) throws -> Void)? = nil) throws {
-        self.object = object
-        self.class = type(of: object)
+    public convenience init(_ object: NSObject, builder: ((Interpose) throws -> Void)? = nil) throws {
+        try self.init(.strong(object), builder: builder)
+    }
+
+    /// Initialize with a single object to interpose.
+    public init(_ objectContainer: AnyObjectContainer, builder: ((Interpose) throws -> Void)? = nil) throws {
+        self.objectContainer = objectContainer
+        self.class = type(of: objectContainer.object)
+
+        let object = objectContainer.object
 
         if let actualClass = checkObjectPosingAsDifferentClass(object) {
             if isKVORuntimeGeneratedClass(actualClass) {
@@ -106,8 +118,8 @@ final public class Interpose {
        _ implementation:(TypedHook<MethodSignature, HookSignature>) -> HookSignature?) throws -> TypedHook<MethodSignature, HookSignature> {
 
         var hook: TypedHook<MethodSignature, HookSignature>
-        if let object = self.object {
-            hook = try ObjectHook(object: object, selector: selector, implementation: implementation)
+        if let objectContainer = self.objectContainer {
+            hook = try ObjectHook(objectContainer: objectContainer, selector: selector, implementation: implementation)
         } else {
             hook = try ClassHook(class: `class`, selector: selector, implementation: implementation)
         }

Sources/InterposeKit/ObjectHook.swift

@@ -46,7 +46,11 @@ extension Interpose {
     final public class ObjectHook<MethodSignature, HookSignature>: TypedHook<MethodSignature, HookSignature> {
 
         /// The object that is being hooked.
-        public let object: AnyObject
+        public var object: AnyObject {
+            objectContainer.object
+        }
+
+        private let objectContainer: AnyObjectContainer
 
         /// Subclass that we create on the fly
         var dynamicSubclass: AnyClass?
@@ -55,9 +59,14 @@ extension Interpose {
         let generatesSuperIMP = isSupportedArchitectureForSuper()
 
         /// Initialize a new hook to interpose an instance method.
-        public init(object: AnyObject, selector: Selector, implementation:(ObjectHook<MethodSignature, HookSignature>) -> HookSignature?) throws {
-            self.object = object
-            try super.init(class: type(of: object), selector: selector)
+        public convenience init(object: AnyObject, selector: Selector, implementation:(ObjectHook<MethodSignature, HookSignature>) -> HookSignature?) throws {
+            try self.init(objectContainer: .strong(object), selector: selector, implementation: implementation)
+        }
+
+        /// Initialize a new hook to interpose an instance method.
+        public init(objectContainer: AnyObjectContainer, selector: Selector, implementation:(ObjectHook<MethodSignature, HookSignature>) -> HookSignature?) throws {
+            self.objectContainer = objectContainer
+            try super.init(class: type(of: objectContainer.object), selector: selector)
             let block = implementation(self) as AnyObject
             replacementIMP = imp_implementationWithBlock(block)
             guard replacementIMP != nil else {